origin_sdk/
crypto.rs

1use crate::random::Random;
2use aes::cipher::{
3    block_padding::{Pkcs7, UnpadError},
4    BlockDecryptMut, BlockEncryptMut, KeyInit,
5};
6use thiserror::Error;
7use tracing::debug;
8
9type EcbEncryptor = ecb::Encryptor<aes::Aes128>;
10type EcbDecryptor = ecb::Decryptor<aes::Aes128>;
11
12const KEY_SIZE: usize = 16;
13const DEFAULT_SEED: u32 = 7;
14
15#[derive(Debug, Error)]
16pub enum CryptoError {
17    #[error("Input cannot be empty")]
18    EmptyInput,
19
20    #[error("Decryption failed: {0}")]
21    DecryptionFailed(String),
22
23    #[error("UTF-8 conversion error")]
24    Utf8Error(#[from] std::string::FromUtf8Error),
25}
26#[derive(Clone, Debug)]
27pub struct Crypto {
28    key: [u8; KEY_SIZE],
29}
30
31impl Crypto {
32    pub fn new(seed: u32) -> Self {
33        let mut crypto = Self {
34            key: [0u8; KEY_SIZE],
35        };
36
37        crypto.set_key(seed);
38        crypto
39    }
40
41    pub fn set_key(&mut self, seed: u32) {
42        let mut key = [0u8; KEY_SIZE];
43
44        if seed == 0 {
45            for (i, byte) in key.iter_mut().enumerate() {
46                *byte = i as u8
47            }
48        } else {
49            let mut rng = Random::new(DEFAULT_SEED);
50            let new_seed = rng.next().wrapping_add(seed);
51            rng.set_seed(new_seed);
52
53            for byte in key.iter_mut() {
54                *byte = rng.next() as u8
55            }
56        }
57
58        debug!("Setting new encryption key: {:?}", hex::encode(key));
59        self.key = key;
60    }
61
62    pub fn encrypt(&self, plain_text: &str) -> Result<Vec<u8>, CryptoError> {
63        if plain_text.is_empty() {
64            return Err(CryptoError::EmptyInput);
65        }
66
67        let res = EcbEncryptor::new(&self.key.into())
68            .encrypt_padded_vec_mut::<Pkcs7>(plain_text.as_bytes());
69
70        Ok(res)
71    }
72
73    pub fn decrypt(&self, cipher_text: &[u8]) -> Result<String, CryptoError> {
74        if cipher_text.is_empty() {
75            return Err(CryptoError::EmptyInput);
76        }
77
78        let res = EcbDecryptor::new(&self.key.into())
79            .decrypt_padded_vec_mut::<Pkcs7>(cipher_text)
80            .map_err(|e: UnpadError| CryptoError::DecryptionFailed(e.to_string()))?;
81
82        String::from_utf8(res).map_err(CryptoError::from)
83    }
84
85    pub fn prepare_challenge_response(&mut self, key: &str) -> Result<String, CryptoError> {
86        let response_key = self.encrypt(key)?;
87        let response_str = hex::encode(&response_key);
88        let response_bytes = response_str.as_bytes();
89
90        let seed = ((response_bytes[0] as u32) << 8) | (response_bytes[1] as u32);
91        self.set_key(seed);
92
93        Ok(response_str)
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[test]
102    fn test_key() {
103        let crypto = Crypto::new(0);
104
105        assert_eq!(crypto.key.len(), KEY_SIZE);
106        assert_eq!(
107            &crypto.key as &[u8],
108            [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
109        );
110
111        let crypto = Crypto::new(1337);
112
113        assert_eq!(crypto.key.len(), KEY_SIZE);
114        assert_eq!(
115            &crypto.key as &[u8],
116            [251, 135, 22, 197, 214, 181, 148, 115, 149, 93, 40, 78, 123, 141, 60, 108]
117        );
118    }
119
120    #[test]
121    fn test_encrypt_decrypt() {
122        let crypto = Crypto::new(1337);
123
124        let plain_text = "hello world".to_string();
125        let cipher_text = crypto.encrypt(&plain_text).expect("Failed to encrypt");
126        let decrypted_text = crypto.decrypt(&cipher_text).expect("Failed to decrypt");
127
128        assert_eq!(plain_text, decrypted_text);
129    }
130}
origin_sdk/crypto.rs

origin_sdk/
crypto.rs
