Skip to main content

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::{Result, bail, ensure};
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!("Invalid encryption algorithm {value}"))),
39        }
40    }
41}
42
43impl From<EncryptionOption> for FileEncryption {
44    fn from(option: EncryptionOption) -> Self {
45        let random_bytes = uuid::Uuid::new_v4().into_bytes().to_vec();
46
47        match option {
48            EncryptionOption::AES128 => FileEncryption::AES128(random_bytes),
49            EncryptionOption::RC4 => FileEncryption::RC4(random_bytes),
50            EncryptionOption::Xor => FileEncryption::Xor(random_bytes),
51        }
52    }
53}
54
55impl Display for EncryptionOption {
56    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
57        match self {
58            EncryptionOption::Xor => write!(f, "Xor"),
59            EncryptionOption::RC4 => write!(f, "RC4"),
60            EncryptionOption::AES128 => write!(f, "AES-128"),
61        }
62    }
63}
64
65/// Some of these algorithms are not secure, and that's fine, since the goal isn't necessarily
66/// data secrecy. The purpose is to store malware on disk without upsetting antivirus or other
67/// endpoint security systems. It would be annoying if our carefully curated data were deleted!
68#[derive(Zeroize, ZeroizeOnDrop, Eq, PartialEq, Hash)]
69#[cfg_attr(test, derive(Clone))]
70pub enum FileEncryption {
71    /// AES-128, to protect from prying eyes.
72    AES128(Vec<u8>),
73
74    /// RC4, to protect from antivirus
75    RC4(Vec<u8>),
76
77    /// Exclusive OR, to protect from antivirus with the best performance
78    Xor(Vec<u8>),
79}
80
81impl FileEncryption {
82    /// Create a key object given the encryption algorithm and key bytes
83    ///
84    /// # Errors
85    ///
86    /// An error occurs if the key isn't 16 bytes.
87    pub fn new(option: EncryptionOption, bytes: Vec<u8>) -> Result<Self> {
88        ensure!(bytes.len() == 16);
89
90        match option {
91            EncryptionOption::AES128 => Ok(FileEncryption::AES128(bytes)),
92            EncryptionOption::RC4 => Ok(FileEncryption::RC4(bytes)),
93            EncryptionOption::Xor => Ok(FileEncryption::Xor(bytes)),
94        }
95    }
96
97    /// Return the name of the algorithm used
98    #[must_use]
99    pub fn name(&self) -> &'static str {
100        match self {
101            FileEncryption::AES128(_) => "aes128",
102            FileEncryption::RC4(_) => "rc4",
103            FileEncryption::Xor(_) => "xor",
104        }
105    }
106
107    /// Return the related [`EncryptionOption`] type
108    #[must_use]
109    pub fn key_type(&self) -> EncryptionOption {
110        match self {
111            FileEncryption::AES128(_) => EncryptionOption::AES128,
112            FileEncryption::RC4(_) => EncryptionOption::RC4,
113            FileEncryption::Xor(_) => EncryptionOption::Xor,
114        }
115    }
116
117    /// Return the bytes for the key
118    #[must_use]
119    pub fn key(&self) -> &[u8] {
120        match self {
121            FileEncryption::AES128(key) | FileEncryption::RC4(key) | FileEncryption::Xor(key) => {
122                key.as_ref()
123            }
124        }
125    }
126
127    /// Decrypt a sample
128    ///
129    /// # Errors
130    ///
131    /// * If the data is corrupted and decryption fails
132    pub fn decrypt(&self, data: &[u8], nonce: Option<Vec<u8>>) -> Result<Vec<u8>> {
133        match self {
134            FileEncryption::AES128(key) => {
135                if let Some(nonce) = nonce {
136                    ensure!(nonce.len() == 12, "AES nonce but be 12 bytes");
137                    let nonce = Nonce::<AesGcm<Aes128, U12>>::from_slice(&nonce);
138                    let key = Key::<Aes128Gcm>::from_slice(key);
139                    let cipher = Aes128Gcm::new(key);
140                    let decrypted = cipher.decrypt(nonce, data)?;
141                    Ok(decrypted)
142                } else {
143                    bail!("Nonce required for AES");
144                }
145            }
146            FileEncryption::RC4(key) => {
147                use rc4::KeyInit;
148
149                let mut key = Rc4::new_from_slice(key)?;
150                let mut output = vec![0u8; data.len()];
151                key.apply_keystream_b2b(data, &mut output);
152                Ok(output)
153            }
154            FileEncryption::Xor(key) => {
155                let mut reader = Cursor::new(data.to_vec());
156                let result = reader.by_ref().xor(key);
157                Ok(result)
158            }
159        }
160    }
161
162    /// Encrypt a sample
163    ///
164    /// # Errors
165    ///
166    /// If AES is missing the nonce, or if the nonce isn't 12 bytes
167    pub fn encrypt(&self, data: &[u8], nonce: Option<Vec<u8>>) -> Result<Vec<u8>> {
168        match self {
169            FileEncryption::AES128(key) => {
170                if let Some(nonce) = nonce {
171                    let nonce = Nonce::<AesGcm<Aes128, U12>>::from_slice(&nonce);
172                    let key = Key::<Aes128Gcm>::from_slice(key);
173                    let cipher = Aes128Gcm::new(key);
174                    let encrypted = cipher.encrypt(nonce, data)?;
175                    Ok(encrypted)
176                } else {
177                    bail!("Nonce required for AES");
178                }
179            }
180            FileEncryption::RC4(key) => {
181                use rc4::KeyInit;
182
183                let mut key = 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
208impl Display for FileEncryption {
209    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
210        write!(f, "{}", self.name())
211    }
212}
213
214#[cfg(test)]
215mod tests {
216    use super::{EncryptionOption, FileEncryption};
217    use malwaredb_types::utils::EntropyCalc;
218
219    use std::time::Instant;
220
221    use rstest::rstest;
222
223    #[rstest]
224    #[case::rc4(EncryptionOption::RC4)]
225    #[case::xor(EncryptionOption::Xor)]
226    #[case::aes128(EncryptionOption::AES128)]
227    #[test]
228    fn enc_dec(#[case] option: EncryptionOption) {
229        const BYTES: &[u8] = include_bytes!("../../types/testdata/exe/pe32_dotnet.exe");
230        let original_entropy = BYTES.entropy();
231
232        let encryptor = FileEncryption::from(option);
233
234        let start = Instant::now();
235        let nonce = encryptor.nonce();
236        let encrypted = encryptor.encrypt(BYTES, nonce.clone()).unwrap();
237        assert_ne!(BYTES, encrypted);
238
239        let encrypted_entropy = encrypted.entropy();
240        assert!(
241            encrypted_entropy > original_entropy,
242            "{option}: Encrypted entropy {encrypted_entropy} should be higher than the original entropy {original_entropy}"
243        );
244        if option != EncryptionOption::Xor {
245            assert!(
246                encrypted_entropy > 7.0,
247                "{option}: Entropy was {encrypted_entropy}, expected >7"
248            );
249        }
250
251        let decrypted = encryptor.decrypt(&encrypted, nonce).unwrap();
252        let duration = start.elapsed();
253        println!(
254            "{option} Time elapsed: {duration:?}, entropy increase: {:+.4}",
255            encrypted_entropy - original_entropy
256        );
257        assert_eq!(BYTES, decrypted);
258    }
259}