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 => {
52 let random_bytes = random_bytes.as_slice()[0..16].to_vec();
53 FileEncryption::RC4(random_bytes)
54 }
55 EncryptionOption::Xor => FileEncryption::Xor(random_bytes),
56 }
57 }
58}
59
60impl Display for EncryptionOption {
61 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
62 match self {
63 EncryptionOption::Xor => write!(f, "Xor"),
64 EncryptionOption::RC4 => write!(f, "RC4"),
65 EncryptionOption::AES128 => write!(f, "AES-128"),
66 }
67 }
68}
69
70#[derive(Zeroize, ZeroizeOnDrop, Eq, PartialEq, Hash)]
74pub enum FileEncryption {
75 AES128(Vec<u8>),
77
78 RC4(Vec<u8>),
80
81 Xor(Vec<u8>),
83}
84
85impl FileEncryption {
86 pub fn new(option: EncryptionOption, bytes: Vec<u8>) -> Result<Self> {
92 match option {
93 EncryptionOption::AES128 => Ok(FileEncryption::AES128(bytes)),
94 EncryptionOption::RC4 => {
95 ensure!(bytes.len() == 16);
96 Ok(FileEncryption::RC4(bytes))
97 }
98 EncryptionOption::Xor => Ok(FileEncryption::Xor(bytes)),
99 }
100 }
101
102 #[must_use]
104 pub fn name(&self) -> &'static str {
105 match self {
106 FileEncryption::AES128(_) => "aes128",
107 FileEncryption::RC4(_) => "rc4",
108 FileEncryption::Xor(_) => "xor",
109 }
110 }
111
112 #[must_use]
114 pub fn key_type(&self) -> EncryptionOption {
115 match self {
116 FileEncryption::AES128(_) => EncryptionOption::AES128,
117 FileEncryption::RC4(_) => EncryptionOption::RC4,
118 FileEncryption::Xor(_) => EncryptionOption::Xor,
119 }
120 }
121
122 #[must_use]
124 pub fn key(&self) -> &[u8] {
125 match self {
126 FileEncryption::AES128(key) | FileEncryption::RC4(key) | FileEncryption::Xor(key) => {
127 key.as_ref()
128 }
129 }
130 }
131
132 pub fn decrypt(&self, data: &[u8], nonce: Option<Vec<u8>>) -> Result<Vec<u8>> {
138 match self {
139 FileEncryption::AES128(key) => {
140 if let Some(nonce) = nonce {
141 let nonce = Nonce::<AesGcm<Aes128, U12>>::from_slice(&nonce);
142 let key = Key::<Aes128Gcm>::from_slice(key);
143 let cipher = Aes128Gcm::new(key);
144 let decrypted = cipher.decrypt(nonce, data)?;
145 Ok(decrypted)
146 } else {
147 bail!("Nonce required for AES");
148 }
149 }
150 FileEncryption::RC4(key) => {
151 let mut key: Rc4<rc4::consts::U16> = Rc4::new_from_slice(key)?;
152 let mut output = vec![0u8; data.len()];
153 key.apply_keystream_b2b(data, &mut output)?;
154 Ok(output)
155 }
156 FileEncryption::Xor(key) => {
157 let mut reader = Cursor::new(data.to_vec());
158 let result = reader.by_ref().xor(key);
159 Ok(result)
160 }
161 }
162 }
163
164 pub fn encrypt(&self, data: &[u8], nonce: Option<Vec<u8>>) -> Result<Vec<u8>> {
170 match self {
171 FileEncryption::AES128(key) => {
172 if let Some(nonce) = nonce {
173 let nonce = Nonce::<AesGcm<Aes128, U12>>::from_slice(&nonce);
174 let key = Key::<Aes128Gcm>::from_slice(key);
175 let cipher = Aes128Gcm::new(key);
176 let encrypted = cipher.encrypt(nonce, data)?;
177 Ok(encrypted)
178 } else {
179 bail!("Nonce required for AES");
180 }
181 }
182 FileEncryption::RC4(key) => {
183 let mut key: Rc4<rc4::consts::U16> = 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
208#[cfg(test)]
209mod tests {
210 use super::{EncryptionOption, FileEncryption};
211 use malwaredb_types::utils::EntropyCalc;
212
213 use std::time::Instant;
214
215 use rstest::rstest;
216
217 #[rstest]
218 #[case::rc4(EncryptionOption::RC4)]
219 #[case::xor(EncryptionOption::Xor)]
220 #[case::aes128(EncryptionOption::AES128)]
221 #[test]
222 fn enc_dec(#[case] option: EncryptionOption) {
223 const BYTES: &[u8] = include_bytes!("../../types/testdata/exe/pe32_dotnet.exe");
224 let original_entropy = BYTES.entropy();
225
226 let encryptor = FileEncryption::from(option);
227
228 let start = Instant::now();
229 let nonce = encryptor.nonce();
230 let encrypted = encryptor.encrypt(BYTES, nonce.clone()).unwrap();
231 assert_ne!(BYTES, encrypted);
232
233 let encrypted_entropy = encrypted.entropy();
234 assert!(encrypted_entropy > original_entropy, "{option}: Encrypted entropy {encrypted_entropy} should be higher than the original entropy {original_entropy}");
235 if option != EncryptionOption::Xor {
236 assert!(
237 encrypted_entropy > 7.0,
238 "{option}: Entropy was {encrypted_entropy}, expected >7"
239 );
240 }
241
242 let decrypted = encryptor.decrypt(&encrypted, nonce).unwrap();
243 let duration = start.elapsed();
244 println!("{option} Time elapsed: {duration:?}");
245 assert_eq!(BYTES, decrypted);
246 }
247}