webterm_core/cryptography/
cryptographer.rs

1use crate::compress::{compress, decompress};
2use crate::cryptography::iv_counter::IvCounter;
3use crate::models::webterm_error::WebtermError;
4use crate::types::{Bits256, Bits96};
5use aes_gcm::{
6    aead::{Aead, KeyInit},
7    Aes256Gcm, Nonce,
8};
9use pbkdf2::pbkdf2_hmac_array;
10use sha2::Sha256;
11
12pub struct EncryptedPayload {
13    pub iv: Bits96,
14    pub ciphertext: Vec<u8>,
15    pub compressed: bool,
16}
17
18pub struct Cryptographer {
19    iv_counter: IvCounter,
20    derived_key: Bits256,
21}
22
23const COMPRESSION_THRESHOLD: usize = 512;
24
25impl Cryptographer {
26    pub fn new(salt: Bits256, secret_key: &str, pbkdf2_iterations: u32) -> Self {
27        let derived_key = Self::generate_key(salt, secret_key, pbkdf2_iterations);
28        Self {
29            iv_counter: IvCounter::new(),
30            derived_key,
31        }
32    }
33
34    fn generate_key(salt: Bits256, secret_key: &str, pbkdf2_iterations: u32) -> Bits256 {
35        let derived_key =
36            pbkdf2_hmac_array::<Sha256, 32>(secret_key.as_bytes(), &salt.0, pbkdf2_iterations);
37        Bits256(derived_key)
38    }
39
40    pub fn encrypt(
41        &self,
42        plaintext: &[u8],
43        may_compress: bool,
44    ) -> Result<EncryptedPayload, WebtermError> {
45        let cipher = Aes256Gcm::new_from_slice(&self.derived_key.0).map_err(|_| {
46            WebtermError::RuntimeError("Failed to create encryption key".to_string())
47        })?;
48
49        let iv = self.iv_counter.next();
50        let nonce = Nonce::try_from(iv.0.as_ref())?;
51
52        let (payload, compressed) = if may_compress && plaintext.len() > COMPRESSION_THRESHOLD {
53            (compress(plaintext)?, true)
54        } else {
55            (plaintext.to_vec(), false)
56        };
57
58        if compressed {
59            // println!("compressed payload is: {:?} ", payload);
60        }
61
62        let ciphertext = cipher
63            .encrypt(&nonce, payload.as_ref())
64            .map_err(|_| WebtermError::EncryptionError("Encryption failed".to_string()))?;
65
66        Ok(EncryptedPayload {
67            ciphertext,
68            iv,
69            compressed,
70        })
71    }
72
73    pub fn decrypt(
74        &self,
75        ciphertext: &[u8],
76        iv: &Bits96,
77        compressed: bool,
78    ) -> Result<Vec<u8>, WebtermError> {
79        let cipher = Aes256Gcm::new_from_slice(&self.derived_key.0).map_err(|_| {
80            WebtermError::RuntimeError("Failed to create decryption key".to_string())
81        })?;
82
83        let nonce = Nonce::try_from(iv.0.as_ref())?;
84
85        let decrypted = cipher
86            .decrypt(&nonce, ciphertext)
87            .map_err(|_| WebtermError::DecryptionError("Decryption failed".to_string()))?;
88
89        if compressed {
90            decompress(&decrypted)
91        } else {
92            Ok(decrypted)
93        }
94    }
95}