1use aes::{Aes128, Aes192, Aes256};
2use aes_gcm::{
3 Aes128Gcm, Aes256Gcm, AesGcm, Nonce,
4 aead::{Aead, KeyInit, consts::U12},
5};
6use base64::{Engine, engine::general_purpose::STANDARD};
7use cbc::Encryptor as CbcEncryptor;
8use cipher::{BlockEncryptMut, KeyIvInit, block_padding::Pkcs7};
9use ecb::Encryptor as EcbEncryptor;
10
11use crate::error::{Error, Result};
12
13#[derive(Clone, Copy, Debug, Eq, PartialEq)]
18pub enum EncryptionAlgorithm {
19 AES128,
21 AES192,
23 AES256,
25}
26
27impl EncryptionAlgorithm {
28 pub const fn key_len(self) -> usize {
30 match self {
31 Self::AES128 => 16,
32 Self::AES192 => 24,
33 Self::AES256 => 32,
34 }
35 }
36
37 pub const fn as_bark_str(self) -> &'static str {
39 match self {
40 Self::AES128 => "AES128",
41 Self::AES192 => "AES192",
42 Self::AES256 => "AES256",
43 }
44 }
45}
46
47#[derive(Clone, Copy, Debug, Eq, PartialEq)]
53pub enum EncryptionMode {
54 CBC,
56 ECB,
58 GCM,
60}
61
62impl EncryptionMode {
63 pub const fn iv_len(self) -> Option<usize> {
67 match self {
68 Self::CBC => Some(16),
69 Self::ECB => None,
70 Self::GCM => Some(12),
71 }
72 }
73
74 pub const fn as_bark_str(self) -> &'static str {
76 match self {
77 Self::CBC => "CBC",
78 Self::ECB => "ECB",
79 Self::GCM => "GCM",
80 }
81 }
82}
83
84#[derive(Clone, Debug, Eq, PartialEq)]
95pub struct Encryption {
96 algorithm: EncryptionAlgorithm,
97 mode: EncryptionMode,
98 key: String,
99 iv: String,
100}
101
102impl Encryption {
103 pub fn new<K>(algorithm: EncryptionAlgorithm, mode: EncryptionMode, key: K) -> Result<Self>
118 where
119 K: Into<String>,
120 {
121 let iv = match mode.iv_len() {
122 Some(len) => random_ascii_iv(len)?,
123 None => String::new(),
124 };
125
126 Self::with_iv(algorithm, mode, key, iv)
127 }
128
129 pub fn with_iv<K, I>(
146 algorithm: EncryptionAlgorithm,
147 mode: EncryptionMode,
148 key: K,
149 iv: I,
150 ) -> Result<Self>
151 where
152 K: Into<String>,
153 I: Into<String>,
154 {
155 let key = key.into();
156 let iv = iv.into();
157 let expected_key_len = algorithm.key_len();
158 let actual_key_len = key.len();
159 if actual_key_len != expected_key_len {
160 return Err(Error::InvalidKeyLength {
161 algorithm: algorithm.as_bark_str(),
162 expected: expected_key_len,
163 actual: actual_key_len,
164 });
165 }
166
167 let expected_iv_len = mode.iv_len().unwrap_or(0);
168 let actual_iv_len = iv.len();
169 if actual_iv_len != expected_iv_len {
170 return Err(Error::InvalidIvLength {
171 mode: mode.as_bark_str(),
172 expected: expected_iv_len,
173 actual: actual_iv_len,
174 });
175 }
176
177 Ok(Self {
178 algorithm,
179 mode,
180 key,
181 iv,
182 })
183 }
184
185 pub const fn algorithm(&self) -> EncryptionAlgorithm {
187 self.algorithm
188 }
189
190 pub const fn mode(&self) -> EncryptionMode {
192 self.mode
193 }
194
195 pub fn key(&self) -> &str {
200 &self.key
201 }
202
203 pub fn iv(&self) -> Option<&str> {
207 self.mode.iv_len().map(|_| self.iv.as_str())
208 }
209
210 pub(crate) fn apns_iv(&self) -> Option<&str> {
211 self.iv()
212 }
213
214 pub(crate) fn encrypt_bark_json(&self, plaintext: &[u8]) -> Result<String> {
215 let encrypted = match (self.algorithm, self.mode) {
216 (EncryptionAlgorithm::AES128, EncryptionMode::CBC) => {
217 CbcEncryptor::<Aes128>::new_from_slices(self.key.as_bytes(), self.iv.as_bytes())
218 .map_err(|_| Error::Encryption)?
219 .encrypt_padded_vec_mut::<Pkcs7>(plaintext)
220 }
221 (EncryptionAlgorithm::AES192, EncryptionMode::CBC) => {
222 CbcEncryptor::<Aes192>::new_from_slices(self.key.as_bytes(), self.iv.as_bytes())
223 .map_err(|_| Error::Encryption)?
224 .encrypt_padded_vec_mut::<Pkcs7>(plaintext)
225 }
226 (EncryptionAlgorithm::AES256, EncryptionMode::CBC) => {
227 CbcEncryptor::<Aes256>::new_from_slices(self.key.as_bytes(), self.iv.as_bytes())
228 .map_err(|_| Error::Encryption)?
229 .encrypt_padded_vec_mut::<Pkcs7>(plaintext)
230 }
231 (EncryptionAlgorithm::AES128, EncryptionMode::ECB) => {
232 EcbEncryptor::<Aes128>::new_from_slice(self.key.as_bytes())
233 .map_err(|_| Error::Encryption)?
234 .encrypt_padded_vec_mut::<Pkcs7>(plaintext)
235 }
236 (EncryptionAlgorithm::AES192, EncryptionMode::ECB) => {
237 EcbEncryptor::<Aes192>::new_from_slice(self.key.as_bytes())
238 .map_err(|_| Error::Encryption)?
239 .encrypt_padded_vec_mut::<Pkcs7>(plaintext)
240 }
241 (EncryptionAlgorithm::AES256, EncryptionMode::ECB) => {
242 EcbEncryptor::<Aes256>::new_from_slice(self.key.as_bytes())
243 .map_err(|_| Error::Encryption)?
244 .encrypt_padded_vec_mut::<Pkcs7>(plaintext)
245 }
246 (EncryptionAlgorithm::AES128, EncryptionMode::GCM) => {
247 let cipher = Aes128Gcm::new_from_slice(self.key.as_bytes())
248 .map_err(|_| Error::Encryption)?;
249 cipher
250 .encrypt(Nonce::from_slice(self.iv.as_bytes()), plaintext)
251 .map_err(|_| Error::Encryption)?
252 }
253 (EncryptionAlgorithm::AES192, EncryptionMode::GCM) => {
254 let cipher = AesGcm::<Aes192, U12>::new_from_slice(self.key.as_bytes())
255 .map_err(|_| Error::Encryption)?;
256 cipher
257 .encrypt(Nonce::from_slice(self.iv.as_bytes()), plaintext)
258 .map_err(|_| Error::Encryption)?
259 }
260 (EncryptionAlgorithm::AES256, EncryptionMode::GCM) => {
261 let cipher = Aes256Gcm::new_from_slice(self.key.as_bytes())
262 .map_err(|_| Error::Encryption)?;
263 cipher
264 .encrypt(Nonce::from_slice(self.iv.as_bytes()), plaintext)
265 .map_err(|_| Error::Encryption)?
266 }
267 };
268
269 Ok(STANDARD.encode(encrypted))
270 }
271}
272
273fn random_ascii_iv(len: usize) -> Result<String> {
274 const ALPHABET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";
275
276 let mut bytes = vec![0; len];
277 getrandom::getrandom(&mut bytes)?;
278
279 Ok(bytes
280 .into_iter()
281 .map(|byte| ALPHABET[usize::from(byte) % ALPHABET.len()] as char)
282 .collect())
283}
284
285#[cfg(test)]
286mod tests {
287 use super::*;
288
289 #[test]
290 fn validates_key_length() {
291 let err =
292 Encryption::new(EncryptionAlgorithm::AES128, EncryptionMode::CBC, "short").unwrap_err();
293
294 assert!(matches!(
295 err,
296 Error::InvalidKeyLength {
297 algorithm: "AES128",
298 expected: 16,
299 actual: 5
300 }
301 ));
302 }
303
304 #[test]
305 fn validates_cbc_iv_length() {
306 let err = Encryption::with_iv(
307 EncryptionAlgorithm::AES128,
308 EncryptionMode::CBC,
309 "1234567890123456",
310 "short",
311 )
312 .unwrap_err();
313
314 assert!(matches!(
315 err,
316 Error::InvalidIvLength {
317 mode: "CBC",
318 expected: 16,
319 actual: 5
320 }
321 ));
322 }
323
324 #[test]
325 fn new_generates_mode_specific_iv() {
326 let cbc = Encryption::new(
327 EncryptionAlgorithm::AES128,
328 EncryptionMode::CBC,
329 "1234567890123456",
330 )
331 .unwrap();
332 let gcm = Encryption::new(
333 EncryptionAlgorithm::AES128,
334 EncryptionMode::GCM,
335 "1234567890123456",
336 )
337 .unwrap();
338 let ecb = Encryption::new(
339 EncryptionAlgorithm::AES128,
340 EncryptionMode::ECB,
341 "1234567890123456",
342 )
343 .unwrap();
344
345 assert_eq!(cbc.iv().unwrap().len(), 16);
346 assert_eq!(gcm.iv().unwrap().len(), 12);
347 assert_eq!(ecb.iv(), None);
348 }
349
350 #[test]
351 fn encrypts_like_bark_docs_cbc_example() {
352 let encryption = Encryption::with_iv(
353 EncryptionAlgorithm::AES128,
354 EncryptionMode::CBC,
355 "1234567890123456",
356 "1111111111111111",
357 )
358 .unwrap();
359
360 let ciphertext = encryption
361 .encrypt_bark_json(br#"{"body": "test", "sound": "birdsong"}"#)
362 .unwrap();
363
364 assert_eq!(
365 ciphertext,
366 "d3QhjQjP5majvNt5CjsvFWwqqj2gKl96RFj5OO+u6ynTt7lkyigDYNA3abnnCLpr"
367 );
368 }
369}