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