1#[cfg(feature = "encryption")]
2use aes_gcm::{
3 aead::{Aead, KeyInit, OsRng},
4 Aes256Gcm, Nonce,
5};
6
7use crate::error::{PackError, Result};
8
9#[cfg(feature = "encryption")]
10#[derive(Clone)]
11pub struct EncryptionKey {
12 key: [u8; 32],
13}
14
15#[cfg(feature = "encryption")]
16impl EncryptionKey {
17 pub fn new(key: [u8; 32]) -> Self {
18 Self { key }
19 }
20
21 pub fn generate() -> Self {
22 use aes_gcm::aead::rand_core::RngCore;
23 let mut key = [0u8; 32];
24 OsRng.fill_bytes(&mut key);
25 Self { key }
26 }
27
28 pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
29 if bytes.len() != 32 {
30 return Err(PackError::Encryption(
31 "Key must be exactly 32 bytes".to_string()
32 ));
33 }
34 let mut key = [0u8; 32];
35 key.copy_from_slice(bytes);
36 Ok(Self { key })
37 }
38
39 pub fn as_bytes(&self) -> &[u8; 32] {
40 &self.key
41 }
42}
43
44#[cfg(feature = "encryption")]
45pub fn encrypt_snapshot(data: &[u8], key: &EncryptionKey) -> Result<Vec<u8>> {
46 use aes_gcm::aead::rand_core::RngCore;
47
48 let cipher = Aes256Gcm::new_from_slice(&key.key)
49 .map_err(|e| PackError::Encryption(e.to_string()))?;
50
51 let mut nonce_bytes = [0u8; 12];
52 OsRng.fill_bytes(&mut nonce_bytes);
53 let nonce = Nonce::from_slice(&nonce_bytes);
54
55 let ciphertext = cipher
56 .encrypt(nonce, data)
57 .map_err(|e| PackError::Encryption(e.to_string()))?;
58
59 let mut result = Vec::with_capacity(12 + ciphertext.len());
60 result.extend_from_slice(&nonce_bytes);
61 result.extend_from_slice(&ciphertext);
62
63 Ok(result)
64}
65
66#[cfg(feature = "encryption")]
67pub fn decrypt_snapshot(data: &[u8], key: &EncryptionKey) -> Result<Vec<u8>> {
68 if data.len() < 12 {
69 return Err(PackError::Decryption(
70 "Encrypted data too short".to_string()
71 ));
72 }
73
74 let cipher = Aes256Gcm::new_from_slice(&key.key)
75 .map_err(|e| PackError::Decryption(e.to_string()))?;
76
77 let nonce = Nonce::from_slice(&data[0..12]);
78 let ciphertext = &data[12..];
79
80 let plaintext = cipher
81 .decrypt(nonce, ciphertext)
82 .map_err(|e| PackError::Decryption(e.to_string()))?;
83
84 Ok(plaintext)
85}
86
87#[cfg(not(feature = "encryption"))]
88pub struct EncryptionKey;
89
90#[cfg(not(feature = "encryption"))]
91impl EncryptionKey {
92 pub fn new(_key: [u8; 32]) -> Self {
93 Self
94 }
95
96 pub fn generate() -> Self {
97 Self
98 }
99}
100
101#[cfg(not(feature = "encryption"))]
102pub fn encrypt_snapshot(_data: &[u8], _key: &EncryptionKey) -> Result<Vec<u8>> {
103 Err(PackError::Encryption(
104 "Encryption feature not enabled".to_string()
105 ))
106}
107
108#[cfg(not(feature = "encryption"))]
109pub fn decrypt_snapshot(_data: &[u8], _key: &EncryptionKey) -> Result<Vec<u8>> {
110 Err(PackError::Decryption(
111 "Encryption feature not enabled".to_string()
112 ))
113}
114
115#[cfg(all(test, feature = "encryption"))]
116mod tests {
117 use super::*;
118
119 #[test]
120 fn test_encryption_decryption() {
121 let data = b"Hello, World! This is sensitive data.";
122 let key = EncryptionKey::generate();
123
124 let encrypted = encrypt_snapshot(data, &key).unwrap();
125 assert_ne!(data.as_slice(), encrypted.as_slice());
126
127 let decrypted = decrypt_snapshot(&encrypted, &key).unwrap();
128 assert_eq!(data.as_slice(), decrypted.as_slice());
129 }
130
131 #[test]
132 fn test_wrong_key() {
133 let data = b"Hello, World!";
134 let key1 = EncryptionKey::generate();
135 let key2 = EncryptionKey::generate();
136
137 let encrypted = encrypt_snapshot(data, &key1).unwrap();
138 let result = decrypt_snapshot(&encrypted, &key2);
139
140 assert!(result.is_err());
141 }
142}