vanity_ssh_rs/core/
pattern.rs

1use ed25519_dalek::VerifyingKey;
2use regex::Regex;
3use ssh_key::private::Ed25519Keypair;
4use std::time::Duration;
5
6pub enum Pattern {
7    Suffix(String),
8    Regex(Regex),
9}
10
11impl Pattern {
12    pub fn new(pattern: String) -> Result<Self, regex::Error> {
13        if pattern.starts_with('/') && pattern.ends_with('/') {
14            let pattern = pattern[1..pattern.len() - 1].to_string();
15            Ok(Pattern::Regex(Regex::new(&pattern)?))
16        } else {
17            Ok(Pattern::Suffix(pattern))
18        }
19    }
20
21    pub fn to_filename(&self) -> String {
22        match self {
23            Pattern::Suffix(suffix) => suffix.clone(),
24            Pattern::Regex(regex) => {
25                let pattern = regex.as_str();
26                // Remove special characters and limit length
27                let clean = pattern
28                    .chars()
29                    .take(20)
30                    .map(|c| if c.is_alphanumeric() { c } else { '_' })
31                    .collect::<String>();
32                format!("regex_{}", clean)
33            }
34        }
35    }
36
37    pub fn probability(&self) -> Option<f64> {
38        match self {
39            Pattern::Suffix(suffix) => {
40                // Base64 has 64 possible characters
41                let base: f64 = 64.0;
42                // Probability is (1/64)^n where n is the length of the suffix
43                Some(1.0 / base.powi(suffix.len() as i32))
44            }
45            Pattern::Regex(_) => None, // Regex patterns are too complex to calculate probability
46        }
47    }
48
49    pub fn estimate_time(&self, keys_per_second: f64) -> Option<String> {
50        self.probability().map(|prob| {
51            let expected_attempts = 1.0 / prob;
52            let seconds = expected_attempts / keys_per_second;
53            let duration = Duration::from_secs_f64(seconds);
54            humantime::format_duration(duration).to_string()
55        })
56    }
57}
58
59pub fn public_key_matches_pattern(public_key: &VerifyingKey, pattern: &Pattern) -> bool {
60    let openssh_pubkey = create_openssh_public_key_from_keypair(public_key);
61    let openssh_pubkey_str = openssh_pubkey.to_string();
62    let base64_part = extract_base64_from_openssh_string(&openssh_pubkey_str);
63
64    match pattern {
65        Pattern::Suffix(suffix) => base64_part.ends_with(suffix),
66        Pattern::Regex(regex) => regex.is_match(base64_part),
67    }
68}
69
70fn create_openssh_public_key_from_keypair(
71    verifying_key: &VerifyingKey,
72) -> ssh_key::public::PublicKey {
73    let mut key_bytes = [0u8; 64];
74    key_bytes[32..].copy_from_slice(&verifying_key.to_bytes());
75
76    let ed25519_keypair = Ed25519Keypair::from_bytes(&key_bytes).unwrap();
77    let openssh_pub = ssh_key::public::Ed25519PublicKey::from(&ed25519_keypair);
78    ssh_key::public::PublicKey::from(openssh_pub)
79}
80
81fn extract_base64_from_openssh_string(openssh_string: &str) -> &str {
82    openssh_string.split_whitespace().nth(1).unwrap_or("")
83}