rustywallet_vanity/
address_type.rs

1//! Address type definitions and utilities.
2
3use crate::error::PatternError;
4use rustywallet_address::{
5    EthereumAddress, Network as AddrNetwork, P2PKHAddress, P2TRAddress, P2WPKHAddress,
6};
7use rustywallet_keys::prelude::*;
8
9/// Supported address types for vanity generation.
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
11pub enum AddressType {
12    /// Legacy P2PKH address (starts with 1 on mainnet).
13    #[default]
14    P2PKH,
15    /// Native SegWit P2WPKH address (starts with bc1q on mainnet).
16    P2WPKH,
17    /// Taproot P2TR address (starts with bc1p on mainnet).
18    P2TR,
19    /// Ethereum address (starts with 0x).
20    Ethereum,
21}
22
23impl AddressType {
24    /// Get the valid characters for this address type's variable portion.
25    pub fn valid_chars(&self) -> &'static str {
26        match self {
27            // Base58 alphabet (no 0, O, I, l)
28            AddressType::P2PKH => "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz",
29            // Bech32 alphabet (lowercase only, no 1, b, i, o)
30            AddressType::P2WPKH | AddressType::P2TR => "023456789acdefghjklmnpqrstuvwxyz",
31            // Hex alphabet
32            AddressType::Ethereum => "0123456789abcdefABCDEF",
33        }
34    }
35
36    /// Get the fixed prefix for this address type on mainnet.
37    pub fn fixed_prefix(&self, testnet: bool) -> &'static str {
38        match (self, testnet) {
39            (AddressType::P2PKH, false) => "1",
40            (AddressType::P2PKH, true) => "m", // or 'n'
41            (AddressType::P2WPKH, false) => "bc1q",
42            (AddressType::P2WPKH, true) => "tb1q",
43            (AddressType::P2TR, false) => "bc1p",
44            (AddressType::P2TR, true) => "tb1p",
45            (AddressType::Ethereum, _) => "0x",
46        }
47    }
48
49    /// Validate that a pattern is compatible with this address type.
50    pub fn validate_pattern(&self, pattern: &str, testnet: bool) -> Result<(), PatternError> {
51        if pattern.is_empty() {
52            return Err(PatternError::EmptyPattern);
53        }
54
55        let fixed_prefix = self.fixed_prefix(testnet);
56        let valid_chars = self.valid_chars();
57
58        // Check if pattern conflicts with fixed prefix
59        if pattern.len() <= fixed_prefix.len() {
60            // Pattern must match the beginning of fixed prefix
61            let prefix_start = &fixed_prefix[..pattern.len().min(fixed_prefix.len())];
62            if !prefix_start.eq_ignore_ascii_case(pattern) && !pattern.starts_with(prefix_start) {
63                return Err(PatternError::ConflictsWithPrefix(
64                    pattern.to_string(),
65                    fixed_prefix.to_string(),
66                ));
67            }
68        }
69
70        // For patterns longer than fixed prefix, validate remaining chars
71        if pattern.len() > fixed_prefix.len() {
72            let variable_part = &pattern[fixed_prefix.len()..];
73            for c in variable_part.chars() {
74                // For case-insensitive matching, check both cases
75                let c_lower = c.to_ascii_lowercase();
76                let c_upper = c.to_ascii_uppercase();
77                if !valid_chars.contains(c_lower) && !valid_chars.contains(c_upper) {
78                    return Err(PatternError::InvalidCharacter(c));
79                }
80            }
81        }
82
83        // Warn about very long patterns
84        if pattern.len() > fixed_prefix.len() + 8 {
85            return Err(PatternError::PatternTooLong(
86                pattern.len() - fixed_prefix.len(),
87            ));
88        }
89
90        Ok(())
91    }
92
93    /// Derive an address from a private key.
94    pub fn derive_address(&self, key: &PrivateKey, testnet: bool) -> Result<String, String> {
95        let pubkey = key.public_key();
96        let network = if testnet {
97            AddrNetwork::BitcoinTestnet
98        } else {
99            AddrNetwork::BitcoinMainnet
100        };
101
102        match self {
103            AddressType::P2PKH => P2PKHAddress::from_public_key(&pubkey, network)
104                .map(|a| a.to_string())
105                .map_err(|e| e.to_string()),
106            AddressType::P2WPKH => P2WPKHAddress::from_public_key(&pubkey, network)
107                .map(|a| a.to_string())
108                .map_err(|e| e.to_string()),
109            AddressType::P2TR => P2TRAddress::from_public_key(&pubkey, network)
110                .map(|a| a.to_string())
111                .map_err(|e| e.to_string()),
112            AddressType::Ethereum => EthereumAddress::from_public_key(&pubkey)
113                .map(|a| a.to_checksum_string())
114                .map_err(|e| e.to_string()),
115        }
116    }
117}
118
119impl std::fmt::Display for AddressType {
120    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
121        match self {
122            AddressType::P2PKH => write!(f, "P2PKH"),
123            AddressType::P2WPKH => write!(f, "P2WPKH"),
124            AddressType::P2TR => write!(f, "P2TR"),
125            AddressType::Ethereum => write!(f, "Ethereum"),
126        }
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133
134    #[test]
135    fn test_fixed_prefixes() {
136        assert_eq!(AddressType::P2PKH.fixed_prefix(false), "1");
137        assert_eq!(AddressType::P2WPKH.fixed_prefix(false), "bc1q");
138        assert_eq!(AddressType::P2TR.fixed_prefix(false), "bc1p");
139        assert_eq!(AddressType::Ethereum.fixed_prefix(false), "0x");
140    }
141
142    #[test]
143    fn test_validate_pattern_p2pkh() {
144        // Valid patterns
145        assert!(AddressType::P2PKH.validate_pattern("1Love", false).is_ok());
146        assert!(AddressType::P2PKH.validate_pattern("1BTC", false).is_ok());
147
148        // Invalid character (0 not in Base58)
149        assert!(AddressType::P2PKH.validate_pattern("10", false).is_err());
150    }
151
152    #[test]
153    fn test_validate_pattern_bech32() {
154        // Valid patterns
155        assert!(AddressType::P2WPKH
156            .validate_pattern("bc1qtest", false)
157            .is_ok());
158
159        // Invalid - uppercase not allowed in bech32 variable part
160        // Actually bech32 is case-insensitive but typically lowercase
161    }
162
163    #[test]
164    fn test_derive_address() {
165        let key = PrivateKey::random();
166
167        let p2pkh = AddressType::P2PKH.derive_address(&key, false).unwrap();
168        assert!(p2pkh.starts_with('1'));
169
170        let p2wpkh = AddressType::P2WPKH.derive_address(&key, false).unwrap();
171        assert!(p2wpkh.starts_with("bc1q"));
172
173        let eth = AddressType::Ethereum.derive_address(&key, false).unwrap();
174        assert!(eth.starts_with("0x"));
175    }
176}