1use crate::error::SecurityError;
22use crate::{P2PError, Result};
23use argon2::{
24 Algorithm, Argon2, Params, Version,
25 password_hash::{PasswordHasher, SaltString, rand_core::RngCore},
26};
27use saorsa_pqc::{ChaCha20Poly1305Cipher, SymmetricEncryptedMessage, SymmetricKey};
28use hkdf::Hkdf;
30use serde::{Deserialize, Serialize};
31use sha2::Sha256;
32
33const CHACHA_KEY_SIZE: usize = 32;
35
36const SALT_SIZE: usize = 32;
38
39const DEVICE_ARGON2_MEMORY: u32 = 32768; const DEVICE_ARGON2_TIME: u32 = 2;
42const DEVICE_ARGON2_PARALLELISM: u32 = 2;
43
44#[derive(Debug, Clone, Serialize, Deserialize)]
46pub struct EncryptedData {
47 pub encrypted_message: SymmetricEncryptedMessage,
49 pub salt: [u8; SALT_SIZE],
51}
52
53pub fn encrypt_with_device_password(data: &[u8], device_password: &str) -> Result<EncryptedData> {
55 let mut salt = [0u8; SALT_SIZE];
57 let mut rng = rand::thread_rng();
58 rng.fill_bytes(&mut salt);
59
60 let key_bytes = derive_key_from_password(device_password, &salt)?;
62 let symmetric_key = SymmetricKey::from_bytes(key_bytes);
63
64 let cipher = ChaCha20Poly1305Cipher::new(&symmetric_key);
66 let (ciphertext, nonce) = cipher.encrypt(data, None).map_err(|e| {
67 P2PError::Security(SecurityError::EncryptionFailed(
68 format!("ChaCha20Poly1305 encryption failed: {:?}", e).into(),
69 ))
70 })?;
71
72 let encrypted_message = SymmetricEncryptedMessage::new(ciphertext, nonce, None);
73
74 Ok(EncryptedData {
75 encrypted_message,
76 salt,
77 })
78}
79
80pub fn decrypt_with_device_password(
82 encrypted: &EncryptedData,
83 device_password: &str,
84) -> Result<Vec<u8>> {
85 let key_bytes = derive_key_from_password(device_password, &encrypted.salt)?;
87 let symmetric_key = SymmetricKey::from_bytes(key_bytes);
88
89 let cipher = ChaCha20Poly1305Cipher::new(&symmetric_key);
91 let plaintext = cipher
92 .decrypt(
93 &encrypted.encrypted_message.ciphertext,
94 &encrypted.encrypted_message.nonce,
95 None,
96 )
97 .map_err(|e| {
98 P2PError::Security(SecurityError::DecryptionFailed(
99 format!("ChaCha20Poly1305 decryption failed: {:?}", e).into(),
100 ))
101 })?;
102
103 Ok(plaintext)
104}
105
106fn derive_key_from_password(
108 password: &str,
109 salt: &[u8; SALT_SIZE],
110) -> Result<[u8; CHACHA_KEY_SIZE]> {
111 let argon2 = Argon2::new(
113 Algorithm::Argon2id,
114 Version::V0x13,
115 Params::new(
116 DEVICE_ARGON2_MEMORY,
117 DEVICE_ARGON2_TIME,
118 DEVICE_ARGON2_PARALLELISM,
119 Some(CHACHA_KEY_SIZE),
120 )
121 .map_err(|e| {
122 P2PError::Security(SecurityError::InvalidKey(
123 format!("Invalid Argon2 params: {}", e).into(),
124 ))
125 })?,
126 );
127
128 let salt_string = SaltString::encode_b64(salt).map_err(|e| {
130 P2PError::Security(SecurityError::InvalidKey(
131 format!("Failed to encode salt: {}", e).into(),
132 ))
133 })?;
134
135 let hash = argon2
137 .hash_password(password.as_bytes(), &salt_string)
138 .map_err(|e| {
139 P2PError::Security(SecurityError::KeyGenerationFailed(
140 format!("Argon2id key derivation failed: {}", e).into(),
141 ))
142 })?;
143
144 let hash_output = hash.hash.ok_or_else(|| {
145 P2PError::Security(SecurityError::KeyGenerationFailed(
146 "No hash output from Argon2".to_string().into(),
147 ))
148 })?;
149
150 let key_bytes = hash_output.as_bytes();
151 if key_bytes.len() < CHACHA_KEY_SIZE {
152 return Err(P2PError::Security(SecurityError::KeyGenerationFailed(
153 "Insufficient key material from Argon2".to_string().into(),
154 )));
155 }
156
157 let mut result = [0u8; CHACHA_KEY_SIZE];
158 result.copy_from_slice(&key_bytes[..CHACHA_KEY_SIZE]);
159 Ok(result)
160}
161
162pub fn encrypt_with_shared_secret(
164 data: &[u8],
165 shared_secret: &[u8; 32],
166 info: &[u8],
167) -> Result<EncryptedData> {
168 let mut salt = [0u8; SALT_SIZE];
170 let mut rng = rand::thread_rng();
171 rng.fill_bytes(&mut salt);
172
173 let hkdf = Hkdf::<Sha256>::new(Some(&salt), shared_secret);
175 let mut key_bytes = [0u8; CHACHA_KEY_SIZE];
176 hkdf.expand(info, &mut key_bytes).map_err(|e| {
177 P2PError::Security(SecurityError::KeyGenerationFailed(
178 format!("HKDF-SHA3 expansion failed: {:?}", e).into(),
179 ))
180 })?;
181
182 let symmetric_key = SymmetricKey::from_bytes(key_bytes);
183
184 let cipher = ChaCha20Poly1305Cipher::new(&symmetric_key);
186 let (ciphertext, nonce) = cipher.encrypt(data, None).map_err(|e| {
187 P2PError::Security(SecurityError::EncryptionFailed(
188 format!("ChaCha20Poly1305 encryption failed: {:?}", e).into(),
189 ))
190 })?;
191
192 let encrypted_message = SymmetricEncryptedMessage::new(ciphertext, nonce, None);
193
194 Ok(EncryptedData {
195 encrypted_message,
196 salt,
197 })
198}
199
200pub fn decrypt_with_shared_secret(
202 encrypted: &EncryptedData,
203 shared_secret: &[u8; 32],
204 info: &[u8],
205) -> Result<Vec<u8>> {
206 let hkdf = Hkdf::<Sha256>::new(Some(&encrypted.salt), shared_secret);
208 let mut key_bytes = [0u8; CHACHA_KEY_SIZE];
209 hkdf.expand(info, &mut key_bytes).map_err(|e| {
210 P2PError::Security(SecurityError::KeyGenerationFailed(
211 format!("HKDF-SHA3 expansion failed: {:?}", e).into(),
212 ))
213 })?;
214
215 let symmetric_key = SymmetricKey::from_bytes(key_bytes);
216
217 let cipher = ChaCha20Poly1305Cipher::new(&symmetric_key);
219 let plaintext = cipher
220 .decrypt(
221 &encrypted.encrypted_message.ciphertext,
222 &encrypted.encrypted_message.nonce,
223 None,
224 )
225 .map_err(|e| {
226 P2PError::Security(SecurityError::DecryptionFailed(
227 format!("ChaCha20Poly1305 decryption failed: {:?}", e).into(),
228 ))
229 })?;
230
231 Ok(plaintext)
232}
233
234#[cfg(test)]
235mod tests {
236 use super::*;
237
238 #[test]
239 fn test_device_password_encryption() {
240 let data = b"Secret identity data";
241 let password = "MyDevicePassword123!";
242
243 let encrypted =
245 encrypt_with_device_password(data, password).expect("Encryption should succeed");
246
247 assert!(!encrypted.encrypted_message.ciphertext.is_empty());
249
250 let decrypted =
252 decrypt_with_device_password(&encrypted, password).expect("Decryption should succeed");
253
254 assert_eq!(decrypted, data);
255 }
256
257 #[test]
258 fn test_encryption_serialization() {
259 let data = b"Test data for serialization";
260 let password = "SerializeTest123!";
261
262 let encrypted =
264 encrypt_with_device_password(data, password).expect("Encryption should succeed");
265
266 let serialized = bincode::serialize(&encrypted).expect("Serialization should succeed");
268
269 let deserialized: EncryptedData =
271 bincode::deserialize(&serialized).expect("Deserialization should succeed");
272
273 assert_eq!(
275 encrypted.encrypted_message.ciphertext,
276 deserialized.encrypted_message.ciphertext
277 );
278 assert_eq!(encrypted.salt, deserialized.salt);
279
280 let decrypted = decrypt_with_device_password(&deserialized, password)
282 .expect("Decryption should succeed");
283
284 assert_eq!(decrypted, data);
285 }
286
287 #[test]
288 fn test_wrong_password_fails() {
289 let data = b"Secret identity data";
290 let password = "MyDevicePassword123!";
291 let wrong_password = "WrongPassword456!";
292
293 let encrypted =
295 encrypt_with_device_password(data, password).expect("Encryption should succeed");
296
297 let result = decrypt_with_device_password(&encrypted, wrong_password);
299
300 assert!(result.is_err());
301 }
302
303 #[test]
304 fn test_shared_secret_encryption() {
305 let data = b"Peer to peer message";
306 let shared_secret = [42u8; 32];
307 let info = b"p2p-identity-sync";
308
309 let encrypted = encrypt_with_shared_secret(data, &shared_secret, info)
311 .expect("Encryption should succeed");
312
313 let decrypted = decrypt_with_shared_secret(&encrypted, &shared_secret, info)
315 .expect("Decryption should succeed");
316
317 assert_eq!(decrypted, data);
318 }
319
320 #[test]
321 fn test_different_info_fails() {
322 let data = b"Peer to peer message";
323 let shared_secret = [42u8; 32];
324 let info1 = b"p2p-identity-sync";
325 let info2 = b"different-context";
326
327 let encrypted = encrypt_with_shared_secret(data, &shared_secret, info1)
329 .expect("Encryption should succeed");
330
331 let result = decrypt_with_shared_secret(&encrypted, &shared_secret, info2);
333
334 assert!(result.is_err());
335 }
336
337 #[test]
338 fn test_encryption_produces_unique_nonces() {
339 let data = b"Test data";
340 let password = "TestPassword123!";
341
342 let encrypted1 =
344 encrypt_with_device_password(data, password).expect("First encryption should succeed");
345 let encrypted2 =
346 encrypt_with_device_password(data, password).expect("Second encryption should succeed");
347
348 assert_ne!(encrypted1.salt, encrypted2.salt);
350 assert_ne!(
352 encrypted1.encrypted_message.ciphertext,
353 encrypted2.encrypted_message.ciphertext
354 );
355 }
356}