kabu_types_entities/
keystore.rs

1use aes::cipher::{Block, BlockDecrypt, KeyInit};
2use aes::Aes128;
3use eyre::{ErrReport, Result};
4use sha2::{Digest, Sha512};
5
6use crate::private::KEY_ENCRYPTION_PWD;
7
8const BLOCK_SIZE: usize = 16;
9
10#[derive(Clone, Default)]
11pub struct KeyStore {
12    pwd: Vec<u8>,
13}
14
15impl KeyStore {
16    pub fn new() -> KeyStore {
17        KeyStore { pwd: KEY_ENCRYPTION_PWD.to_vec() }
18    }
19
20    pub fn new_from_string(pwd: String) -> KeyStore {
21        KeyStore { pwd: pwd.as_bytes().to_vec() }
22    }
23    pub fn new_from_bytes(pwd: Vec<u8>) -> KeyStore {
24        KeyStore { pwd }
25    }
26
27    pub fn encrypt_once(&self, data: &[u8]) -> Result<Vec<u8>> {
28        if self.pwd.is_empty() {
29            return Err(ErrReport::msg("NOT_INITIALIZED"));
30        }
31
32        let mut hasher = Sha512::new();
33        hasher.update(self.pwd.clone());
34        let pwd_hash = hasher.finalize();
35
36        let cipher = Aes128::new_from_slice(&pwd_hash[0..16]).unwrap();
37
38        //println!("{:?}", pwd_hash);
39
40        let mut ret = Vec::new();
41        let mut block: Block<Aes128> = [0u8; BLOCK_SIZE].into();
42
43        let mut a = 0;
44        while a + BLOCK_SIZE <= data.len() {
45            block.copy_from_slice(&data[a..a + BLOCK_SIZE]);
46            cipher.decrypt_block(&mut block);
47            ret.extend_from_slice(&block);
48            a += BLOCK_SIZE;
49        }
50
51        let mut sha = Sha512::new();
52        sha.update(&ret);
53        let crc = &sha.finalize()[0..4];
54
55        if &data[a..a + 4] != crc {
56            return Err(ErrReport::msg("BAD_CHECKSUM"));
57        }
58
59        Ok(ret)
60    }
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66
67    #[test]
68    fn test_encrypt_once_not_initialized() {
69        let key_store = KeyStore::new_from_string(String::from(""));
70        let data = vec![0u8; 36];
71
72        match key_store.encrypt_once(&data) {
73            Ok(_) => panic!("Expected an error, but didn't get one"),
74            Err(e) => assert_eq!(format!("{}", e), "NOT_INITIALIZED"),
75        }
76    }
77
78    #[test]
79    fn test_encrypt_once_bad_checksum() {
80        let key_store = KeyStore::new_from_string(String::from("password"));
81        let data = vec![0u8; 36];
82
83        match key_store.encrypt_once(&data) {
84            Ok(_) => panic!("Expected an error, but didn't get one"),
85            Err(e) => assert_eq!(format!("{}", e), "BAD_CHECKSUM"),
86        }
87    }
88
89    // For this test, you'll need some valid encrypted data to pass and a correct password.
90    #[test]
91    fn test_encrypt_once_valid_data() {
92        let key: Vec<u8> = vec![0x41, 0x8f, 0x2, 0xe4, 0x7e, 0xe4, 0x6, 0xaa, 0xee, 0x71, 0x9e, 0x30, 0xea, 0xe6, 0x64, 0x23];
93        let key_store = KeyStore::new_from_bytes(key);
94        //let encrypted_data = vec![0u8;36]; // Provide valid encrypted data here
95
96        let encrypted_data = hex::decode("51d9dc302b02a02a94d3c7f3057549cd0c990f4c7cc822b61af584fb85afdf209084f48a").unwrap();
97
98        match key_store.encrypt_once(&encrypted_data) {
99            Ok(decrypted_data) => {
100                println!("{}", hex::encode(decrypted_data));
101            }
102            Err(_) => {
103                //println!("{}", hex::encode(decrypted_data));
104                panic!("BAD_CHECKSUM")
105            }
106        }
107    }
108}
kabu_types_entities/keystore.rs

kabu_types_entities/
keystore.rs
