diff --git a/src/source/wordlist.rs b/src/source/wordlist.rs index 598cf08..939e471 100644 --- a/src/source/wordlist.rs +++ b/src/source/wordlist.rs @@ -1,11 +1,13 @@ //! Wordlist source - generate keys from file of passphrases. +//! +//! Streams the file in chunks to avoid loading entire wordlists into memory. use anyhow::Result; use indicatif::ProgressBar; use rayon::prelude::*; -use std::fs::File; -use std::io::{BufRead, BufReader}; -use std::path::Path; +use std::fs; +use std::io::{BufRead, BufReader, Seek}; +use std::path::{Path, PathBuf}; use super::{ProcessStats, Source}; use crate::derive::KeyDeriver; @@ -13,28 +15,25 @@ use crate::matcher::Matcher; use crate::output::Output; use crate::transform::{Input, Transform}; +const CHUNK_SIZE: usize = 100_000; +const BATCH_SIZE: usize = 1000; + /// Generate keys from a wordlist file pub struct WordlistSource { - lines: Vec, + path: PathBuf, } impl WordlistSource { pub fn from_file(path: &Path) -> Result { - let file = File::open(path)?; - let reader = BufReader::new(file); - let lines: Vec = reader - .lines() - .filter_map(|line| match line { - Ok(s) => { - let trimmed = s.trim().to_string(); - if trimmed.is_empty() { None } else { Some(Ok(trimmed)) } - } - Err(e) if e.kind() == std::io::ErrorKind::InvalidData => None, - Err(e) => Some(Err(e)), - }) - .collect::>>()?; - - Ok(Self { lines }) + if !path.exists() { + anyhow::bail!("Wordlist file not found: {}", path.display()); + } + if !path.is_file() { + anyhow::bail!("Not a file: {}", path.display()); + } + Ok(Self { + path: path.to_path_buf(), + }) } } @@ -46,51 +45,176 @@ impl Source for WordlistSource { matcher: Option<&Matcher>, output: &dyn Output, ) -> Result { - let pb = ProgressBar::new(self.lines.len() as u64); + let file_size = fs::metadata(&self.path)?.len(); + let pb = ProgressBar::new(file_size); pb.set_style(crate::default_progress_style()); - let stats = std::sync::atomic::AtomicU64::new(0); let matches = std::sync::atomic::AtomicU64::new(0); + let mut inputs_processed = 0u64; + let mut bytes_consumed = 0u64; + + let file = std::fs::File::open(&self.path)?; + let mut reader = BufReader::new(file); + let mut chunk = Vec::with_capacity(CHUNK_SIZE); + let mut line_buf = String::new(); + + loop { + line_buf.clear(); + let bytes_read = match reader.read_line(&mut line_buf) { + Ok(0) => break, + Ok(n) => n as u64, + Err(e) if e.kind() == std::io::ErrorKind::InvalidData => { + // read_line() already consumed the bytes, sync position from reader + bytes_consumed = reader.stream_position().unwrap_or(bytes_consumed); + pb.set_position(bytes_consumed); + continue; + } + Err(e) => return Err(e.into()), + }; - self.lines.par_chunks(1000).for_each(|chunk| { - let inputs: Vec = chunk - .iter() - .map(|s| Input::from_string(s.clone())) - .collect(); - let mut buffer = Vec::with_capacity(inputs.len() * 2); - - for transform in transforms { - buffer.clear(); - transform.apply_batch(&inputs, &mut buffer); - - for (source, key) in &buffer { - let derived = deriver.derive(key); - - if let Some(m) = matcher { - if let Some(match_info) = m.check(&derived) { - output - .hit(source, transform.name(), &derived, &match_info) - .ok(); - matches.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - } - } else { - output.key(source, transform.name(), &derived).ok(); - } + bytes_consumed += bytes_read; + pb.set_position(bytes_consumed); - stats.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - } + let trimmed = line_buf.trim().to_string(); + if trimmed.is_empty() { + continue; } - pb.inc(chunk.len() as u64); - }); + chunk.push(trimmed); + inputs_processed += 1; + + if chunk.len() >= CHUNK_SIZE { + process_chunk( + &chunk, transforms, deriver, matcher, output, &stats, &matches, + ); + chunk.clear(); + } + } + + if !chunk.is_empty() { + process_chunk( + &chunk, transforms, deriver, matcher, output, &stats, &matches, + ); + } pb.finish_and_clear(); Ok(ProcessStats { - inputs_processed: self.lines.len() as u64, + inputs_processed, keys_generated: stats.load(std::sync::atomic::Ordering::Relaxed), matches_found: matches.load(std::sync::atomic::Ordering::Relaxed), }) } } + +fn process_chunk( + lines: &[String], + transforms: &[Box], + deriver: &KeyDeriver, + matcher: Option<&Matcher>, + output: &dyn Output, + stats: &std::sync::atomic::AtomicU64, + matches: &std::sync::atomic::AtomicU64, +) { + lines.par_chunks(BATCH_SIZE).for_each(|batch| { + let inputs: Vec = batch + .iter() + .map(|s| Input::from_string(s.clone())) + .collect(); + let mut buffer = Vec::with_capacity(inputs.len() * 2); + + for transform in transforms { + buffer.clear(); + transform.apply_batch(&inputs, &mut buffer); + + for (source, key) in &buffer { + let derived = deriver.derive(key); + + if let Some(m) = matcher { + if let Some(match_info) = m.check(&derived) { + output + .hit(source, transform.name(), &derived, &match_info) + .ok(); + matches.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + } else { + output.key(source, transform.name(), &derived).ok(); + } + + stats.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + } + }); +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::derive::KeyDeriver; + use crate::output::ConsoleOutput; + use std::io::Write; + use tempfile::NamedTempFile; + + #[test] + fn from_file_not_found() { + let result = WordlistSource::from_file(Path::new("/nonexistent/path/file.txt")); + assert!(result.is_err()); + } + + #[test] + fn from_file_not_a_file() { + let dir = tempfile::tempdir().unwrap(); + let result = WordlistSource::from_file(dir.path()); + assert!(result.is_err()); + } + + #[test] + fn process_empty_file() { + let mut file = NamedTempFile::new().unwrap(); + file.write_all(b"").unwrap(); + + let source = WordlistSource::from_file(file.path()).unwrap(); + let deriver = KeyDeriver::new(); + let output = ConsoleOutput::new(); + let transforms: Vec> = Vec::new(); + + let stats = source + .process(&transforms, &deriver, None, &output) + .unwrap(); + assert_eq!(stats.inputs_processed, 0); + } + + #[test] + fn process_skips_blank_lines() { + let mut file = NamedTempFile::new().unwrap(); + file.write_all(b"hello\n\n \nworld\n").unwrap(); + + let source = WordlistSource::from_file(file.path()).unwrap(); + let deriver = KeyDeriver::new(); + let output = ConsoleOutput::new(); + let transforms: Vec> = Vec::new(); + + let stats = source + .process(&transforms, &deriver, None, &output) + .unwrap(); + assert_eq!(stats.inputs_processed, 2); + } + + #[test] + fn process_skips_invalid_utf8() { + let mut file = NamedTempFile::new().unwrap(); + file.write_all(b"valid\n\xff\xfe\ninvalid bytes\ntest\n") + .unwrap(); + + let source = WordlistSource::from_file(file.path()).unwrap(); + let deriver = KeyDeriver::new(); + let output = ConsoleOutput::new(); + let transforms: Vec> = Vec::new(); + + let stats = source + .process(&transforms, &deriver, None, &output) + .unwrap(); + assert_eq!(stats.inputs_processed, 3); + } +}