1use aes_gcm::{aead::Aead, Aes256Gcm, KeyInit, Nonce};
19use rand::RngCore;
20
21use crate::{
22 error::{Error, Result},
23 PrivateKey,
24};
25
26pub trait Aes {
29 fn encrypt(&self, data: &str) -> String;
31
32 fn encrypt_with_nonce(&self, data: &str, nonce: &str) -> Result<String>;
34
35 fn decrypt(&self, data: &str) -> Result<String>;
37
38 fn decrypt_with_nonce(&self, data: &str, nonce: &str) -> Result<String>;
40}
41
42impl From<&PrivateKey> for Aes256Gcm {
43 fn from(value: &PrivateKey) -> Self {
44 Self::new_from_slice(value.as_slice()).unwrap()
45 }
46}
47
48impl Aes for PrivateKey {
49 fn encrypt(&self, data: &str) -> String {
50 let mut nonce_bytes = [0u8; 12];
51 rand::thread_rng().fill_bytes(&mut nonce_bytes);
52 let nonce = Nonce::from_slice(&nonce_bytes);
53 let key: Aes256Gcm = self.into();
54 let ciphertext = key.encrypt(nonce, data.as_bytes()).unwrap();
55 let mut result = vec![];
56 result.extend_from_slice(&nonce_bytes);
57 result.extend_from_slice(&ciphertext);
58 hex::encode(result)
59 }
60
61 fn encrypt_with_nonce(&self, data: &str, nonce: &str) -> Result<String> {
62 let nonce = Nonce::from_slice(nonce.as_bytes());
63 let key: Aes256Gcm = self.into();
64 let ciphertext = key
65 .encrypt(nonce, data.as_bytes())
66 .map_err(|e| Error::EncryptError(e.to_string()))?;
67 Ok(hex::encode(ciphertext))
68 }
69
70 fn decrypt(&self, data: &str) -> Result<String> {
71 let data = hex::decode(data)?;
72 if data.len() < 12 {
73 return Err(Error::LengthError(
74 "Encrypted data is too short".to_string(),
75 ));
76 }
77 let nonce = Nonce::from_slice(&data[0..12]);
78 let key: Aes256Gcm = self.into();
79 let plaintext = key
80 .decrypt(nonce, &data[12..])
81 .map_err(|e| Error::DecryptError(e.to_string()))?;
82 Ok(String::from_utf8(plaintext).unwrap())
83 }
84
85 fn decrypt_with_nonce(&self, data: &str, nonce: &str) -> Result<String> {
86 let data = hex::decode(data)?;
87 let nonce = Nonce::from_slice(nonce.as_bytes());
88 let key: Aes256Gcm = self.into();
89 let plaintext = key
90 .decrypt(nonce, data.as_ref())
91 .map_err(|e| Error::DecryptError(e.to_string()))?;
92 Ok(String::from_utf8(plaintext).unwrap())
93 }
94}
95
96#[cfg(test)]
97mod tests {
98 use super::*;
99
100 #[test]
101 fn test_aes() {
102 let text = "Hello World!";
103 let key = PrivateKey::new();
104 let encrypted = key.encrypt(text);
105 let decrypted = key.decrypt(&encrypted).unwrap();
106 assert_eq!(text, decrypted);
107 let nonce = "unique nonce";
108 let encrypted = key.encrypt_with_nonce(text, nonce).unwrap();
109 let decrypted = key.decrypt_with_nonce(&encrypted, nonce).unwrap();
110 assert_eq!(text, decrypted);
111 }
112}