diff --git a/src/lib.rs b/src/lib.rs index 367ed93..3b4f806 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,6 +7,7 @@ use std::process::exit; use indicatif::{ProgressBar, ProgressStyle}; use std::fmt::{Display, Formatter, Result}; +use std::path::Path; use walkdir::WalkDir; const BATCH_SIZE: usize = 8192; @@ -47,6 +48,7 @@ impl<'a> Shredder<'a> { } else { if self.configuration.is_recursive { Shredder::run_directory_shredding(&configuration, &configuration.raw_path); + fs::remove_dir_all(Path::new(&configuration.raw_path)).unwrap(); } else { println!("Target is a directory!"); exit(1); @@ -79,14 +81,12 @@ impl<'a> Shredder<'a> { } println!("{}", absolute_path); - for iteration in 0..configuration.rewrite_iterations { - >::shred_file( - &file, - file_length, - absolute_path, - &iteration, - ); - } + >::shred_file( + &file, + file_length, + absolute_path, + &configuration.rewrite_iterations, + ); if !configuration.keep_files { fs::remove_file(absolute_path).unwrap(); @@ -109,37 +109,39 @@ impl<'a> Shredder<'a> { }; } - fn shred_file(file: &File, file_length: u64, absolute_path: &str, iteration: &u8) { + fn shred_file(file: &File, file_length: u64, absolute_path: &str, iterations: &u8) { let mut buffer = BufWriter::new(file); - let mut bytes_processed = 0; - let pb = ProgressBar::new(file_length); - pb.set_prefix(&format!("Iteration #{}", iteration + 1)); pb.set_message(absolute_path); pb.set_style(ProgressStyle::default_bar() .template("{prefix} {spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {bytes}/{total_bytes} {bytes_per_sec} ({eta})") .progress_chars("#>-")); - while bytes_processed < file_length { - let bytes_to_write = if file_length - bytes_processed > BATCH_SIZE as u64 { - BATCH_SIZE - } else { - (file_length - bytes_processed) as usize - }; + for i in 0..*iterations { + pb.set_prefix(&format!("Iteration {}/{}", i + 1, *iterations)); + let mut bytes_processed = 0; + while bytes_processed < file_length { + let bytes_to_write = if file_length - bytes_processed > BATCH_SIZE as u64 { + BATCH_SIZE + } else { + (file_length - bytes_processed) as usize + }; - let random_bytes: Vec = (0..bytes_to_write).map(|_| rand::random::()).collect(); - buffer.write(&random_bytes).unwrap(); + let random_bytes: Vec = + (0..bytes_to_write).map(|_| rand::random::()).collect(); + buffer.write(&random_bytes).unwrap(); - bytes_processed = bytes_processed + bytes_to_write as u64; + bytes_processed = bytes_processed + bytes_to_write as u64; - pb.set_position(bytes_processed); + pb.set_position(bytes_processed); + } + pb.finish_with_message("shredded"); + + buffer.flush().unwrap(); + file.sync_all().unwrap(); + buffer.seek(SeekFrom::Start(0)).unwrap(); } - pb.finish_with_message("shredded"); - - buffer.flush().unwrap(); - file.sync_all().unwrap(); - buffer.seek(SeekFrom::Start(0)).unwrap(); } fn run_directory_shredding(configuration: &ShredConfiguration, relative_path: &str) { diff --git a/src/main.rs b/src/main.rs index f4457f4..09321ba 100644 --- a/src/main.rs +++ b/src/main.rs @@ -43,6 +43,7 @@ fn main() { .arg( Arg::with_name("n") .short("n") + .takes_value(true) .help("How many times the file must be overridden"), ) .get_matches();