tx2_pack/
encryption.rs

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}