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)]
91pub struct Device {
92 token: String,
93 encryption_algorithm: Option<EncryptionAlgorithm>,
94 encryption_mode: Option<EncryptionMode>,
95 encryption_key: Option<String>,
96}
97
98impl Device {
99 pub fn new<T>(token: T) -> Self
101 where
102 T: Into<String>,
103 {
104 Self {
105 token: normalize_device_token(token.into()),
106 encryption_algorithm: None,
107 encryption_mode: None,
108 encryption_key: None,
109 }
110 }
111
112 pub fn encrypt<K>(
136 mut self,
137 algorithm: EncryptionAlgorithm,
138 mode: EncryptionMode,
139 key: K,
140 ) -> Result<Self>
141 where
142 K: Into<String>,
143 {
144 let key = key.into();
145 let expected = algorithm.key_len();
146 let actual = key.len();
147 if actual != expected {
148 return Err(Error::InvalidKeyLength {
149 algorithm: algorithm.as_bark_str(),
150 expected,
151 actual,
152 });
153 }
154
155 self.encryption_algorithm = Some(algorithm);
156 self.encryption_mode = Some(mode);
157 self.encryption_key = Some(key);
158 Ok(self)
159 }
160
161 pub fn token(&self) -> &str {
163 &self.token
164 }
165
166 pub(crate) fn has_encryption(&self) -> bool {
168 self.encryption().is_some()
169 }
170
171 pub(crate) fn encrypt_bark_json(&self, plaintext: &[u8]) -> Result<EncryptedPayload> {
172 let (algorithm, mode, key) =
173 self.encryption()
174 .ok_or_else(|| Error::MissingDeviceEncryption {
175 device: self.token.clone(),
176 })?;
177 let iv = match mode.iv_len() {
178 Some(len) => random_ascii_iv(len)?,
179 None => String::new(),
180 };
181
182 let encrypted = match (algorithm, mode) {
183 (EncryptionAlgorithm::AES128, EncryptionMode::CBC) => {
184 CbcEncryptor::<Aes128>::new_from_slices(key.as_bytes(), iv.as_bytes())
185 .map_err(|_| Error::Encryption)?
186 .encrypt_padded_vec_mut::<Pkcs7>(plaintext)
187 }
188 (EncryptionAlgorithm::AES192, EncryptionMode::CBC) => {
189 CbcEncryptor::<Aes192>::new_from_slices(key.as_bytes(), iv.as_bytes())
190 .map_err(|_| Error::Encryption)?
191 .encrypt_padded_vec_mut::<Pkcs7>(plaintext)
192 }
193 (EncryptionAlgorithm::AES256, EncryptionMode::CBC) => {
194 CbcEncryptor::<Aes256>::new_from_slices(key.as_bytes(), iv.as_bytes())
195 .map_err(|_| Error::Encryption)?
196 .encrypt_padded_vec_mut::<Pkcs7>(plaintext)
197 }
198 (EncryptionAlgorithm::AES128, EncryptionMode::ECB) => {
199 EcbEncryptor::<Aes128>::new_from_slice(key.as_bytes())
200 .map_err(|_| Error::Encryption)?
201 .encrypt_padded_vec_mut::<Pkcs7>(plaintext)
202 }
203 (EncryptionAlgorithm::AES192, EncryptionMode::ECB) => {
204 EcbEncryptor::<Aes192>::new_from_slice(key.as_bytes())
205 .map_err(|_| Error::Encryption)?
206 .encrypt_padded_vec_mut::<Pkcs7>(plaintext)
207 }
208 (EncryptionAlgorithm::AES256, EncryptionMode::ECB) => {
209 EcbEncryptor::<Aes256>::new_from_slice(key.as_bytes())
210 .map_err(|_| Error::Encryption)?
211 .encrypt_padded_vec_mut::<Pkcs7>(plaintext)
212 }
213 (EncryptionAlgorithm::AES128, EncryptionMode::GCM) => {
214 let cipher =
215 Aes128Gcm::new_from_slice(key.as_bytes()).map_err(|_| Error::Encryption)?;
216 cipher
217 .encrypt(Nonce::from_slice(iv.as_bytes()), plaintext)
218 .map_err(|_| Error::Encryption)?
219 }
220 (EncryptionAlgorithm::AES192, EncryptionMode::GCM) => {
221 let cipher = AesGcm::<Aes192, U12>::new_from_slice(key.as_bytes())
222 .map_err(|_| Error::Encryption)?;
223 cipher
224 .encrypt(Nonce::from_slice(iv.as_bytes()), plaintext)
225 .map_err(|_| Error::Encryption)?
226 }
227 (EncryptionAlgorithm::AES256, EncryptionMode::GCM) => {
228 let cipher =
229 Aes256Gcm::new_from_slice(key.as_bytes()).map_err(|_| Error::Encryption)?;
230 cipher
231 .encrypt(Nonce::from_slice(iv.as_bytes()), plaintext)
232 .map_err(|_| Error::Encryption)?
233 }
234 };
235
236 Ok(EncryptedPayload {
237 ciphertext: STANDARD.encode(encrypted),
238 iv: mode.iv_len().map(|_| iv),
239 })
240 }
241
242 fn encryption(&self) -> Option<(EncryptionAlgorithm, EncryptionMode, &str)> {
243 Some((
244 self.encryption_algorithm?,
245 self.encryption_mode?,
246 self.encryption_key.as_deref()?,
247 ))
248 }
249}
250
251pub(crate) struct EncryptedPayload {
253 pub(crate) ciphertext: String,
255 pub(crate) iv: Option<String>,
257}
258
259fn random_ascii_iv(len: usize) -> Result<String> {
260 const ALPHABET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";
261
262 let mut bytes = vec![0; len];
263 getrandom::getrandom(&mut bytes)?;
264
265 Ok(bytes
266 .into_iter()
267 .map(|byte| ALPHABET[usize::from(byte) % ALPHABET.len()] as char)
268 .collect())
269}
270
271fn normalize_device_token(device: String) -> String {
272 device
273 .trim()
274 .trim_start_matches('<')
275 .trim_end_matches('>')
276 .chars()
277 .filter(|ch| !ch.is_ascii_whitespace())
278 .collect()
279}
280
281#[cfg(test)]
282mod tests {
283 use super::*;
284
285 #[test]
286 fn normalizes_device_tokens() {
287 let device = Device::new("<aa bb>");
288
289 assert_eq!(device.token(), "aabb");
290 }
291
292 #[test]
293 fn keeps_device_encryption() {
294 let device = Device::new("aabb")
295 .encrypt(
296 EncryptionAlgorithm::AES128,
297 EncryptionMode::CBC,
298 "1234567890123456",
299 )
300 .unwrap();
301
302 assert_eq!(
303 device.encryption_algorithm,
304 Some(EncryptionAlgorithm::AES128)
305 );
306 assert_eq!(device.encryption_mode, Some(EncryptionMode::CBC));
307 assert_eq!(device.encryption_key.as_deref(), Some("1234567890123456"));
308 assert!(device.has_encryption());
309 }
310
311 #[test]
312 fn validates_encryption_key_length() {
313 let err = Device::new("aabb")
314 .encrypt(EncryptionAlgorithm::AES128, EncryptionMode::CBC, "short")
315 .unwrap_err();
316
317 assert!(matches!(
318 err,
319 crate::Error::InvalidKeyLength {
320 algorithm: "AES128",
321 expected: 16,
322 actual: 5
323 }
324 ));
325 }
326
327 #[test]
328 fn generates_mode_specific_payload_iv() {
329 let cbc = Device::new("aabb")
330 .encrypt(
331 EncryptionAlgorithm::AES128,
332 EncryptionMode::CBC,
333 "1234567890123456",
334 )
335 .unwrap();
336 let gcm = Device::new("aabb")
337 .encrypt(
338 EncryptionAlgorithm::AES128,
339 EncryptionMode::GCM,
340 "1234567890123456",
341 )
342 .unwrap();
343 let ecb = Device::new("aabb")
344 .encrypt(
345 EncryptionAlgorithm::AES128,
346 EncryptionMode::ECB,
347 "1234567890123456",
348 )
349 .unwrap();
350
351 assert_eq!(
352 cbc.encrypt_bark_json(b"test").unwrap().iv.unwrap().len(),
353 16
354 );
355 assert_eq!(
356 gcm.encrypt_bark_json(b"test").unwrap().iv.unwrap().len(),
357 12
358 );
359 assert_eq!(ecb.encrypt_bark_json(b"test").unwrap().iv, None);
360 }
361
362 #[test]
363 fn generates_fresh_iv_for_each_encrypted_payload() {
364 let device = Device::new("aabb")
365 .encrypt(
366 EncryptionAlgorithm::AES128,
367 EncryptionMode::CBC,
368 "1234567890123456",
369 )
370 .unwrap();
371
372 let first = device.encrypt_bark_json(b"test").unwrap();
373 let second = device.encrypt_bark_json(b"test").unwrap();
374
375 assert_ne!(first.iv, second.iv);
376 assert_ne!(first.ciphertext, second.ciphertext);
377 }
378}