1use aes_gcm::{
4 Aes256Gcm, Nonce as AesNonce,
5 aead::{Aead, KeyInit},
6};
7use argon2::{Algorithm as Argon2Algorithm, Argon2, Params, Version};
8use chacha20poly1305::{ChaCha20Poly1305, Nonce as ChaChaNonce};
9use rand::RngCore;
10use zeroize::Zeroizing;
11
12use crate::crypto::EncryptionAlgorithm;
13use crate::error::AgentError;
14
15pub const TAG_LEN: usize = 1;
17pub const NONCE_LEN: usize = 12; pub const SALT_LEN: usize = 16;
21pub const SYMMETRIC_KEY_LEN: usize = 32;
23
24pub const ARGON2_TAG: u8 = 3;
26const ARGON2_PARAMS_LEN: usize = 12;
28
29pub fn get_kdf_params() -> Result<Params, AgentError> {
46 #[cfg(not(any(test, feature = "test-utils")))]
47 let params = Params::new(65536, 3, 1, Some(SYMMETRIC_KEY_LEN));
48 #[cfg(any(test, feature = "test-utils"))]
49 let params = Params::new(8, 1, 1, Some(SYMMETRIC_KEY_LEN));
50 params.map_err(|e| AgentError::CryptoError(format!("Invalid Argon2 params: {}", e)))
51}
52
53pub fn validate_passphrase(passphrase: &str) -> Result<(), AgentError> {
58 if passphrase.len() < 12 {
59 return Err(AgentError::WeakPassphrase(format!(
60 "Passphrase must be at least 12 characters (got {})",
61 passphrase.len()
62 )));
63 }
64
65 let has_lower = passphrase.chars().any(|c| c.is_ascii_lowercase());
66 let has_upper = passphrase.chars().any(|c| c.is_ascii_uppercase());
67 let has_digit = passphrase.chars().any(|c| c.is_ascii_digit());
68 let has_symbol = passphrase.chars().any(|c| {
69 c.is_ascii_punctuation() || (c.is_ascii() && !c.is_ascii_alphanumeric() && c != ' ')
70 });
71
72 let class_count = has_lower as u8 + has_upper as u8 + has_digit as u8 + has_symbol as u8;
73
74 if class_count < 3 {
75 return Err(AgentError::WeakPassphrase(format!(
76 "Passphrase must contain at least 3 of 4 character classes \
77 (lowercase, uppercase, digit, symbol); found {}",
78 class_count
79 )));
80 }
81
82 Ok(())
83}
84
85pub fn encrypt_bytes(
89 data: &[u8],
90 passphrase: &str,
91 algo: EncryptionAlgorithm,
92) -> Result<Vec<u8>, AgentError> {
93 validate_passphrase(passphrase)?;
94
95 let mut salt = [0u8; SALT_LEN];
96 rand::rngs::OsRng.fill_bytes(&mut salt);
97
98 let params = get_kdf_params()?;
99 let m_cost = params.m_cost();
100 let t_cost = params.t_cost();
101 let p_cost = params.p_cost();
102 let argon2 = Argon2::new(Argon2Algorithm::Argon2id, Version::V0x13, params);
103 let mut key = Zeroizing::new([0u8; SYMMETRIC_KEY_LEN]);
104 argon2
105 .hash_password_into(passphrase.as_bytes(), &salt, &mut *key)
106 .map_err(|e| AgentError::CryptoError(format!("Argon2 key derivation failed: {}", e)))?;
107
108 let mut nonce = [0u8; NONCE_LEN];
109 rand::rngs::OsRng.fill_bytes(&mut nonce);
110
111 let ciphertext = match algo {
112 EncryptionAlgorithm::AesGcm256 => {
113 let cipher = Aes256Gcm::new_from_slice(&*key)
114 .map_err(|_| AgentError::CryptoError("Invalid AES key".into()))?;
115 cipher
116 .encrypt(AesNonce::from_slice(&nonce), data)
117 .map_err(|_| AgentError::CryptoError("AES encryption failed".into()))?
118 }
119 EncryptionAlgorithm::ChaCha20Poly1305 => {
120 let cipher = ChaCha20Poly1305::new_from_slice(&*key)
121 .map_err(|_| AgentError::CryptoError("Invalid ChaCha key".into()))?;
122 cipher
123 .encrypt(ChaChaNonce::from_slice(&nonce), data)
124 .map_err(|_| AgentError::CryptoError("ChaCha encryption failed".into()))?
125 }
126 };
127
128 let mut out = Vec::with_capacity(
129 TAG_LEN + SALT_LEN + ARGON2_PARAMS_LEN + TAG_LEN + NONCE_LEN + ciphertext.len(),
130 );
131 out.push(ARGON2_TAG);
132 out.extend_from_slice(&salt);
133 out.extend_from_slice(&m_cost.to_le_bytes());
134 out.extend_from_slice(&t_cost.to_le_bytes());
135 out.extend_from_slice(&p_cost.to_le_bytes());
136 out.push(algo.tag());
137 out.extend_from_slice(&nonce);
138 out.extend_from_slice(&ciphertext);
139 Ok(out)
140}
141
142pub fn decrypt_bytes(encrypted: &[u8], passphrase: &str) -> Result<Vec<u8>, AgentError> {
149 if encrypted.is_empty() {
150 return Err(AgentError::CryptoError("Encrypted data too short".into()));
151 }
152
153 let tag = encrypted[0];
154
155 if tag == ARGON2_TAG {
156 return decrypt_bytes_argon2(encrypted, passphrase);
157 }
158
159 if tag == 1 || tag == 2 {
161 return Err(AgentError::CryptoError(
162 "This key was encrypted with a legacy format (HKDF). \
163 Re-encrypt it with: auths key migrate"
164 .into(),
165 ));
166 }
167
168 Err(AgentError::CryptoError(format!(
169 "Unknown encryption tag: {}",
170 tag
171 )))
172}
173
174fn decrypt_bytes_argon2(encrypted: &[u8], passphrase: &str) -> Result<Vec<u8>, AgentError> {
178 const MIN_LEN: usize = TAG_LEN + SALT_LEN + ARGON2_PARAMS_LEN + TAG_LEN + NONCE_LEN;
179 if encrypted.len() < MIN_LEN {
180 return Err(AgentError::CryptoError(
181 "Argon2id encrypted data too short".into(),
182 ));
183 }
184
185 let mut offset = TAG_LEN; let salt = &encrypted[offset..offset + SALT_LEN];
188 offset += SALT_LEN;
189
190 let m_cost = u32::from_le_bytes(
191 encrypted[offset..offset + 4]
192 .try_into()
193 .map_err(|_| AgentError::CryptoError("invalid m_cost bytes".into()))?,
194 );
195 offset += 4;
196 let t_cost = u32::from_le_bytes(
197 encrypted[offset..offset + 4]
198 .try_into()
199 .map_err(|_| AgentError::CryptoError("invalid t_cost bytes".into()))?,
200 );
201 offset += 4;
202 let p_cost = u32::from_le_bytes(
203 encrypted[offset..offset + 4]
204 .try_into()
205 .map_err(|_| AgentError::CryptoError("invalid p_cost bytes".into()))?,
206 );
207 offset += 4;
208
209 let algo_tag = encrypted[offset];
210 offset += 1;
211 let algo = EncryptionAlgorithm::from_tag(algo_tag)
212 .ok_or_else(|| AgentError::CryptoError(format!("Unknown encryption tag: {}", algo_tag)))?;
213
214 let nonce = &encrypted[offset..offset + NONCE_LEN];
215 offset += NONCE_LEN;
216
217 let ciphertext = &encrypted[offset..];
218
219 let params = Params::new(m_cost, t_cost, p_cost, Some(SYMMETRIC_KEY_LEN))
220 .map_err(|e| AgentError::CryptoError(format!("Invalid Argon2 params: {}", e)))?;
221 let argon2 = Argon2::new(Argon2Algorithm::Argon2id, Version::V0x13, params);
222 let mut key = Zeroizing::new([0u8; SYMMETRIC_KEY_LEN]);
223 argon2
224 .hash_password_into(passphrase.as_bytes(), salt, &mut *key)
225 .map_err(|e| AgentError::CryptoError(format!("Argon2 key derivation failed: {}", e)))?;
226
227 let result = match algo {
228 EncryptionAlgorithm::AesGcm256 => Aes256Gcm::new_from_slice(&*key)
229 .map_err(|_| AgentError::CryptoError("Invalid AES key".into()))?
230 .decrypt(AesNonce::from_slice(nonce), ciphertext),
231
232 EncryptionAlgorithm::ChaCha20Poly1305 => ChaCha20Poly1305::new_from_slice(&*key)
233 .map_err(|_| AgentError::CryptoError("Invalid ChaCha key".into()))?
234 .decrypt(ChaChaNonce::from_slice(nonce), ciphertext),
235 };
236
237 match result {
238 Ok(plaintext) => Ok(plaintext),
239 Err(_) => Err(AgentError::IncorrectPassphrase),
240 }
241}
242
243#[cfg(test)]
244mod tests {
245 use super::*;
246 use crate::crypto::EncryptionAlgorithm;
247
248 const STRONG_PASS: &str = "MyStr0ng!Pass";
249
250 #[test]
251 fn test_argon2_roundtrip_aes() {
252 let data = b"hello argon2 aes";
253 let encrypted = encrypt_bytes(data, STRONG_PASS, EncryptionAlgorithm::AesGcm256).unwrap();
254 let decrypted = decrypt_bytes(&encrypted, STRONG_PASS).unwrap();
255 assert_eq!(data.as_slice(), decrypted.as_slice());
256 }
257
258 #[test]
259 fn test_argon2_roundtrip_chacha() {
260 let data = b"hello argon2 chacha";
261 let encrypted =
262 encrypt_bytes(data, STRONG_PASS, EncryptionAlgorithm::ChaCha20Poly1305).unwrap();
263 let decrypted = decrypt_bytes(&encrypted, STRONG_PASS).unwrap();
264 assert_eq!(data.as_slice(), decrypted.as_slice());
265 }
266
267 #[test]
268 fn test_legacy_hkdf_tag_returns_migration_error() {
269 let mut blob = vec![1u8];
271 blob.extend_from_slice(&[0u8; 64]);
272 let result = decrypt_bytes(&blob, "any-passphrase");
273 match result {
274 Err(AgentError::CryptoError(msg)) => {
275 assert!(msg.contains("legacy format"), "got: {}", msg);
276 assert!(msg.contains("auths key migrate"), "got: {}", msg);
277 }
278 other => panic!(
279 "expected CryptoError with migration message, got: {:?}",
280 other
281 ),
282 }
283 }
284
285 #[test]
286 fn test_legacy_hkdf_tag2_returns_migration_error() {
287 let mut blob = vec![2u8];
289 blob.extend_from_slice(&[0u8; 64]);
290 let result = decrypt_bytes(&blob, "any-passphrase");
291 assert!(matches!(result, Err(AgentError::CryptoError(_))));
292 }
293
294 #[test]
295 fn test_argon2_wrong_passphrase() {
296 let data = b"secret";
297 let encrypted = encrypt_bytes(data, STRONG_PASS, EncryptionAlgorithm::AesGcm256).unwrap();
298 let result = decrypt_bytes(&encrypted, "Wr0ng!Passphrase");
299 assert!(matches!(result, Err(AgentError::IncorrectPassphrase)));
300 }
301
302 #[test]
303 fn test_argon2_blob_starts_with_tag_3() {
304 let data = b"tag check";
305 let encrypted = encrypt_bytes(data, STRONG_PASS, EncryptionAlgorithm::AesGcm256).unwrap();
306 assert_eq!(encrypted[0], ARGON2_TAG);
307 }
308
309 #[test]
310 fn test_unknown_tag_returns_error() {
311 let blob = vec![0xFF; 64];
312 let result = decrypt_bytes(&blob, "irrelevant");
313 assert!(matches!(result, Err(AgentError::CryptoError(_))));
314 }
315
316 #[test]
317 fn test_validate_passphrase_too_short() {
318 let result = validate_passphrase("Short1!");
319 assert!(matches!(result, Err(AgentError::WeakPassphrase(_))));
320 }
321
322 #[test]
323 fn test_validate_passphrase_insufficient_classes() {
324 let result = validate_passphrase("abcdefABCDEFgh");
326 assert!(matches!(result, Err(AgentError::WeakPassphrase(_))));
327 }
328
329 #[test]
330 fn test_validate_passphrase_strong() {
331 assert!(validate_passphrase(STRONG_PASS).is_ok());
332 }
333
334 #[test]
335 fn test_argon2_encrypt_rejects_weak() {
336 let result = encrypt_bytes(b"data", "weak", EncryptionAlgorithm::AesGcm256);
337 assert!(matches!(result, Err(AgentError::WeakPassphrase(_))));
338 }
339}