webterm_core/cryptography/
cryptographer.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
use crate::compress::{compress, decompress};
use crate::cryptography::iv_counter::IvCounter;
use crate::models::webterm_error::WebtermError;
use crate::types::{Bits256, Bits96};
use aes_gcm::{
    aead::{Aead, KeyInit},
    Aes256Gcm, Nonce,
};
use pbkdf2::pbkdf2_hmac_array;
use sha2::Sha256;

pub struct EncryptedPayload {
    pub iv: Bits96,
    pub ciphertext: Vec<u8>,
    pub compressed: bool,
}

pub struct Cryptographer {
    iv_counter: IvCounter,
    derived_key: Bits256,
}

const COMPRESSION_THRESHOLD: usize = 512;

impl Cryptographer {
    pub fn new(salt: Bits256, secret_key: &str, pbkdf2_iterations: u32) -> Self {
        let derived_key = Self::generate_key(salt, secret_key, pbkdf2_iterations);
        Self {
            iv_counter: IvCounter::new(),
            derived_key,
        }
    }

    fn generate_key(salt: Bits256, secret_key: &str, pbkdf2_iterations: u32) -> Bits256 {
        let derived_key =
            pbkdf2_hmac_array::<Sha256, 32>(secret_key.as_bytes(), &salt.0, pbkdf2_iterations);
        Bits256(derived_key)
    }

    pub fn encrypt(
        &self,
        plaintext: &[u8],
        may_compress: bool,
    ) -> Result<EncryptedPayload, WebtermError> {
        let cipher = Aes256Gcm::new_from_slice(&self.derived_key.0).map_err(|_| {
            WebtermError::RuntimeError("Failed to create encryption key".to_string())
        })?;

        let iv = self.iv_counter.next();
        let nonce = Nonce::try_from(iv.0.as_ref())?;

        let (payload, compressed) = if may_compress && plaintext.len() > COMPRESSION_THRESHOLD {
            (compress(plaintext)?, true)
        } else {
            (plaintext.to_vec(), false)
        };

        if compressed {
            // println!("compressed payload is: {:?} ", payload);
        }

        let ciphertext = cipher
            .encrypt(&nonce, payload.as_ref())
            .map_err(|_| WebtermError::EncryptionError("Encryption failed".to_string()))?;

        Ok(EncryptedPayload {
            ciphertext,
            iv,
            compressed,
        })
    }

    pub fn decrypt(
        &self,
        ciphertext: &[u8],
        iv: &Bits96,
        compressed: bool,
    ) -> Result<Vec<u8>, WebtermError> {
        let cipher = Aes256Gcm::new_from_slice(&self.derived_key.0).map_err(|_| {
            WebtermError::RuntimeError("Failed to create decryption key".to_string())
        })?;

        let nonce = Nonce::try_from(iv.0.as_ref())?;

        let decrypted = cipher
            .decrypt(&nonce, ciphertext)
            .map_err(|_| WebtermError::DecryptionError("Decryption failed".to_string()))?;

        if compressed {
            decompress(&decrypted)
        } else {
            Ok(decrypted)
        }
    }
}