malwaredb_server/
crypto.rs1use 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#[derive(Debug, Copy, Clone, Eq, PartialEq, ValueEnum, Hash, ToSql, FromSql)]
19#[postgres(name = "encryptionkey_algorithm", rename_all = "lowercase")]
20pub enum EncryptionOption {
21 AES128,
23
24 RC4,
26
27 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#[derive(Zeroize, ZeroizeOnDrop, Eq, PartialEq, Hash)]
69#[cfg_attr(test, derive(Clone))]
70pub enum FileEncryption {
71 AES128(Vec<u8>),
73
74 RC4(Vec<u8>),
76
77 Xor(Vec<u8>),
79}
80
81impl FileEncryption {
82 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 #[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 #[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 #[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 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 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 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}