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