From 38569e632ed3ede5a43e6c94f0e054469b22afa0 Mon Sep 17 00:00:00 2001 From: Trivernis Date: Fri, 6 Mar 2020 11:49:35 +0100 Subject: [PATCH] Change mapping function to iterators only --- src/lib/hash.rs | 16 ++---------- src/main.rs | 67 ++++++++++++++++++++++++------------------------- 2 files changed, 35 insertions(+), 48 deletions(-) diff --git a/src/lib/hash.rs b/src/lib/hash.rs index a5fd159..c194220 100644 --- a/src/lib/hash.rs +++ b/src/lib/hash.rs @@ -1,27 +1,15 @@ -use rayon::prelude::*; use sha2::{Digest, Sha256}; pub type PassKey = (String, Vec); /// Hashes a text to a 32 bytes long key. -pub fn create_key(pw: String) -> Vec { +pub fn create_key(pw: &str) -> Vec { let mut hasher = Sha256::default(); hasher.input(pw); let result = hasher.result(); let key = &result[0..8]; - key.to_vec() -} - -/// Maps a list of passwords to keys and returns a vector of pairs -pub fn map_to_keys(passwords: Vec) -> Vec { - passwords - .par_iter() - .map(|pw| { - let pw_str = (*pw).clone(); - (pw_str.clone(), create_key(pw_str)) - }) - .collect() + key.to_vec().clone() } /// Creates a sha256 hashsum from the input data diff --git a/src/main.rs b/src/main.rs index 46b19c5..fad9563 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,7 +3,7 @@ pub mod lib; use crate::lib::crypt::{ decrypt_brute_brute_force, decrypt_data, decrypt_with_dictionary, encrypt_data, }; -use crate::lib::hash::{create_key, map_to_keys, sha_checksum, PassKey}; +use crate::lib::hash::{create_key, sha_checksum, PassKey}; use pbr::ProgressBar; use rayon::prelude::*; use rpassword; @@ -12,6 +12,8 @@ use spinners::{Spinner, Spinners}; use std::fs; use std::fs::File; use std::io::Write; +use std::sync::mpsc::channel; +use std::thread; use std::time::Duration; use structopt::StructOpt; @@ -106,7 +108,7 @@ fn encrypt(_opts: &Opts, args: &Encrypt) { .expect("Failed to write checksum file!"); } let pass = read_password_from_tty(Some("Password: ")).unwrap(); - let key = create_key(pass); + let key = create_key(&pass); let enc_data = encrypt_data(data.as_slice(), key.as_slice()); fs::write(output, enc_data.as_slice()).expect("Failed to write output file!"); } @@ -161,7 +163,7 @@ fn decrypt(_opts: &Opts, args: &Decrypt) { } } else { let pass = read_password_from_tty(Some("Password: ")).unwrap(); - let key = create_key(pass); + let key = create_key(&pass); let result = decrypt_data(&data, key.as_slice()); fs::write(output, &result).expect("Failed to write output file!"); } @@ -171,43 +173,40 @@ fn decrypt(_opts: &Opts, args: &Decrypt) { fn create_dictionary(_opts: &Opts, args: &CreateDictionary) { let input: String = (*args.input).parse().unwrap(); let sp = spinner("Parsing passwords..."); - let dictionary: Vec; - { - let passwords = get_lines(&input); - // TODO: Some form of removing duplicates (without itertools) - sp.message("Mapping passwords to keys".into()); - dictionary = map_to_keys(passwords); - } - sp.message("Encoding data to lines".into()); - let dict_lines: Vec = dictionary - .par_iter() - .map(|(pw, key): &PassKey| -> String { - let key_base64 = base64::encode(key.as_slice()); - format!("{},{}\n", pw, key_base64) - }) - .collect(); + let contents = fs::read_to_string(input).expect("Failed to read input file!"); + let passwords = contents.lines().collect::>(); + // TODO: Some form of removing duplicates (without itertools) sp.stop(); - println!("\rWriting passwords to file..."); let mut fout = File::create(args.output.clone()).unwrap(); - let mut pb = ProgressBar::new(dict_lines.len() as u64); - pb.set_max_refresh_rate(Some(Duration::from_millis(200))); - for line in dict_lines { - fout.write(&line.into_bytes()).unwrap(); - pb.inc(); + let handle; + { + let mut pb = ProgressBar::new(passwords.len() as u64); + pb.set_max_refresh_rate(Some(Duration::from_millis(200))); + let (rx, tx) = channel::(); + handle = thread::spawn(move || { + for line in tx { + fout.write(&line.as_bytes()).unwrap(); + pb.inc(); + } + pb.finish(); + }); + passwords + .par_iter() + .map(|pw| -> String { + let key = create_key(pw); + let key_base64 = base64::encode(key.as_slice()); + format!("{},{}\n", pw, key_base64) + }) + .for_each_with(rx, |rx, line| { + rx.send(line).expect("Failed to send value to channel."); + }); + } + if let Err(_err) = handle.join() { + println!("Failed to join!"); } - pb.finish(); println!("Finished!"); } -fn get_lines(filename: &str) -> Vec { - let contents = fs::read_to_string(filename).expect("Failed to read input file!"); - let lines = contents.lines().collect::>(); - lines - .par_iter() - .map(|s| -> String { s.parse().unwrap() }) - .collect() -} - /// Creates a new spinner with a given text fn spinner(text: &str) -> Spinner { Spinner::new(Spinners::Dots2, text.into())