Pablu/folder_encryptor
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Added Parallelism, also cleaned up code a bit

Browse Source
This commit is contained in:
Pablu23
2023-03-09 19:45:56 +01:00
parent 11d3eab05c
commit aec8725474
1 changed files with 105 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

153
src/main.rs
View File

@@ -9,7 +9,9 @@ use std::{
fs::File,
io,
io::{Read, Write},
// time::SystemTime,
sync::Arc, // time::SystemTime,
thread,
thread::{available_parallelism, JoinHandle},
};
use uuid::Uuid;
@@ -22,10 +24,10 @@ const BUFFER_LEN: usize = 50 * 1024 * 1024; // 50 MiB
fn encrypt_file(
source_path: String,
dest_path: String,
nonce: &[u8; 24],
key: &Vec<u8>,
nonce: [u8; 24],
key: &[u8],
) -> io::Result<()> {
let mut cipher = XChaCha20::new(key[..32].as_ref().into(), nonce.into());
let mut cipher = XChaCha20::new(key[..32].as_ref().into(), &nonce.into());
let source_file_path = Path::new(&source_path);
@@ -48,10 +50,12 @@ fn encrypt_file(
// Stack allocated buffer
// let mut buffer = [0u8; BUFFER_LEN];
println!("Start encrypting File: {}", file_name.to_str().unwrap());
// Heap allocated buffer (Allows larger sized buffer, up to 50 % max ram)
let mut buffer = vec![0u8; BUFFER_LEN].into_boxed_slice();
dest_file.write(nonce)?;
dest_file.write(&nonce)?;
let mut f_name_bytes = file_name.to_str().unwrap_or_default().as_bytes().to_owned();
@@ -83,29 +87,23 @@ fn encrypt_file(
fs::remove_file(source_file_path)?;
return Ok(());
println!("Finished encrypting File: {}", file_name.to_str().unwrap());
Ok(())
}
fn decrypt_file(source_path: String, pwd: &String, config: &argon2::Config) -> io::Result<()> {
fn decrypt_file(source_path: &Path, pwd: &String, config: &argon2::Config) -> io::Result<()> {
let mut nonce = [0u8; 24];
let path_info = Path::new(&source_path);
if !path_info.try_exists()? {
if !source_path.try_exists()? {
return Err(io::Error::from(io::ErrorKind::NotFound));
}
//let path = &source_path[..&source_path.len() - 4];
// if Path::new(&path).exists() {
// fs::remove_file(&path)?;
// }
let mut source_file = File::open(&source_path)?;
source_file.read(&mut nonce)?;
let key = argon2::hash_raw(pwd.as_bytes(), &nonce, &config).unwrap();
let key = argon2::hash_raw(pwd.as_bytes(), &nonce, config).unwrap();
let mut cipher = XChaCha20::new(key[..32].as_ref().into(), &nonce.into());
@@ -116,14 +114,12 @@ fn decrypt_file(source_path: String, pwd: &String, config: &argon2::Config) -> i
source_file.read(&mut file_name_size_buffer)?;
let file_name_size = u16::from_le_bytes(file_name_size_buffer);
println!("File Name Size: {file_name_size}");
let mut file_name_bytes = vec![0u8; file_name_size.into()];
source_file.read_exact(&mut file_name_bytes)?;
cipher.apply_keystream(&mut file_name_bytes);
let private_dir_path = match path_info.parent() {
let private_dir_path = match source_path.parent() {
Some(p) => Ok(p),
None => Err(io::Error::from(io::ErrorKind::AddrNotAvailable)),
}?;
@@ -134,22 +130,18 @@ fn decrypt_file(source_path: String, pwd: &String, config: &argon2::Config) -> i
}?;
let file_name = String::from_utf8(file_name_bytes).unwrap_or_default();
let path = root_dir_path.join(&file_name); //to_str().unwrap_or_default().to_owned() + &file_name;
let path = root_dir_path.join(&file_name);
println!("File Name: {file_name}\n Path: {path:?}");
println!("Start decrypting File: {file_name}");
let mut dest_file = File::create(path)?;
println!("Worked!");
// Heap allocated buffer (Allows larger sized buffer, up to 50 % max ram)
// Heap allocated buffer (Allows larger sized buffer, up to 50 % max ram, technically more)
let mut buffer = vec![0u8; BUFFER_LEN].into_boxed_slice();
loop {
let read_count = source_file.read(&mut buffer)?;
println!("Read Bytes {read_count}");
if read_count == BUFFER_LEN {
cipher.apply_keystream(&mut buffer);
dest_file.write(&buffer)?;
@@ -160,43 +152,77 @@ fn decrypt_file(source_path: String, pwd: &String, config: &argon2::Config) -> i
}
}
println!("Fully written");
return Ok(());
println!("Finished decrypting File: {file_name}");
Ok(())
}
fn main() -> io::Result<()> {
let config = argon2::Config {
let config = Arc::new(argon2::Config {
variant: argon2::Variant::Argon2id,
hash_length: 32,
lanes: 8,
mem_cost: 16 * 1024,
time_cost: 8,
..Default::default()
};
});
let exe = env::current_exe().unwrap();
let cwd = env::current_dir().unwrap();
let private = cwd.join("private");
let mut max_threads: usize = available_parallelism().unwrap().into();
let percent = max_threads as f32 * 0.7;
max_threads = percent.floor() as usize;
if max_threads < 1 {
max_threads = 1;
}
println!("Parallelism: {max_threads}");
if private.exists() {
let paths = fs::read_dir(&private).unwrap();
print!("Type password for encrypted files: ");
std::io::stdout().flush().unwrap();
let pwd = read_password().unwrap();
let pwd = Arc::new(read_password().unwrap());
let config = Arc::new(config);
let mut handles: Vec<JoinHandle<()>> = Vec::with_capacity(max_threads);
let mut current_threads = 0;
for path_result in paths {
let path = path_result.unwrap().path();
let pwd = pwd.clone();
let config = config.clone();
if path.is_file() {
decrypt_file(String::from(path.to_str().unwrap()), &pwd, &config)?;
fs::remove_file(String::from(path.to_str().unwrap()))?;
handles.push(thread::spawn(move || {
if path.is_file() {
decrypt_file(path.as_path(), &pwd, &config).unwrap();
fs::remove_file(String::from(path.to_str().unwrap())).unwrap();
}
}));
current_threads += 1;
if current_threads >= max_threads {
while let Some(handle) = handles.pop() {
handle.join().unwrap();
current_threads -= 1;
}
}
}
if current_threads > 0 {
while let Some(handle) = handles.pop() {
handle.join().unwrap();
current_threads -= 1;
}
}
fs::remove_dir(private)?;
//fs::remove_dir_all(private).unwrap();
} else {
let paths = fs::read_dir(&cwd).unwrap();
@@ -222,26 +248,57 @@ fn main() -> io::Result<()> {
print!("Type password to encrypt files: ");
std::io::stdout().flush().unwrap();
let pwd = read_password().unwrap();
let pwd = Arc::new(read_password().unwrap());
let mut nonce = [0u8; 24];
OsRng.fill_bytes(&mut nonce);
let key = argon2::hash_raw(pwd.as_bytes(), &nonce, &config).unwrap();
// let key = Arc::new(argon2::hash_raw(pwd.as_bytes(), &nonce, &config).unwrap());
let paths = fs::read_dir(cwd).unwrap();
let exe = Arc::new(exe);
let private = Arc::new(private);
let mut handles: Vec<JoinHandle<()>> = Vec::with_capacity(max_threads);
let mut current_threads = 0;
for path_result in paths {
let path = path_result.unwrap().path();
let pwd = pwd.clone();
let exe = exe.clone();
let private = private.clone();
let config = config.clone();
if path.is_file() && path != exe {
encrypt_file(
String::from(path.to_str().unwrap()),
String::from(private.to_str().unwrap()),
&nonce,
&key,
)
.unwrap();
handles.push(thread::spawn(move || {
let mut nonce = [0u8; 24];
OsRng.fill_bytes(&mut nonce);
let key = argon2::hash_raw(pwd.as_bytes(), &nonce, &config).unwrap();
let path = path_result.unwrap().path();
if path.is_file() && path.as_os_str() != exe.as_os_str() {
encrypt_file(
String::from(path.to_str().unwrap()),
String::from(private.to_str().unwrap()),
nonce,
&key,
)
.unwrap();
}
}));
current_threads += 1;
if current_threads >= max_threads {
while let Some(handle) = handles.pop() {
handle.join().unwrap();
current_threads -= 1;
}
}
}
if current_threads > 0 {
while let Some(handle) = handles.pop() {
handle.join().unwrap();
current_threads -= 1;
}
}
}