keys/
main.rs

1use aes_gcm::{
2    Aes256Gcm, Nonce,
3    aead::{Aead, AeadCore, KeyInit, OsRng},
4};
5use clap::Parser;
6use eyre::{Result, eyre};
7use pbkdf2::pbkdf2_hmac;
8use rand::{Rng, rng};
9use sha2::Sha256;
10
11/// Size of the salt used for key derivation (32 bytes)
12const SALT_SIZE: usize = 32;
13/// Size of the nonce used for AES-GCM encryption (12 bytes)
14const NONCE_SIZE: usize = 12;
15/// Size of the derived encryption key (32 bytes for AES-256)
16const KEY_SIZE: usize = 32;
17/// Number of PBKDF2 iterations for key derivation
18const PBKDF2_ITERATIONS: u32 = 100_000;
19
20/// Command-line interface for the keys utility
21#[derive(Parser, Debug)]
22enum Commands {
23    /// Generate a cryptographically secure random password
24    GeneratePassword,
25    /// Encrypt a private key with a password
26    Encrypt {
27        /// The private key to encrypt (hex format, with or without 0x prefix)
28        #[arg(short, long)]
29        key: String,
30    },
31}
32
33/// Encrypts a private key using AES-256-GCM with PBKDF2 key derivation
34///
35/// # Arguments
36/// * `private_key` - The private key bytes to encrypt
37/// * `pwd` - The password bytes to derive the encryption key from
38///
39/// # Returns
40/// * `Result<Vec<u8>>` - The encrypted data containing salt, nonce, and ciphertext
41///
42/// # Format
43/// The returned encrypted data has the following structure:
44/// - Bytes 0-31: Salt (32 bytes)
45/// - Bytes 32-43: Nonce (12 bytes)
46/// - Bytes 44+: Ciphertext with authentication tag
47///
48/// # Security
49/// - Uses AES-256-GCM for authenticated encryption
50/// - Derives key using PBKDF2-HMAC-SHA256 with 100,000 iterations
51/// - Generates random salt and nonce for each encryption
52fn encrypt_key(private_key: Vec<u8>, pwd: Vec<u8>) -> Result<Vec<u8>> {
53    // Generate random salt for key derivation
54    let mut salt = [0u8; SALT_SIZE];
55    rng().fill(&mut salt);
56
57    // Derive encryption key using PBKDF2 with 100,000 iterations
58    let mut derived_key = [0u8; KEY_SIZE];
59    pbkdf2_hmac::<Sha256>(&pwd, &salt, PBKDF2_ITERATIONS, &mut derived_key);
60
61    // Generate random nonce for AES-GCM
62    let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
63
64    // Create cipher and encrypt
65    let cipher = Aes256Gcm::new_from_slice(&derived_key).map_err(|e| eyre!("Failed to create cipher: {}", e))?;
66
67    let ciphertext = cipher.encrypt(&nonce, private_key.as_ref()).map_err(|e| eyre!("Encryption failed: {}", e))?;
68
69    // Format: salt || nonce || ciphertext (includes auth tag)
70    let mut result = Vec::with_capacity(SALT_SIZE + NONCE_SIZE + ciphertext.len());
71    result.extend_from_slice(&salt);
72    result.extend_from_slice(&nonce);
73    result.extend_from_slice(&ciphertext);
74
75    Ok(result)
76}
77
78/// Decrypts data that was encrypted with `encrypt_key`
79///
80/// # Arguments
81/// * `encrypted_data` - The encrypted data containing salt, nonce, and ciphertext
82/// * `pwd` - The password bytes to derive the decryption key from
83///
84/// # Returns
85/// * `Result<Vec<u8>>` - The decrypted private key bytes
86///
87/// # Security
88/// - Verifies authentication tag to ensure data integrity
89/// - Uses the same PBKDF2 parameters as encryption
90fn decrypt_key(encrypted_data: &[u8], pwd: &[u8]) -> Result<Vec<u8>> {
91    if encrypted_data.len() < SALT_SIZE + NONCE_SIZE {
92        return Err(eyre!("Invalid encrypted data: too short"));
93    }
94
95    // Extract components
96    let salt = &encrypted_data[..SALT_SIZE];
97    let nonce = &encrypted_data[SALT_SIZE..SALT_SIZE + NONCE_SIZE];
98    let ciphertext = &encrypted_data[SALT_SIZE + NONCE_SIZE..];
99
100    // Derive same encryption key using PBKDF2
101    let mut derived_key = [0u8; KEY_SIZE];
102    pbkdf2_hmac::<Sha256>(pwd, salt, PBKDF2_ITERATIONS, &mut derived_key);
103
104    // Create cipher and decrypt
105    let cipher = Aes256Gcm::new_from_slice(&derived_key).map_err(|e| eyre!("Failed to create cipher: {}", e))?;
106
107    let nonce = Nonce::from_slice(nonce);
108    let plaintext = cipher.decrypt(nonce, ciphertext).map_err(|e| eyre!("Decryption failed: {}", e))?;
109
110    Ok(plaintext)
111}
112
113fn main() -> Result<()> {
114    let args = Commands::parse();
115    match args {
116        Commands::GeneratePassword => {
117            // Use cryptographically secure RNG for password generation
118            let mut random = rng();
119            let pwd: Vec<u8> = (0..16).map(|_| random.random::<u8>()).collect();
120            println!("{pwd:?}");
121        }
122        Commands::Encrypt { key } => {
123            let pwd = kabu_types_entities::private::KEY_ENCRYPTION_PWD.to_vec();
124
125            let private_key = hex::decode(key.strip_prefix("0x").unwrap_or(key.clone().as_str()))?;
126            let encrypted_key = encrypt_key(private_key.clone(), pwd.clone())?;
127
128            // Verify decryption with new method
129            let decrypted_key = decrypt_key(&encrypted_key, &pwd)?;
130            if decrypted_key == private_key {
131                println!("Encrypted private key : {}", hex::encode(encrypted_key))
132            } else {
133                println!("Error encrypting private key");
134            }
135        }
136    }
137
138    Ok(())
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144
145    #[test]
146    fn test_encrypt_decrypt_roundtrip() {
147        let private_key = vec![0x42; 32];
148        let password = b"test_password_123!";
149
150        let encrypted = encrypt_key(private_key.clone(), password.to_vec()).unwrap();
151        let decrypted = decrypt_key(&encrypted, password).unwrap();
152
153        assert_eq!(decrypted, private_key);
154    }
155
156    #[test]
157    fn test_wrong_password_fails() {
158        let private_key = vec![0x42; 32];
159        let password = b"correct_password";
160        let wrong_password = b"wrong_password";
161
162        let encrypted = encrypt_key(private_key.clone(), password.to_vec()).unwrap();
163        let result = decrypt_key(&encrypted, wrong_password);
164
165        assert!(result.is_err());
166    }
167
168    #[test]
169    fn test_tampered_data_fails() {
170        let private_key = vec![0x42; 32];
171        let password = b"test_password";
172
173        let mut encrypted = encrypt_key(private_key, password.to_vec()).unwrap();
174        // Tamper with the ciphertext
175        let last_idx = encrypted.len() - 1;
176        encrypted[last_idx] ^= 0xFF;
177
178        let result = decrypt_key(&encrypted, password);
179        assert!(result.is_err());
180    }
181
182    #[test]
183    fn test_encryption_randomness() {
184        let private_key = vec![0xff; 32];
185        let password = b"same_password";
186
187        let encrypted1 = encrypt_key(private_key.clone(), password.to_vec()).unwrap();
188        let encrypted2 = encrypt_key(private_key.clone(), password.to_vec()).unwrap();
189
190        // Should produce different ciphertexts due to random salt/nonce
191        assert_ne!(encrypted1, encrypted2);
192
193        // But both should decrypt to the same private key
194        let decrypted1 = decrypt_key(&encrypted1, password).unwrap();
195        let decrypted2 = decrypt_key(&encrypted2, password).unwrap();
196        assert_eq!(decrypted1, private_key);
197        assert_eq!(decrypted2, private_key);
198    }
199}
keys/main.rs

keys/
main.rs
