1use aes::{Aes128, Aes192, Aes256};
2use cfb_mode::Cfb;
3use chacha20::ChaCha20;
4use cipher::{consts::U16, AsyncStreamCipher, BlockCipher, BlockEncrypt, NewCipher, StreamCipher};
5use ctr::Ctr128BE;
6use rand::distributions::Standard;
7use rand::{thread_rng, Rng};
8
9use crate::util::generate_key;
10
11pub trait SymmetricCipher {
12 fn encrypt(&mut self, data: &mut [u8]);
14
15 fn decrypt(&mut self, data: &mut [u8]);
17}
18
19impl<C: BlockCipher + BlockEncrypt> SymmetricCipher for Cfb<C> {
20 fn encrypt(&mut self, data: &mut [u8]) {
22 AsyncStreamCipher::encrypt(self, data)
23 }
24
25 fn decrypt(&mut self, data: &mut [u8]) {
27 AsyncStreamCipher::decrypt(self, data)
28 }
29}
30
31impl<B: BlockEncrypt + BlockCipher<BlockSize = U16>> SymmetricCipher for Ctr128BE<B> {
32 fn encrypt(&mut self, data: &mut [u8]) {
34 StreamCipher::apply_keystream(self, data)
35 }
36
37 fn decrypt(&mut self, data: &mut [u8]) {
39 StreamCipher::apply_keystream(self, data)
40 }
41}
42
43impl SymmetricCipher for ChaCha20 {
44 fn encrypt(&mut self, data: &mut [u8]) {
46 StreamCipher::apply_keystream(self, data)
47 }
48
49 fn decrypt(&mut self, data: &mut [u8]) {
51 StreamCipher::apply_keystream(self, data)
52 }
53}
54
55type Aes128Cfb = Cfb<Aes128>;
56type Aes192Cfb = Cfb<Aes192>;
57type Aes256Cfb = Cfb<Aes256>;
58type Aes128Ctr = Ctr128BE<Aes128>;
59type Aes192Ctr = Ctr128BE<Aes192>;
60type Aes256Ctr = Ctr128BE<Aes256>;
61
62pub struct Cipher {
63 pub key: Vec<u8>,
64 pub key_len: usize,
65 pub iv: Vec<u8>,
66 pub iv_len: usize,
67 pub enc: Option<Box<dyn SymmetricCipher + Send + 'static>>,
68 pub dec: Option<Box<dyn SymmetricCipher + Send + 'static>>,
69 cipher_method: CipherMethod,
70}
71
72#[derive(Clone, Copy, Debug)]
73enum CipherMethod {
74 Aes128Cfb,
75 Aes192Cfb,
76 Aes256Cfb,
77 Aes128Ctr,
78 Aes192Ctr,
79 Aes256Ctr,
80 ChaCha20,
81}
82
83impl Cipher {
84 pub fn new(method: &str, password: &str) -> Cipher {
85 let (key_len, cipher_method, iv_len) = match method {
86 "aes-128-cfb" => (16, CipherMethod::Aes128Cfb, 16),
87 "aes-192-cfb" => (24, CipherMethod::Aes192Cfb, 16),
88 "aes-256-cfb" => (32, CipherMethod::Aes256Cfb, 16),
89 "aes-128-ctr" => (16, CipherMethod::Aes128Ctr, 16),
90 "aes-192-ctr" => (24, CipherMethod::Aes192Ctr, 16),
91 "aes-256-ctr" => (32, CipherMethod::Aes256Ctr, 16),
92 "chacha20" => (32, CipherMethod::ChaCha20, 12),
93 _ => panic!("method not supported"),
94 };
95
96 let key = generate_key(password.as_bytes(), key_len);
97 Cipher {
98 key: Vec::from(&key[..]),
99 key_len,
100 iv_len,
101 iv: vec![0u8; iv_len],
102 enc: None,
103 dec: None,
104 cipher_method,
105 }
106 }
107
108 pub fn init_encrypt(&mut self) {
109 if self.iv.is_empty() {
110 let rng = thread_rng();
111 self.iv = rng.sample_iter(&Standard).take(self.iv_len).collect();
112 }
113 self.enc = Some(self.new_cipher(&self.iv));
114 }
115
116 fn new_cipher(&self, iv: &[u8]) -> Box<dyn SymmetricCipher + Send + 'static> {
117 let key: &[u8] = &self.key;
118 match self.cipher_method {
119 CipherMethod::Aes128Cfb => {
120 Box::new(Aes128Cfb::new_from_slices(key, iv).expect("init cipher error"))
121 }
122 CipherMethod::Aes192Cfb => {
123 Box::new(Aes192Cfb::new_from_slices(key, iv).expect("init cipher error"))
124 }
125 CipherMethod::Aes256Cfb => {
126 Box::new(Aes256Cfb::new_from_slices(key, iv).expect("init cipher error"))
127 }
128 CipherMethod::Aes128Ctr => Box::new(Aes128Ctr::new(key.into(), iv.into())),
129 CipherMethod::Aes192Ctr => Box::new(Aes192Ctr::new(key.into(), iv.into())),
130 CipherMethod::Aes256Ctr => Box::new(Aes256Ctr::new(key.into(), iv.into())),
131 CipherMethod::ChaCha20 => Box::new(ChaCha20::new(key.into(), iv.into())),
132 }
133 }
134
135 pub fn init_decrypt(&mut self) {
136 self.dec = Some(self.new_cipher(&self.iv));
137 }
138
139 pub fn encrypt(&mut self, input: &mut [u8]) {
140 if let Some(enc) = &mut self.enc {
141 enc.encrypt(input);
142 }
143 }
144
145 pub fn decrypt(&mut self, input: &mut [u8]) {
146 if let Some(dec) = &mut self.dec {
147 dec.decrypt(input)
148 }
149 }
150
151 pub fn reset(&self) -> Cipher {
152 Cipher {
153 key: self.key.clone(),
154 iv: vec![0u8; self.iv_len],
155 iv_len: self.iv_len,
156 key_len: self.key_len,
157 enc: None,
158 dec: None,
159 cipher_method: self.cipher_method,
160 }
161 }
162}