From 89b5d55f62bcfdc7c13b43981e0725b6860ebad2 Mon Sep 17 00:00:00 2001
From: ThinkChaos <ThinkChaos@users.noreply.github.com>
Date: Fri, 24 Apr 2020 22:21:55 +0200
Subject: [PATCH] Add `--case-sensitive` and allow `--in 0` meaning actual
 prefix

---
 src/bin.rs | 20 +++++++++++++-------
 src/lib.rs | 15 +++++++++++----
 2 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/src/bin.rs b/src/bin.rs
index a5a6f35..c18d430 100644
--- a/src/bin.rs
+++ b/src/bin.rs
@@ -8,12 +8,11 @@ use num_cpus;
 use rayon::prelude::*;
 use wireguard_vanity_lib::trial;
 
-fn estimate_one_trial() -> Duration {
-    let prefix = "prefix";
+fn estimate_one_trial(prefix: &str, end: usize, case_sensitive: bool) -> Duration {
     let start = SystemTime::now();
     const COUNT: u32 = 100;
     (0..COUNT).for_each(|_| {
-        trial(&prefix, 0, 10);
+        trial(&prefix, end, case_sensitive);
     });
     let elapsed = start.elapsed().unwrap();
     elapsed.checked_div(COUNT).unwrap()
@@ -86,11 +85,16 @@ fn main() -> Result<(), Box<dyn Error>> {
         .version("0.3.1")
         .author("Brian Warner <warner@lothar.com>")
         .about("finds Wireguard keypairs with a given string prefix")
+        .arg(
+            Arg::with_name("CASE")
+                .long("case-sensitive")
+                .help("Use case-sensitive matching"),
+        )
         .arg(
             Arg::with_name("RANGE")
                 .long("in")
                 .takes_value(true)
-                .help("NAME must be found within first RANGE chars of pubkey (default: 10)"),
+                .help("NAME must be found within first RANGE chars of pubkey (default: 10, 0 means actual prefix)"),
         )
         .arg(
             Arg::with_name("NAME")
@@ -98,6 +102,7 @@ fn main() -> Result<(), Box<dyn Error>> {
                 .help("string to find near the start of the pubkey"),
         )
         .get_matches();
+    let case_sensitive = matches.is_present("CASE");
     let prefix = matches.value_of("NAME").unwrap().to_ascii_lowercase();
     let len = prefix.len();
     let end: usize = 44.min(match matches.value_of("RANGE") {
@@ -110,6 +115,7 @@ fn main() -> Result<(), Box<dyn Error>> {
             }
         }
     });
+    let end = if end == 0 { len } else { end };
     if end < len {
         return Err(ParseError(format!("range {} is too short for len={}", end, len)).into());
     }
@@ -119,7 +125,7 @@ fn main() -> Result<(), Box<dyn Error>> {
     let mut num = offsets;
     let mut denom = 1u64;
     prefix.chars().for_each(|c| {
-        if c.is_ascii_alphabetic() {
+        if !case_sensitive && c.is_ascii_alphabetic() {
             num *= 2; // letters can match both uppercase and lowercase
         }
         denom *= 64; // base64
@@ -136,7 +142,7 @@ fn main() -> Result<(), Box<dyn Error>> {
     // run at half the speed that this predicts.
 
     if trials_per_key < 2u64.pow(32) {
-        let est = estimate_one_trial();
+        let est = estimate_one_trial(&prefix, end, case_sensitive);
         println!(
             "one trial takes {}, CPU cores available: {}",
             format_time(duration_to_f64(est)),
@@ -162,7 +168,7 @@ fn main() -> Result<(), Box<dyn Error>> {
     // 1M trials takes about 10s on my laptop, so let it run for 1000s
     (0..100_000_000)
         .into_par_iter()
-        .map(|_| trial(&prefix, 0, end))
+        .map(|_| trial(&prefix, end, case_sensitive))
         .filter_map(|r| r)
         .try_for_each(print)?;
     Ok(())
diff --git a/src/lib.rs b/src/lib.rs
index 31138c5..fb95581 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,15 +1,22 @@
 use base64;
 use rand::thread_rng;
+use std::borrow::Cow;
 use x25519_dalek::{PublicKey, StaticSecret};
 
-pub fn trial(prefix: &str, start: usize, end: usize) -> Option<(String, String)> {
+pub fn trial(prefix: &str, end: usize, case_sensitive: bool) -> Option<(String, String)> {
     let mut rng = thread_rng();
     let private = StaticSecret::new(&mut rng);
     let public = PublicKey::from(&private);
+
     let public_b64 = base64::encode(public.as_bytes());
-    if public_b64[start..end]
-        .to_ascii_lowercase()
-        .contains(&prefix)
+
+    let b64_prefix = if case_sensitive {
+        Cow::Borrowed(&public_b64[..end])
+    } else {
+        Cow::Owned(public_b64[..end].to_ascii_lowercase())
+    };
+
+    if b64_prefix.contains(prefix)
     {
         let private_b64 = base64::encode(&private.to_bytes());
         Some((private_b64, public_b64))