malwaredb_server/
crypto.rs

1// SPDX-License-Identifier: Apache-2.0
2
3use std::fmt::{Display, Formatter};
4use std::io::{Cursor, Write};
5
6use aes_gcm::aead::{Aead, Nonce, OsRng};
7use aes_gcm::aes::Aes128;
8use aes_gcm::{AeadCore, Aes128Gcm, AesGcm, Key, KeyInit};
9use anyhow::{bail, ensure, Result};
10use clap::ValueEnum;
11use deadpool_postgres::tokio_postgres::types::{FromSql, ToSql};
12use md5::digest::consts::U12;
13use rc4::{Rc4, StreamCipher};
14use xor_utils::Xor;
15use zeroize::{Zeroize, ZeroizeOnDrop};
16
17/// Available options for specifying which algorithm to use
18#[derive(Debug, Copy, Clone, Eq, PartialEq, ValueEnum, Hash, ToSql, FromSql)]
19#[postgres(name = "encryptionkey_algorithm", rename_all = "lowercase")]
20pub enum EncryptionOption {
21    /// AES-128 encryption, the best
22    AES128,
23
24    /// RC4 encryption, pretty weak but effective enough
25    RC4,
26
27    /// Exclusive Or (XOR), also weak but effective and fastest.
28    Xor,
29}
30
31impl TryFrom<&str> for EncryptionOption {
32    type Error = anyhow::Error;
33    fn try_from(value: &str) -> std::result::Result<Self, Self::Error> {
34        match value {
35            "xor" => Ok(EncryptionOption::Xor),
36            "rc4" => Ok(EncryptionOption::RC4),
37            "aes128" => Ok(EncryptionOption::AES128),
38            _ => Err(anyhow::Error::msg(format!(
39                "Invalid encryption algorithm {value}"
40            ))),
41        }
42    }
43}
44
45impl From<EncryptionOption> for FileEncryption {
46    fn from(option: EncryptionOption) -> Self {
47        let random_bytes = uuid::Uuid::new_v4().into_bytes().to_vec();
48
49        match option {
50            EncryptionOption::AES128 => FileEncryption::AES128(random_bytes),
51            EncryptionOption::RC4 => {
52                let random_bytes = random_bytes.as_slice()[0..16].to_vec();
53                FileEncryption::RC4(random_bytes)
54            }
55            EncryptionOption::Xor => FileEncryption::Xor(random_bytes),
56        }
57    }
58}
59
60impl Display for EncryptionOption {
61    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
62        match self {
63            EncryptionOption::Xor => write!(f, "Xor"),
64            EncryptionOption::RC4 => write!(f, "RC4"),
65            EncryptionOption::AES128 => write!(f, "AES-128"),
66        }
67    }
68}
69
70/// Some of these algorithms are not secure, and that's fine, since the goal isn't necessarily
71/// data secrecy. The purpose is to store malware on disk without upsetting anti-virus systems.
72/// It would be annoying if security software deleted our carefully curated data!
73#[derive(Zeroize, ZeroizeOnDrop, Eq, PartialEq, Hash)]
74pub enum FileEncryption {
75    /// AES-128, to protect from prying eyes.
76    AES128(Vec<u8>),
77
78    /// RC4 with a 16-byte key, to protect from anti-virus
79    RC4(Vec<u8>),
80
81    /// Exclusive OR, to protect from anti-virus
82    Xor(Vec<u8>),
83}
84
85impl FileEncryption {
86    /// Create a key object given the encryption algorithm and key bytes
87    ///
88    /// # Errors
89    ///
90    /// The only error is if the RC4 key isn't exactly 16 bytes long.
91    pub fn new(option: EncryptionOption, bytes: Vec<u8>) -> Result<Self> {
92        match option {
93            EncryptionOption::AES128 => Ok(FileEncryption::AES128(bytes)),
94            EncryptionOption::RC4 => {
95                ensure!(bytes.len() == 16);
96                Ok(FileEncryption::RC4(bytes))
97            }
98            EncryptionOption::Xor => Ok(FileEncryption::Xor(bytes)),
99        }
100    }
101
102    /// Return the name of the algorithm used
103    #[must_use]
104    pub fn name(&self) -> &'static str {
105        match self {
106            FileEncryption::AES128(_) => "aes128",
107            FileEncryption::RC4(_) => "rc4",
108            FileEncryption::Xor(_) => "xor",
109        }
110    }
111
112    /// Return the related `EncryptionOption` type
113    #[must_use]
114    pub fn key_type(&self) -> EncryptionOption {
115        match self {
116            FileEncryption::AES128(_) => EncryptionOption::AES128,
117            FileEncryption::RC4(_) => EncryptionOption::RC4,
118            FileEncryption::Xor(_) => EncryptionOption::Xor,
119        }
120    }
121
122    /// Return the bytes for the key
123    #[must_use]
124    pub fn key(&self) -> &[u8] {
125        match self {
126            FileEncryption::AES128(key) | FileEncryption::RC4(key) | FileEncryption::Xor(key) => {
127                key.as_ref()
128            }
129        }
130    }
131
132    /// Decrypt some data returning the clear bytes
133    ///
134    /// # Errors
135    ///
136    /// * If the data is corrupted and decryption fails
137    pub fn decrypt(&self, data: &[u8], nonce: Option<Vec<u8>>) -> Result<Vec<u8>> {
138        match self {
139            FileEncryption::AES128(key) => {
140                if let Some(nonce) = nonce {
141                    let nonce = Nonce::<AesGcm<Aes128, U12>>::from_slice(&nonce);
142                    let key = Key::<Aes128Gcm>::from_slice(key);
143                    let cipher = Aes128Gcm::new(key);
144                    let decrypted = cipher.decrypt(nonce, data)?;
145                    Ok(decrypted)
146                } else {
147                    bail!("Nonce required for AES");
148                }
149            }
150            FileEncryption::RC4(key) => {
151                let mut key: Rc4<rc4::consts::U16> = Rc4::new_from_slice(key)?;
152                let mut output = vec![0u8; data.len()];
153                key.apply_keystream_b2b(data, &mut output)?;
154                Ok(output)
155            }
156            FileEncryption::Xor(key) => {
157                let mut reader = Cursor::new(data.to_vec());
158                let result = reader.by_ref().xor(key);
159                Ok(result)
160            }
161        }
162    }
163
164    /// Encrypt some data returning the ciphertext bytes
165    ///
166    /// # Errors
167    ///
168    /// If AES somehow doesn't have a nonce.
169    pub fn encrypt(&self, data: &[u8], nonce: Option<Vec<u8>>) -> Result<Vec<u8>> {
170        match self {
171            FileEncryption::AES128(key) => {
172                if let Some(nonce) = nonce {
173                    let nonce = Nonce::<AesGcm<Aes128, U12>>::from_slice(&nonce);
174                    let key = Key::<Aes128Gcm>::from_slice(key);
175                    let cipher = Aes128Gcm::new(key);
176                    let encrypted = cipher.encrypt(nonce, data)?;
177                    Ok(encrypted)
178                } else {
179                    bail!("Nonce required for AES");
180                }
181            }
182            FileEncryption::RC4(key) => {
183                let mut key: Rc4<rc4::consts::U16> = Rc4::new_from_slice(key)?;
184                let mut output = vec![0u8; data.len()];
185                key.apply_keystream_b2b(data, &mut output)?;
186                Ok(output)
187            }
188            FileEncryption::Xor(key) => {
189                let mut reader = Cursor::new(data.to_vec());
190                let result = reader.by_ref().xor(key);
191                Ok(result)
192            }
193        }
194    }
195
196    /// Generate nonce bytes, if used by the algorithm
197    pub fn nonce(&self) -> Option<Vec<u8>> {
198        match self {
199            FileEncryption::AES128(_) => {
200                let nonce = Aes128Gcm::generate_nonce(&mut OsRng);
201                Some(nonce.to_vec())
202            }
203            _ => None,
204        }
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use super::{EncryptionOption, FileEncryption};
211    use malwaredb_types::utils::EntropyCalc;
212
213    use std::time::Instant;
214
215    use rstest::rstest;
216
217    #[rstest]
218    #[case::rc4(EncryptionOption::RC4)]
219    #[case::xor(EncryptionOption::Xor)]
220    #[case::aes128(EncryptionOption::AES128)]
221    #[test]
222    fn enc_dec(#[case] option: EncryptionOption) {
223        const BYTES: &[u8] = include_bytes!("../../types/testdata/exe/pe32_dotnet.exe");
224        let original_entropy = BYTES.entropy();
225
226        let encryptor = FileEncryption::from(option);
227
228        let start = Instant::now();
229        let nonce = encryptor.nonce();
230        let encrypted = encryptor.encrypt(BYTES, nonce.clone()).unwrap();
231        assert_ne!(BYTES, encrypted);
232
233        let encrypted_entropy = encrypted.entropy();
234        assert!(encrypted_entropy > original_entropy, "{option}: Encrypted entropy {encrypted_entropy} should be higher than the original entropy {original_entropy}");
235        if option != EncryptionOption::Xor {
236            assert!(
237                encrypted_entropy > 7.0,
238                "{option}: Entropy was {encrypted_entropy}, expected >7"
239            );
240        }
241
242        let decrypted = encryptor.decrypt(&encrypted, nonce).unwrap();
243        let duration = start.elapsed();
244        println!("{option} Time elapsed: {duration:?}");
245        assert_eq!(BYTES, decrypted);
246    }
247}