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::{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#[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!(
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#[derive(Zeroize, ZeroizeOnDrop, Eq, PartialEq, Hash)]
71pub enum FileEncryption {
72 AES128(Vec<u8>),
74
75 RC4(Vec<u8>),
77
78 Xor(Vec<u8>),
80}
81
82impl FileEncryption {
83 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 #[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 #[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 #[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 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 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 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}