From 561a4d220bdcaf61c50cd31df028dcf89e21dc71 Mon Sep 17 00:00:00 2001 From: Niklas Elsbrock Date: Sat, 25 Oct 2025 13:53:38 +0200 Subject: [PATCH] add `--stop-after` option --- src/lib.rs | 95 ++++++++++++++++++++++++++++++++++++----------------- src/main.rs | 10 +++++- 2 files changed, 74 insertions(+), 31 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 98239f4..b32f94b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,10 +2,11 @@ use std::{ fmt::{self, Write}, - io::{self}, + io, + num::NonZeroU64, sync::{ - atomic::{AtomicU64, Ordering}, - mpsc::{self, Receiver}, + atomic::{AtomicBool, AtomicU64, Ordering}, + mpsc, }, thread, time::{Duration, Instant}, @@ -31,12 +32,15 @@ type SecretKey = Key; pub struct Config { pub regex: Regex, pub status_enabled: bool, + pub stop_after: Option, pub password: Option, } #[derive(Debug)] pub struct Fingerprunk { config: Config, + started_instant: Instant, + stop: AtomicBool, counter_tried: AtomicU64, counter_found: AtomicU64, } @@ -52,12 +56,16 @@ impl Fingerprunk { pub fn new_from_config(config: Config) -> Self { Self { config, + started_instant: Instant::now(), + stop: AtomicBool::new(false), counter_tried: AtomicU64::new(0), counter_found: AtomicU64::new(0), } } - pub fn run(self) { + pub fn run(mut self) { + self.started_instant = Instant::now(); + let (tx, rx) = mpsc::channel(); thread::scope(|scope| { @@ -65,12 +73,16 @@ impl Fingerprunk { let ref_self = &self; - if self.config.status_enabled { - thread::Builder::new() - .name("status_displayer".to_string()) - .spawn_scoped(scope, move || ref_self.status_displayer_thread()) - .expect(THREAD_SPAWN_EXPECT_MSG); - } + let status_displayer = if self.config.status_enabled { + Some( + thread::Builder::new() + .name("status_displayer".to_string()) + .spawn_scoped(scope, move || ref_self.status_displayer_thread()) + .expect(THREAD_SPAWN_EXPECT_MSG), + ) + } else { + None + }; for num in 0..num_cpus::get() { let tx = tx.clone(); @@ -81,9 +93,19 @@ impl Fingerprunk { .expect(THREAD_SPAWN_EXPECT_MSG); } + let on_stop = || { + // Ask all other threads to stop + self.stop.store(true, Ordering::Relaxed); + + // Unpark the status displayer thread + if let Some(status_displayer) = status_displayer { + status_displayer.thread().unpark(); + } + }; + thread::Builder::new() .name("finalizer".to_string()) - .spawn_scoped(scope, move || ref_self.finalizer_thread(rx)) + .spawn_scoped(scope, move || ref_self.finalizer_thread(rx, on_stop)) .expect(THREAD_SPAWN_EXPECT_MSG); }); } @@ -91,7 +113,7 @@ impl Fingerprunk { fn worker_thread(&self, matches_tx: mpsc::Sender) { let mut fingerprint_hex = String::with_capacity(20 * 2); - loop { + while !self.stop.load(Ordering::Relaxed) { let key = Key4::generate_ecc(true, Curve::Ed25519).expect("should be able to generate key"); fingerprint_hex.clear(); @@ -114,7 +136,7 @@ impl Fingerprunk { .expect("should check regex without error") } - fn finalizer_thread(&self, matches_rx: Receiver) { + fn finalizer_thread(&self, matches_rx: mpsc::Receiver, on_stop: impl FnOnce()) { let mut stdout = io::stdout().lock(); for key in matches_rx { @@ -125,8 +147,14 @@ impl Fingerprunk { self.serialize_cert(cert, &mut stdout) .expect("should be able to serialize certificate"); - self.counter_found.fetch_add(1, Ordering::Relaxed); + let prev = self.counter_found.fetch_add(1, Ordering::Relaxed); + + if self.config.stop_after.is_some_and(|s| prev + 1 == s.get()) { + break; + } } + + on_stop(); } fn key_to_cert(&self, key: &SecretKey) -> anyhow::Result { @@ -181,6 +209,21 @@ impl Fingerprunk { } fn status_displayer_thread(&self) { + const UPDATE_INTERVAL: Duration = Duration::from_millis(250); + + eprint!("\n\n\n\n\n"); + + while !self.stop.load(Ordering::Relaxed) { + self.print_status(); + // We are parking the thread instead of sleeping so we can unpark it when we want to + // stop the program. + thread::park_timeout(UPDATE_INTERVAL); + } + + self.print_status(); + } + + fn print_status(&self) { struct DurationDhms(Duration); impl fmt::Display for DurationDhms { @@ -194,28 +237,20 @@ impl Fingerprunk { } } - const UPDATE_INTERVAL: Duration = Duration::from_millis(250); const FORMAT_WIDTH: usize = 12; - let start = Instant::now(); - - eprint!("\n\n\n\n\n"); - - loop { - let duration = DurationDhms(start.elapsed()); - let keys = self.counter_tried.load(Ordering::Relaxed); - let keys_per_sec = keys as f64 / duration.0.as_secs_f64(); - let found = self.counter_found.load(Ordering::Relaxed); - eprint!( - "\x1b[F\x1b[F\x1b[F\x1b[F\x1b[F\ + let duration = DurationDhms(self.started_instant.elapsed()); + let keys = self.counter_tried.load(Ordering::Relaxed); + let keys_per_sec = keys as f64 / duration.0.as_secs_f64(); + let found = self.counter_found.load(Ordering::Relaxed); + eprint!( + "\x1b[F\x1b[F\x1b[F\x1b[F\x1b[F\ Time: {duration}\n\ Tried: {keys: >w$} keys\n\ Rate: {keys_per_sec: >w$.0} keys/s\n\ ---\n\ Found: {found: >w$} keys\n", - w = FORMAT_WIDTH - ); - thread::sleep(UPDATE_INTERVAL); - } + w = FORMAT_WIDTH + ); } } diff --git a/src/main.rs b/src/main.rs index b33b881..6f575de 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,7 @@ -use std::io::{self, IsTerminal}; +use std::{ + io::{self, IsTerminal}, + num::NonZeroU64, +}; use anyhow::{Context as AnyhowContext, anyhow}; use clap::{ArgAction, Parser, ValueEnum}; @@ -26,6 +29,10 @@ struct Args { #[arg(long, value_enum, default_value_t)] status: StatusEnabled, + /// Stop once the specified number of matching keys has been found. + #[arg(long)] + stop_after: Option, + /// Prompt for a password and use it to encrypt found keys. /// /// By default, found keys are printed to stdout unencrypted. Use this if you actually plan to @@ -78,6 +85,7 @@ fn main() -> anyhow::Result<()> { let config = fingerprunk::Config { regex: args.regex, status_enabled: args.status.evaluate(), + stop_after: args.stop_after, password, };