diff --git a/crates/dtmt/src/cmd/experiment/brute_force_words.rs b/crates/dtmt/src/cmd/experiment/brute_force_words.rs index d2891f9..bb3aa9e 100644 --- a/crates/dtmt/src/cmd/experiment/brute_force_words.rs +++ b/crates/dtmt/src/cmd/experiment/brute_force_words.rs @@ -1,11 +1,14 @@ +use std::collections::HashSet; use std::path::PathBuf; use clap::{value_parser, Arg, ArgAction, ArgMatches, Command}; use color_eyre::eyre::{self, Context}; use color_eyre::Result; use itertools::Itertools; +use sdk::murmur::Murmur64; use tokio::fs; -use tokio::io::{AsyncWriteExt, BufWriter}; +use tokio::io::AsyncWriteExt; +use tokio::time::Instant; pub(crate) fn command_definition() -> Command { Command::new("brute-force-words") @@ -60,6 +63,15 @@ pub(crate) fn command_definition() -> Command { .required(true) .value_parser(value_parser!(PathBuf)), ) + .arg( + Arg::new("hashes") + .help( + "Path to a file containing the hashes to attempt to brute force. \ + Hashes are expected in hexadecimal notiation. \ + Only 64-bit hashes are supported." + ) + .value_parser(value_parser!(PathBuf)), + ) } #[tracing::instrument(skip_all)] @@ -86,6 +98,25 @@ pub(crate) async fn run(_ctx: sdk::Context, matches: &ArgMatches) -> Result<()> eyre::bail!("Word list must not be empty"); } + let hashes = if let Some(path) = matches.get_one::("hashes") { + let content = fs::read_to_string(&path) + .await + .wrap_err_with(|| format!("Failed to read file '{}'", path.display()))?; + + let hashes: Result, _> = content + .lines() + .map(|s| u64::from_str_radix(s, 16).map(Murmur64::from)) + .collect(); + + let hashes = hashes?; + + tracing::trace!("{:?}", hashes); + + Some(hashes) + } else { + None + }; + let mut delimiters: Vec = matches .get_many::("delimiter") .unwrap_or_default() @@ -163,8 +194,6 @@ pub(crate) async fn run(_ctx: sdk::Context, matches: &ArgMatches) -> Result<()> tracing::debug!("{:?}", delimiter_lists); - let mut count = 0u64; - let mut indices = if let Some(cont) = matches.get_one::("continue").cloned() { let mut splits = vec![cont.clone()]; @@ -204,7 +233,12 @@ pub(crate) async fn run(_ctx: sdk::Context, matches: &ArgMatches) -> Result<()> indices.reserve(max_length); sequence.reserve(max_length); - let mut writer = BufWriter::new(tokio::io::stdout()); + let mut count: usize = 0; + let mut found: usize = 0; + let mut start = Instant::now(); + + // let mut writer = BufWriter::new(tokio::io::stdout()); + let mut writer = tokio::io::stdout(); let mut buf = Vec::with_capacity(1024); const LINE_FEED: u8 = 0x0A; @@ -216,13 +250,6 @@ pub(crate) async fn run(_ctx: sdk::Context, matches: &ArgMatches) -> Result<()> // one. let delimiter_count = sequence.len() as u32 - 1; - tracing::trace!( - "{} | {:?} -> {:?}", - delimiters_len.pow(delimiter_count), - indices, - sequence - ); - for delims in delimiter_lists .iter() .take(delimiters_len.pow(delimiter_count)) @@ -233,16 +260,25 @@ pub(crate) async fn run(_ctx: sdk::Context, matches: &ArgMatches) -> Result<()> .take(delimiter_count as usize); let s = sequence.iter().copied().interleave(delims.clone()); - count = count.wrapping_add(1); - buf.clear(); for prefix in prefixes.iter() { buf.extend_from_slice(prefix.as_bytes()); s.clone() .for_each(|word| buf.extend_from_slice(word.as_bytes())); - // buf.extend_from_slice(s.as_bytes()); - buf.push(LINE_FEED); + + if let Some(hashes) = &hashes { + let hash = Murmur64::hash(&buf); + if hashes.contains(&hash) { + found += 1; + buf.push(LINE_FEED); + writer.write_all(&buf).await?; + } + + buf.clear(); + } else { + buf.push(LINE_FEED); + } for i in 0..=9 { buf.extend_from_slice(prefix.as_bytes()); @@ -250,7 +286,19 @@ pub(crate) async fn run(_ctx: sdk::Context, matches: &ArgMatches) -> Result<()> .for_each(|word| buf.extend_from_slice(word.as_bytes())); buf.push(UNDERSCORE); buf.push(ZERO + i); - buf.push(LINE_FEED); + + if let Some(hashes) = &hashes { + let hash = Murmur64::hash(&buf); + if hashes.contains(&hash) { + found += 1; + buf.push(LINE_FEED); + writer.write_all(&buf).await?; + } + + buf.clear(); + } else { + buf.push(LINE_FEED); + } buf.extend_from_slice(prefix.as_bytes()); s.clone() @@ -258,11 +306,47 @@ pub(crate) async fn run(_ctx: sdk::Context, matches: &ArgMatches) -> Result<()> buf.push(UNDERSCORE); buf.push(ZERO); buf.push(ZERO + i); - buf.push(LINE_FEED); + + if let Some(hashes) = &hashes { + let hash = Murmur64::hash(&buf); + if hashes.contains(&hash) { + found += 1; + buf.push(LINE_FEED); + writer.write_all(&buf).await?; + } + + buf.clear(); + } else { + buf.push(LINE_FEED); + } } } - writer.write_all(&buf).await?; + if let Some(hashes) = &hashes { + count += prefixes.len() * 20; + + let dur = Instant::now() - start; + if dur.as_secs() >= 1 { + let hashes_len = hashes.len(); + // Don't care when it finishes, don't care if it fails. + tokio::spawn(async move { + let _ = tokio::io::stderr() + .write_all( + format!( + "\r{} hashes per second, {}/{} found", + count, found, hashes_len + ) + .as_bytes(), + ) + .await; + }); + + start = Instant::now(); + count = 0; + } + } else { + writer.write_all(&buf).await?; + } } for i in 0..indices_len {