1use 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#[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 AES128,
21
22 RC4,
24
25 #[default]
27 Xor,
28
29 #[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#[derive(Zeroize, ZeroizeOnDrop, Eq, PartialEq, Hash)]
79#[cfg_attr(test, derive(Clone))]
80pub enum FileEncryption {
81 AES128(Vec<u8>),
83
84 RC4(Vec<u8>),
86
87 Xor(Vec<u8>),
89}
90
91impl FileEncryption {
92 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 #[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 #[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 #[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 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 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 #[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}