1use crate::error::{Result, TdbError};
38use aes_gcm::{
39 aead::{rand_core::RngCore, Aead, KeyInit, OsRng},
40 Aes256Gcm, Nonce,
41};
42use pbkdf2::pbkdf2_hmac;
43use sha2::Sha256;
44use std::time::SystemTime;
45
46const PBKDF2_ITERATIONS: u32 = 600_000;
48
49const SALT_SIZE: usize = 32;
51
52const NONCE_SIZE: usize = 12;
54
55const KEY_SIZE: usize = 32;
57
58#[derive(Debug, Clone)]
60pub struct EncryptionConfig {
61 password: String,
63 iterations: u32,
65 compress_before_encrypt: bool,
67}
68
69impl EncryptionConfig {
70 pub fn new(password: impl Into<String>) -> Self {
72 Self {
73 password: password.into(),
74 iterations: PBKDF2_ITERATIONS,
75 compress_before_encrypt: true,
76 }
77 }
78
79 pub fn with_iterations(mut self, iterations: u32) -> Self {
81 self.iterations = iterations;
82 self
83 }
84
85 pub fn with_compression(mut self, enable: bool) -> Self {
87 self.compress_before_encrypt = enable;
88 self
89 }
90}
91
92#[derive(Debug, Clone)]
94pub struct EncryptedData {
95 pub salt: Vec<u8>,
97 pub nonce: Vec<u8>,
99 pub ciphertext: Vec<u8>,
101 pub encrypted_at: SystemTime,
103 pub compressed: bool,
105}
106
107impl EncryptedData {
108 pub fn to_bytes(&self) -> Vec<u8> {
110 let mut bytes = Vec::new();
111
112 bytes.extend_from_slice(&(self.salt.len() as u32).to_le_bytes());
114 bytes.extend_from_slice(&self.salt);
115
116 bytes.extend_from_slice(&(self.nonce.len() as u32).to_le_bytes());
117 bytes.extend_from_slice(&self.nonce);
118
119 bytes.push(if self.compressed { 1 } else { 0 });
120
121 bytes.extend_from_slice(&(self.ciphertext.len() as u32).to_le_bytes());
122 bytes.extend_from_slice(&self.ciphertext);
123
124 bytes
125 }
126
127 pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
129 let mut offset = 0;
130
131 if bytes.len() < offset + 4 {
133 return Err(TdbError::Other(
134 "Invalid encrypted data: too short".to_string(),
135 ));
136 }
137 let salt_len = u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap()) as usize;
138 offset += 4;
139
140 if bytes.len() < offset + salt_len {
141 return Err(TdbError::Other(
142 "Invalid encrypted data: salt truncated".to_string(),
143 ));
144 }
145 let salt = bytes[offset..offset + salt_len].to_vec();
146 offset += salt_len;
147
148 if bytes.len() < offset + 4 {
150 return Err(TdbError::Other(
151 "Invalid encrypted data: nonce missing".to_string(),
152 ));
153 }
154 let nonce_len = u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap()) as usize;
155 offset += 4;
156
157 if bytes.len() < offset + nonce_len {
158 return Err(TdbError::Other(
159 "Invalid encrypted data: nonce truncated".to_string(),
160 ));
161 }
162 let nonce = bytes[offset..offset + nonce_len].to_vec();
163 offset += nonce_len;
164
165 if bytes.len() < offset + 1 {
167 return Err(TdbError::Other(
168 "Invalid encrypted data: compressed flag missing".to_string(),
169 ));
170 }
171 let compressed = bytes[offset] != 0;
172 offset += 1;
173
174 if bytes.len() < offset + 4 {
176 return Err(TdbError::Other(
177 "Invalid encrypted data: ciphertext length missing".to_string(),
178 ));
179 }
180 let ciphertext_len =
181 u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap()) as usize;
182 offset += 4;
183
184 if bytes.len() < offset + ciphertext_len {
185 return Err(TdbError::Other(
186 "Invalid encrypted data: ciphertext truncated".to_string(),
187 ));
188 }
189 let ciphertext = bytes[offset..offset + ciphertext_len].to_vec();
190
191 Ok(Self {
192 salt,
193 nonce,
194 ciphertext,
195 encrypted_at: SystemTime::now(),
196 compressed,
197 })
198 }
199}
200
201pub struct BackupEncryption {
203 config: EncryptionConfig,
205}
206
207impl BackupEncryption {
208 pub fn new(config: EncryptionConfig) -> Result<Self> {
210 if config.password.is_empty() {
211 return Err(TdbError::Other("Password cannot be empty".to_string()));
212 }
213
214 if config.password.len() < 8 {
215 return Err(TdbError::Other(
216 "Password must be at least 8 characters".to_string(),
217 ));
218 }
219
220 Ok(Self { config })
221 }
222
223 pub fn encrypt(&mut self, plaintext: &[u8]) -> Result<EncryptedData> {
225 let mut salt = vec![0u8; SALT_SIZE];
227 OsRng.fill_bytes(&mut salt);
228
229 let key = self.derive_key(&salt)?;
231
232 let cipher = Aes256Gcm::new_from_slice(&key)
234 .map_err(|e| TdbError::Other(format!("Failed to create cipher: {}", e)))?;
235
236 let mut nonce_bytes = vec![0u8; NONCE_SIZE];
238 OsRng.fill_bytes(&mut nonce_bytes);
239 let nonce = Nonce::from_slice(&nonce_bytes);
240
241 let data_to_encrypt = if self.config.compress_before_encrypt {
243 self.compress(plaintext)?
244 } else {
245 plaintext.to_vec()
246 };
247
248 let ciphertext = cipher
250 .encrypt(nonce, data_to_encrypt.as_ref())
251 .map_err(|e| TdbError::Other(format!("Encryption failed: {}", e)))?;
252
253 Ok(EncryptedData {
254 salt,
255 nonce: nonce_bytes,
256 ciphertext,
257 encrypted_at: SystemTime::now(),
258 compressed: self.config.compress_before_encrypt,
259 })
260 }
261
262 pub fn decrypt(&self, encrypted: &EncryptedData) -> Result<Vec<u8>> {
264 let key = self.derive_key(&encrypted.salt)?;
266
267 let cipher = Aes256Gcm::new_from_slice(&key)
269 .map_err(|e| TdbError::Other(format!("Failed to create cipher: {}", e)))?;
270
271 let nonce = Nonce::from_slice(&encrypted.nonce);
273
274 let plaintext = cipher
276 .decrypt(nonce, encrypted.ciphertext.as_ref())
277 .map_err(|e| {
278 TdbError::Other(format!(
279 "Decryption failed (wrong password or corrupted data): {}",
280 e
281 ))
282 })?;
283
284 if encrypted.compressed {
286 self.decompress(&plaintext)
287 } else {
288 Ok(plaintext)
289 }
290 }
291
292 fn derive_key(&self, salt: &[u8]) -> Result<Vec<u8>> {
294 let mut key = vec![0u8; KEY_SIZE];
295 pbkdf2_hmac::<Sha256>(
296 self.config.password.as_bytes(),
297 salt,
298 self.config.iterations,
299 &mut key,
300 );
301 Ok(key)
302 }
303
304 fn compress(&self, data: &[u8]) -> Result<Vec<u8>> {
306 Ok(lz4_flex::compress_prepend_size(data))
307 }
308
309 fn decompress(&self, data: &[u8]) -> Result<Vec<u8>> {
311 lz4_flex::decompress_size_prepended(data)
312 .map_err(|e| TdbError::Other(format!("Decompression failed: {}", e)))
313 }
314
315 pub fn encrypt_file(
317 &mut self,
318 input_path: &std::path::Path,
319 output_path: &std::path::Path,
320 ) -> Result<()> {
321 let plaintext = std::fs::read(input_path).map_err(TdbError::Io)?;
322 let encrypted = self.encrypt(&plaintext)?;
323 let bytes = encrypted.to_bytes();
324 std::fs::write(output_path, bytes).map_err(TdbError::Io)?;
325 Ok(())
326 }
327
328 pub fn decrypt_file(
330 &self,
331 input_path: &std::path::Path,
332 output_path: &std::path::Path,
333 ) -> Result<()> {
334 let bytes = std::fs::read(input_path).map_err(TdbError::Io)?;
335 let encrypted = EncryptedData::from_bytes(&bytes)?;
336 let plaintext = self.decrypt(&encrypted)?;
337 std::fs::write(output_path, plaintext).map_err(TdbError::Io)?;
338 Ok(())
339 }
340
341 pub fn change_password(
343 &mut self,
344 encrypted: &EncryptedData,
345 new_password: impl Into<String>,
346 ) -> Result<EncryptedData> {
347 let plaintext = self.decrypt(encrypted)?;
349
350 self.config.password = new_password.into();
352
353 self.encrypt(&plaintext)
355 }
356}
357
358#[cfg(test)]
359mod tests {
360 use super::*;
361 use std::env;
362
363 fn test_config(password: impl Into<String>) -> EncryptionConfig {
365 EncryptionConfig::new(password).with_iterations(1000) }
367
368 #[test]
369 fn test_encryption_roundtrip() {
370 let config = test_config("test-password-123");
371 let mut encryption = BackupEncryption::new(config).unwrap();
372
373 let plaintext = b"Sensitive backup data that needs encryption";
374 let encrypted = encryption.encrypt(plaintext).unwrap();
375 let decrypted = encryption.decrypt(&encrypted).unwrap();
376
377 assert_eq!(plaintext.as_ref(), decrypted.as_slice());
378 }
379
380 #[test]
381 fn test_wrong_password_fails() {
382 let config1 = test_config("password1");
383 let mut encryption1 = BackupEncryption::new(config1).unwrap();
384
385 let config2 = test_config("password2");
386 let encryption2 = BackupEncryption::new(config2).unwrap();
387
388 let plaintext = b"Secret data";
389 let encrypted = encryption1.encrypt(plaintext).unwrap();
390
391 let result = encryption2.decrypt(&encrypted);
393 assert!(result.is_err());
394 }
395
396 #[test]
397 fn test_serialization_roundtrip() {
398 let config = test_config("serialize-test");
399 let mut encryption = BackupEncryption::new(config).unwrap();
400
401 let plaintext = b"Data to serialize";
402 let encrypted = encryption.encrypt(plaintext).unwrap();
403
404 let bytes = encrypted.to_bytes();
406 let deserialized = EncryptedData::from_bytes(&bytes).unwrap();
407
408 let decrypted = encryption.decrypt(&deserialized).unwrap();
410
411 assert_eq!(plaintext.as_ref(), decrypted.as_slice());
412 }
413
414 #[test]
415 fn test_file_encryption() {
416 let temp_dir = env::temp_dir().join("oxirs_tdb_encryption_test");
417 std::fs::remove_dir_all(&temp_dir).ok();
418 std::fs::create_dir_all(&temp_dir).unwrap();
419
420 let plain_file = temp_dir.join("plaintext.dat");
421 let encrypted_file = temp_dir.join("encrypted.dat");
422 let decrypted_file = temp_dir.join("decrypted.dat");
423
424 let plaintext = b"Confidential backup data for encryption test";
426 std::fs::write(&plain_file, plaintext).unwrap();
427
428 let config = test_config("file-encryption-key");
429 let mut encryption = BackupEncryption::new(config).unwrap();
430
431 encryption
433 .encrypt_file(&plain_file, &encrypted_file)
434 .unwrap();
435
436 let encrypted_content = std::fs::read(&encrypted_file).unwrap();
438 assert_ne!(plaintext.as_ref(), encrypted_content.as_slice());
439
440 encryption
442 .decrypt_file(&encrypted_file, &decrypted_file)
443 .unwrap();
444
445 let decrypted_content = std::fs::read(&decrypted_file).unwrap();
447 assert_eq!(plaintext.as_ref(), decrypted_content.as_slice());
448
449 std::fs::remove_dir_all(&temp_dir).ok();
450 }
451
452 #[test]
453 fn test_compression_before_encryption() {
454 let config = test_config("compress-test").with_compression(true);
455 let mut encryption = BackupEncryption::new(config).unwrap();
456
457 let plaintext = vec![b'A'; 1000];
459 let encrypted = encryption.encrypt(&plaintext).unwrap();
460
461 assert!(encrypted.compressed);
462
463 assert!(encrypted.ciphertext.len() < plaintext.len());
465
466 let decrypted = encryption.decrypt(&encrypted).unwrap();
468 assert_eq!(plaintext, decrypted);
469 }
470
471 #[test]
472 fn test_password_change() {
473 let config = test_config("old-password");
474 let mut encryption = BackupEncryption::new(config).unwrap();
475
476 let plaintext = b"Data encrypted with old password";
477 let encrypted_old = encryption.encrypt(plaintext).unwrap();
478
479 let encrypted_new = encryption
481 .change_password(&encrypted_old, "new-password")
482 .unwrap();
483
484 let config_old = test_config("old-password");
486 let encryption_old = BackupEncryption::new(config_old).unwrap();
487 let result = encryption_old.decrypt(&encrypted_new);
488 assert!(result.is_err());
489
490 let decrypted = encryption.decrypt(&encrypted_new).unwrap();
492 assert_eq!(plaintext.as_ref(), decrypted.as_slice());
493 }
494
495 #[test]
496 fn test_weak_password_rejected() {
497 let config = EncryptionConfig::new("short");
498 let result = BackupEncryption::new(config);
499 assert!(result.is_err());
500 }
501
502 #[test]
503 fn test_empty_password_rejected() {
504 let config = EncryptionConfig::new("");
505 let result = BackupEncryption::new(config);
506 assert!(result.is_err());
507 }
508
509 #[test]
510 fn test_multiple_encryptions_unique() {
511 let config = test_config("uniqueness-test");
512 let mut encryption = BackupEncryption::new(config).unwrap();
513
514 let plaintext = b"Same plaintext";
515 let encrypted1 = encryption.encrypt(plaintext).unwrap();
516 let encrypted2 = encryption.encrypt(plaintext).unwrap();
517
518 assert_ne!(encrypted1.salt, encrypted2.salt);
520 assert_ne!(encrypted1.nonce, encrypted2.nonce);
521 assert_ne!(encrypted1.ciphertext, encrypted2.ciphertext);
522
523 let decrypted1 = encryption.decrypt(&encrypted1).unwrap();
525 let decrypted2 = encryption.decrypt(&encrypted2).unwrap();
526 assert_eq!(decrypted1, decrypted2);
527 assert_eq!(plaintext.as_ref(), decrypted1.as_slice());
528 }
529
530 #[test]
531 fn test_large_data_encryption() {
532 let config = test_config("large-data-test");
533 let mut encryption = BackupEncryption::new(config).unwrap();
534
535 let plaintext = vec![0xAB; 1024 * 1024];
537 let encrypted = encryption.encrypt(&plaintext).unwrap();
538 let decrypted = encryption.decrypt(&encrypted).unwrap();
539
540 assert_eq!(plaintext, decrypted);
541 }
542
543 #[test]
544 fn test_custom_iterations() {
545 let config = test_config("iterations-test").with_iterations(10_000); let mut encryption = BackupEncryption::new(config).unwrap();
547
548 let plaintext = b"Test with custom iterations";
549 let encrypted = encryption.encrypt(plaintext).unwrap();
550 let decrypted = encryption.decrypt(&encrypted).unwrap();
551
552 assert_eq!(plaintext.as_ref(), decrypted.as_slice());
553 }
554
555 #[test]
556 fn test_encryption_without_compression() {
557 let config = test_config("no-compression").with_compression(false);
558 let mut encryption = BackupEncryption::new(config).unwrap();
559
560 let plaintext = vec![b'A'; 1000];
561 let encrypted = encryption.encrypt(&plaintext).unwrap();
562
563 assert!(!encrypted.compressed);
564
565 let decrypted = encryption.decrypt(&encrypted).unwrap();
566 assert_eq!(plaintext, decrypted);
567 }
568}