1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

// This file is part of the laron-crypto
//
// Copyright 2023 Ade M Ramdani
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

use aes_gcm::{aead::Aead, Aes256Gcm, KeyInit, Nonce};
use rand::RngCore;

use crate::{
    error::{Error, Result},
    PrivateKey,
};

/// Aes trait provides the interface for AES encryption and
/// decryption.
pub trait Aes {
    /// Encrypt the given string.
    fn encrypt(&self, data: &str) -> String;

    /// Encrypt with the given nonce.
    fn encrypt_with_nonce(&self, data: &str, nonce: &str) -> Result<String>;

    /// Decrypt the given string.
    fn decrypt(&self, data: &str) -> Result<String>;

    ///  Decrypt with the given nonce.
    fn decrypt_with_nonce(&self, data: &str, nonce: &str) -> Result<String>;
}

impl From<&PrivateKey> for Aes256Gcm {
    fn from(value: &PrivateKey) -> Self {
        Self::new_from_slice(value.as_slice()).unwrap()
    }
}

impl Aes for PrivateKey {
    fn encrypt(&self, data: &str) -> String {
        let mut nonce_bytes = [0u8; 12];
        rand::thread_rng().fill_bytes(&mut nonce_bytes);
        let nonce = Nonce::from_slice(&nonce_bytes);
        let key: Aes256Gcm = self.into();
        let ciphertext = key.encrypt(nonce, data.as_bytes()).unwrap();
        let mut result = vec![];
        result.extend_from_slice(&nonce_bytes);
        result.extend_from_slice(&ciphertext);
        hex::encode(result)
    }

    fn encrypt_with_nonce(&self, data: &str, nonce: &str) -> Result<String> {
        let nonce = Nonce::from_slice(nonce.as_bytes());
        let key: Aes256Gcm = self.into();
        let ciphertext = key
            .encrypt(nonce, data.as_bytes())
            .map_err(|e| Error::EncryptError(e.to_string()))?;
        Ok(hex::encode(ciphertext))
    }

    fn decrypt(&self, data: &str) -> Result<String> {
        let data = hex::decode(data)?;
        if data.len() < 12 {
            return Err(Error::LengthError(
                "Encrypted data is too short".to_string(),
            ));
        }
        let nonce = Nonce::from_slice(&data[0..12]);
        let key: Aes256Gcm = self.into();
        let plaintext = key
            .decrypt(nonce, &data[12..])
            .map_err(|e| Error::DecryptError(e.to_string()))?;
        Ok(String::from_utf8(plaintext).unwrap())
    }

    fn decrypt_with_nonce(&self, data: &str, nonce: &str) -> Result<String> {
        let data = hex::decode(data)?;
        let nonce = Nonce::from_slice(nonce.as_bytes());
        let key: Aes256Gcm = self.into();
        let plaintext = key
            .decrypt(nonce, data.as_ref())
            .map_err(|e| Error::DecryptError(e.to_string()))?;
        Ok(String::from_utf8(plaintext).unwrap())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_aes() {
        let text = "Hello World!";
        let key = PrivateKey::new();
        let encrypted = key.encrypt(text);
        let decrypted = key.decrypt(&encrypted).unwrap();
        assert_eq!(text, decrypted);
        let nonce = "unique nonce";
        let encrypted = key.encrypt_with_nonce(text, nonce).unwrap();
        let decrypted = key.decrypt_with_nonce(&encrypted, nonce).unwrap();
        assert_eq!(text, decrypted);
    }
}