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

use aes::{
    cipher::{NewCipher, StreamCipher, StreamCipherSeek},
    Aes128, Aes128Ctr, Aes256, Aes256Ctr,
};
use anyhow::*;
use block_modes::{block_padding::Pkcs7, BlockMode, Cbc};
use wasm_bindgen::{prelude::*, throw_str};

#[wasm_bindgen]
pub struct AES {}

#[wasm_bindgen]
#[allow(non_camel_case_types)]
#[derive(Debug, Clone, Copy)]
pub enum AESAlgorithms {
    AES128_CBC,
    AES256_CBC,
    AES128_CTR,
    AES256_CTR,
}

impl AES {
    pub fn encrypt_impl(key: &[u8], iv: &[u8], message: &[u8], algo: AESAlgorithms) -> Result<Vec<u8>> {
        let result = match algo {
            AESAlgorithms::AES128_CBC => Cbc::<Aes128, Pkcs7>::new_from_slices(key, iv)?.encrypt_vec(&message),
            AESAlgorithms::AES256_CBC => Cbc::<Aes256, Pkcs7>::new_from_slices(key, iv)?.encrypt_vec(&message),
            AESAlgorithms::AES128_CTR => AES::aes_ctr::<Aes128Ctr>(key, iv, message),
            AESAlgorithms::AES256_CTR => AES::aes_ctr::<Aes256Ctr>(key, iv, message),
        };

        Ok(result)
    }

    pub fn decrypt_impl(key: &[u8], iv: &[u8], message: &[u8], algo: AESAlgorithms) -> Result<Vec<u8>> {
        let result = match algo {
            AESAlgorithms::AES128_CBC => Cbc::<Aes128, Pkcs7>::new_from_slices(key, iv)?.decrypt_vec(message)?,
            AESAlgorithms::AES256_CBC => Cbc::<Aes256, Pkcs7>::new_from_slices(key, iv)?.decrypt_vec(message)?,
            AESAlgorithms::AES128_CTR => AES::aes_ctr::<Aes128Ctr>(key, iv, message),
            AESAlgorithms::AES256_CTR => AES::aes_ctr::<Aes256Ctr>(key, iv, message),
        };

        Ok(result)
    }

    fn aes_ctr<T: NewCipher + StreamCipherSeek + StreamCipher>(key: &[u8], iv: &[u8], message: &[u8]) -> Vec<u8> {
        let data = &mut message.to_vec();
        let mut cipher = T::new(key.into(), iv.into());
        cipher.seek(0);
        cipher.apply_keystream(data);
        data.to_vec()
    }
}

#[cfg(target_arch = "wasm32")]
#[wasm_bindgen]
impl AES {
    #[wasm_bindgen(js_name = encrypt)]
    pub fn encrypt(key: &[u8], iv: &[u8], message: &[u8], algo: AESAlgorithms) -> Result<Vec<u8>, JsValue> {
        match AES::encrypt_impl(key, iv, message, algo) {
            Ok(v) => Ok(v),
            Err(e) => throw_str(&e.to_string()),
        }
    }

    #[wasm_bindgen(js_name = decrypt)]
    pub fn decrypt(key: &[u8], iv: &[u8], message: &[u8], algo: AESAlgorithms) -> Result<Vec<u8>, JsValue> {
        match AES::decrypt_impl(key, iv, message, algo) {
            Ok(v) => Ok(v),
            Err(e) => throw_str(&e.to_string()),
        }
    }
}

#[cfg(not(target_arch = "wasm32"))]
impl AES {
    pub fn encrypt(key: &[u8], iv: &[u8], message: &[u8], algo: AESAlgorithms) -> Result<Vec<u8>> {
        AES::encrypt_impl(key, iv, message, algo)
    }

    pub fn decrypt(key: &[u8], iv: &[u8], message: &[u8], algo: AESAlgorithms) -> Result<Vec<u8>> {
        AES::decrypt_impl(key, iv, message, algo)
    }
}