1use std::{fmt::Display, str::FromStr};
2
3use aes::cipher::typenum::U32;
4use base64::{engine::general_purpose::STANDARD, Engine};
5use generic_array::GenericArray;
6use serde::Deserialize;
7
8use super::{check_length, from_b64, from_b64_vec, split_enc_string};
9use crate::{
10    error::{CryptoError, EncStringParseError, Result},
11    KeyDecryptable, KeyEncryptable, LocateKey, SymmetricCryptoKey,
12};
13
14#[derive(Clone, zeroize::ZeroizeOnDrop, PartialEq)]
48#[allow(unused, non_camel_case_types)]
49pub enum EncString {
50    AesCbc256_B64 { iv: [u8; 16], data: Vec<u8> },
52    AesCbc128_HmacSha256_B64 {
54        iv: [u8; 16],
55        mac: [u8; 32],
56        data: Vec<u8>,
57    },
58    AesCbc256_HmacSha256_B64 {
60        iv: [u8; 16],
61        mac: [u8; 32],
62        data: Vec<u8>,
63    },
64}
65
66impl std::fmt::Debug for EncString {
68    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
69        f.debug_struct("EncString").finish()
70    }
71}
72
73impl FromStr for EncString {
75    type Err = CryptoError;
76
77    fn from_str(s: &str) -> Result<Self, Self::Err> {
78        let (enc_type, parts) = split_enc_string(s);
79        match (enc_type, parts.len()) {
80            ("0", 2) => {
81                let iv = from_b64(parts[0])?;
82                let data = from_b64_vec(parts[1])?;
83
84                Ok(EncString::AesCbc256_B64 { iv, data })
85            }
86            ("1" | "2", 3) => {
87                let iv = from_b64(parts[0])?;
88                let data = from_b64_vec(parts[1])?;
89                let mac = from_b64(parts[2])?;
90
91                if enc_type == "1" {
92                    Ok(EncString::AesCbc128_HmacSha256_B64 { iv, mac, data })
93                } else {
94                    Ok(EncString::AesCbc256_HmacSha256_B64 { iv, mac, data })
95                }
96            }
97
98            (enc_type, parts) => Err(EncStringParseError::InvalidTypeSymm {
99                enc_type: enc_type.to_string(),
100                parts,
101            }
102            .into()),
103        }
104    }
105}
106
107impl EncString {
108    pub fn try_from_optional(s: Option<String>) -> Result<Option<EncString>, CryptoError> {
110        s.map(|s| s.parse()).transpose()
111    }
112
113    pub fn from_buffer(buf: &[u8]) -> Result<Self> {
114        if buf.is_empty() {
115            return Err(EncStringParseError::NoType.into());
116        }
117        let enc_type = buf[0];
118
119        match enc_type {
120            0 => {
121                check_length(buf, 18)?;
122                let iv = buf[1..17].try_into().expect("Valid length");
123                let data = buf[17..].to_vec();
124
125                Ok(EncString::AesCbc256_B64 { iv, data })
126            }
127            1 | 2 => {
128                check_length(buf, 50)?;
129                let iv = buf[1..17].try_into().expect("Valid length");
130                let mac = buf[17..49].try_into().expect("Valid length");
131                let data = buf[49..].to_vec();
132
133                if enc_type == 1 {
134                    Ok(EncString::AesCbc128_HmacSha256_B64 { iv, mac, data })
135                } else {
136                    Ok(EncString::AesCbc256_HmacSha256_B64 { iv, mac, data })
137                }
138            }
139            _ => Err(EncStringParseError::InvalidTypeSymm {
140                enc_type: enc_type.to_string(),
141                parts: 1,
142            }
143            .into()),
144        }
145    }
146
147    pub fn to_buffer(&self) -> Result<Vec<u8>> {
148        let mut buf;
149
150        match self {
151            EncString::AesCbc256_B64 { iv, data } => {
152                buf = Vec::with_capacity(1 + 16 + data.len());
153                buf.push(self.enc_type());
154                buf.extend_from_slice(iv);
155                buf.extend_from_slice(data);
156            }
157            EncString::AesCbc128_HmacSha256_B64 { iv, mac, data }
158            | EncString::AesCbc256_HmacSha256_B64 { iv, mac, data } => {
159                buf = Vec::with_capacity(1 + 16 + 32 + data.len());
160                buf.push(self.enc_type());
161                buf.extend_from_slice(iv);
162                buf.extend_from_slice(mac);
163                buf.extend_from_slice(data);
164            }
165        }
166
167        Ok(buf)
168    }
169}
170
171impl Display for EncString {
172    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
173        let parts: Vec<&[u8]> = match self {
174            EncString::AesCbc256_B64 { iv, data } => vec![iv, data],
175            EncString::AesCbc128_HmacSha256_B64 { iv, mac, data } => vec![iv, data, mac],
176            EncString::AesCbc256_HmacSha256_B64 { iv, mac, data } => vec![iv, data, mac],
177        };
178
179        let encoded_parts: Vec<String> = parts.iter().map(|part| STANDARD.encode(part)).collect();
180
181        write!(f, "{}.{}", self.enc_type(), encoded_parts.join("|"))?;
182
183        Ok(())
184    }
185}
186
187impl<'de> Deserialize<'de> for EncString {
188    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
189    where
190        D: serde::Deserializer<'de>,
191    {
192        deserializer.deserialize_str(super::FromStrVisitor::new())
193    }
194}
195
196impl serde::Serialize for EncString {
197    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
198    where
199        S: serde::Serializer,
200    {
201        serializer.serialize_str(&self.to_string())
202    }
203}
204
205impl EncString {
206    pub(crate) fn encrypt_aes256_hmac(
207        data_dec: &[u8],
208        mac_key: &GenericArray<u8, U32>,
209        key: &GenericArray<u8, U32>,
210    ) -> Result<EncString> {
211        let (iv, mac, data) = crate::aes::encrypt_aes256_hmac(data_dec, mac_key, key)?;
212        Ok(EncString::AesCbc256_HmacSha256_B64 { iv, mac, data })
213    }
214
215    const fn enc_type(&self) -> u8 {
217        match self {
218            EncString::AesCbc256_B64 { .. } => 0,
219            EncString::AesCbc128_HmacSha256_B64 { .. } => 1,
220            EncString::AesCbc256_HmacSha256_B64 { .. } => 2,
221        }
222    }
223}
224
225impl LocateKey for EncString {}
226impl KeyEncryptable<SymmetricCryptoKey, EncString> for &[u8] {
227    fn encrypt_with_key(self, key: &SymmetricCryptoKey) -> Result<EncString> {
228        EncString::encrypt_aes256_hmac(
229            self,
230            key.mac_key.as_ref().ok_or(CryptoError::InvalidMac)?,
231            &key.key,
232        )
233    }
234}
235
236impl KeyDecryptable<SymmetricCryptoKey, Vec<u8>> for EncString {
237    fn decrypt_with_key(&self, key: &SymmetricCryptoKey) -> Result<Vec<u8>> {
238        match self {
239            EncString::AesCbc256_B64 { iv, data } => {
240                if key.mac_key.is_some() {
241                    return Err(CryptoError::MacNotProvided);
242                }
243
244                let dec = crate::aes::decrypt_aes256(iv, data.clone(), &key.key)?;
245                Ok(dec)
246            }
247            EncString::AesCbc128_HmacSha256_B64 { iv, mac, data } => {
248                let enc_key = key.key[0..16].into();
253                let mac_key = key.key[16..32].into();
254                let dec = crate::aes::decrypt_aes128_hmac(iv, mac, data.clone(), mac_key, enc_key)?;
255                Ok(dec)
256            }
257            EncString::AesCbc256_HmacSha256_B64 { iv, mac, data } => {
258                let mac_key = key.mac_key.as_ref().ok_or(CryptoError::InvalidMac)?;
259                let dec =
260                    crate::aes::decrypt_aes256_hmac(iv, mac, data.clone(), mac_key, &key.key)?;
261                Ok(dec)
262            }
263        }
264    }
265}
266
267impl KeyEncryptable<SymmetricCryptoKey, EncString> for String {
268    fn encrypt_with_key(self, key: &SymmetricCryptoKey) -> Result<EncString> {
269        self.as_bytes().encrypt_with_key(key)
270    }
271}
272
273impl KeyEncryptable<SymmetricCryptoKey, EncString> for &str {
274    fn encrypt_with_key(self, key: &SymmetricCryptoKey) -> Result<EncString> {
275        self.as_bytes().encrypt_with_key(key)
276    }
277}
278
279impl KeyDecryptable<SymmetricCryptoKey, String> for EncString {
280    fn decrypt_with_key(&self, key: &SymmetricCryptoKey) -> Result<String> {
281        let dec: Vec<u8> = self.decrypt_with_key(key)?;
282        String::from_utf8(dec).map_err(|_| CryptoError::InvalidUtf8String)
283    }
284}
285
286impl schemars::JsonSchema for EncString {
289    fn schema_name() -> String {
290        "EncString".to_string()
291    }
292
293    fn json_schema(gen: &mut schemars::gen::SchemaGenerator) -> schemars::schema::Schema {
294        gen.subschema_for::<String>()
295    }
296}
297
298#[cfg(test)]
299mod tests {
300    use schemars::schema_for;
301
302    use super::EncString;
303    use crate::{
304        derive_symmetric_key, CryptoError, KeyDecryptable, KeyEncryptable, SymmetricCryptoKey,
305    };
306
307    #[test]
308    fn test_enc_string_roundtrip() {
309        let key = derive_symmetric_key("test");
310
311        let test_string = "encrypted_test_string";
312        let cipher = test_string.to_owned().encrypt_with_key(&key).unwrap();
313
314        let decrypted_str: String = cipher.decrypt_with_key(&key).unwrap();
315        assert_eq!(decrypted_str, test_string);
316    }
317
318    #[test]
319    fn test_enc_string_ref_roundtrip() {
320        let key = derive_symmetric_key("test");
321
322        let test_string = "encrypted_test_string";
323        let cipher = test_string.encrypt_with_key(&key).unwrap();
324
325        let decrypted_str: String = cipher.decrypt_with_key(&key).unwrap();
326        assert_eq!(decrypted_str, test_string);
327    }
328
329    #[test]
330    fn test_enc_string_serialization() {
331        #[derive(serde::Serialize, serde::Deserialize)]
332        struct Test {
333            key: EncString,
334        }
335
336        let cipher = "2.pMS6/icTQABtulw52pq2lg==|XXbxKxDTh+mWiN1HjH2N1w==|Q6PkuT+KX/axrgN9ubD5Ajk2YNwxQkgs3WJM0S0wtG8=";
337        let serialized = format!("{{\"key\":\"{cipher}\"}}");
338
339        let t = serde_json::from_str::<Test>(&serialized).unwrap();
340        assert_eq!(t.key.enc_type(), 2);
341        assert_eq!(t.key.to_string(), cipher);
342        assert_eq!(serde_json::to_string(&t).unwrap(), serialized);
343    }
344
345    #[test]
346    fn test_enc_from_to_buffer() {
347        let enc_str: &str = "2.pMS6/icTQABtulw52pq2lg==|XXbxKxDTh+mWiN1HjH2N1w==|Q6PkuT+KX/axrgN9ubD5Ajk2YNwxQkgs3WJM0S0wtG8=";
348        let enc_string: EncString = enc_str.parse().unwrap();
349
350        let enc_buf = enc_string.to_buffer().unwrap();
351
352        assert_eq!(
353            enc_buf,
354            vec![
355                2, 164, 196, 186, 254, 39, 19, 64, 0, 109, 186, 92, 57, 218, 154, 182, 150, 67,
356                163, 228, 185, 63, 138, 95, 246, 177, 174, 3, 125, 185, 176, 249, 2, 57, 54, 96,
357                220, 49, 66, 72, 44, 221, 98, 76, 209, 45, 48, 180, 111, 93, 118, 241, 43, 16, 211,
358                135, 233, 150, 136, 221, 71, 140, 125, 141, 215
359            ]
360        );
361
362        let enc_string_new = EncString::from_buffer(&enc_buf).unwrap();
363
364        assert_eq!(enc_string_new.to_string(), enc_str)
365    }
366
367    #[test]
368    fn test_from_str_cbc256() {
369        let enc_str = "0.pMS6/icTQABtulw52pq2lg==|XXbxKxDTh+mWiN1HjH2N1w==";
370        let enc_string: EncString = enc_str.parse().unwrap();
371
372        assert_eq!(enc_string.enc_type(), 0);
373        if let EncString::AesCbc256_B64 { iv, data } = &enc_string {
374            assert_eq!(
375                iv,
376                &[164, 196, 186, 254, 39, 19, 64, 0, 109, 186, 92, 57, 218, 154, 182, 150]
377            );
378            assert_eq!(
379                data,
380                &[93, 118, 241, 43, 16, 211, 135, 233, 150, 136, 221, 71, 140, 125, 141, 215]
381            );
382        } else {
383            panic!("Invalid variant")
384        };
385    }
386
387    #[test]
388    fn test_from_str_cbc128_hmac() {
389        let enc_str = "1.Hh8gISIjJCUmJygpKissLQ==|MjM0NTY3ODk6Ozw9Pj9AQUJDREU=|KCkqKywtLi8wMTIzNDU2Nzg5Ojs8PT4/QEFCQ0RFRkc=";
390        let enc_string: EncString = enc_str.parse().unwrap();
391
392        assert_eq!(enc_string.enc_type(), 1);
393        if let EncString::AesCbc128_HmacSha256_B64 { iv, mac, data } = &enc_string {
394            assert_eq!(
395                iv,
396                &[30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45]
397            );
398            assert_eq!(
399                mac,
400                &[
401                    40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59,
402                    60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71
403                ]
404            );
405            assert_eq!(
406                data,
407                &[50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69]
408            );
409        } else {
410            panic!("Invalid variant")
411        };
412    }
413
414    #[test]
415    fn test_decrypt_cbc256() {
416        let key = "hvBMMb1t79YssFZkpetYsM3deyVuQv4r88Uj9gvYe08=".to_string();
417        let key = SymmetricCryptoKey::try_from(key).unwrap();
418
419        let enc_str = "0.NQfjHLr6za7VQVAbrpL81w==|wfrjmyJ0bfwkQlySrhw8dA==";
420        let enc_string: EncString = enc_str.parse().unwrap();
421        assert_eq!(enc_string.enc_type(), 0);
422
423        let dec_str: String = enc_string.decrypt_with_key(&key).unwrap();
424        assert_eq!(dec_str, "EncryptMe!");
425    }
426
427    #[test]
428    fn test_decrypt_downgrade_encstring_prevention() {
429        let key = "hvBMMb1t79YssFZkpetYsM3deyVuQv4r88Uj9gvYe0+G8EwxvW3v1iywVmSl61iwzd17JW5C/ivzxSP2C9h7Tw==".to_string();
432        let key = SymmetricCryptoKey::try_from(key).unwrap();
433
434        let enc_str = "0.NQfjHLr6za7VQVAbrpL81w==|wfrjmyJ0bfwkQlySrhw8dA==";
438        let enc_string: EncString = enc_str.parse().unwrap();
439        assert_eq!(enc_string.enc_type(), 0);
440
441        let result: Result<String, CryptoError> = enc_string.decrypt_with_key(&key);
442        assert!(matches!(result, Err(CryptoError::MacNotProvided)));
443    }
444
445    #[test]
446    fn test_decrypt_cbc128_hmac() {
447        let key = "Gt1aZ8kTTgkF80bLtb7LiMZBcxEA2FA5mbvV4x7K208=".to_string();
448        let key = SymmetricCryptoKey::try_from(key).unwrap();
449
450        let enc_str = "1.CU/oG4VZuxbHoZSDZjCLQw==|kb1HGwAk+fQ275ORfLf5Ew==|8UaEYHyqRZcG37JWhYBOBdEatEXd1u1/wN7OuImolcM=";
451        let enc_string: EncString = enc_str.parse().unwrap();
452        assert_eq!(enc_string.enc_type(), 1);
453
454        let dec_str: String = enc_string.decrypt_with_key(&key).unwrap();
455        assert_eq!(dec_str, "EncryptMe!");
456    }
457
458    #[test]
459    fn test_from_str_invalid() {
460        let enc_str = "7.ABC";
461        let enc_string: Result<EncString, _> = enc_str.parse();
462
463        let err = enc_string.unwrap_err();
464        assert_eq!(
465            err.to_string(),
466            "EncString error, Invalid symmetric type, got type 7 with 1 parts"
467        );
468    }
469
470    #[test]
471    fn test_debug_format() {
472        let enc_str  = "2.pMS6/icTQABtulw52pq2lg==|XXbxKxDTh+mWiN1HjH2N1w==|Q6PkuT+KX/axrgN9ubD5Ajk2YNwxQkgs3WJM0S0wtG8=";
473        let enc_string: EncString = enc_str.parse().unwrap();
474
475        let debug_string = format!("{:?}", enc_string);
476        assert_eq!(debug_string, "EncString");
477    }
478
479    #[test]
480    fn test_json_schema() {
481        let schema = schema_for!(EncString);
482
483        assert_eq!(
484            serde_json::to_string(&schema).unwrap(),
485            r#"{"$schema":"http://json-schema.org/draft-07/schema#","title":"EncString","type":"string"}"#
486        );
487    }
488}