use aes::{Aes128, Aes192, Aes256};
use cfb_mode::Cfb;
use chacha20::ChaCha20;
use cipher::{consts::U16, AsyncStreamCipher, BlockCipher, BlockEncrypt, NewCipher, StreamCipher};
use ctr::Ctr128BE;
use rand::distributions::Standard;
use rand::{thread_rng, Rng};
use crate::util::generate_key;
pub trait SymmetricCipher {
fn encrypt(&mut self, data: &mut [u8]);
fn decrypt(&mut self, data: &mut [u8]);
}
impl<C: BlockCipher + BlockEncrypt> SymmetricCipher for Cfb<C> {
fn encrypt(&mut self, data: &mut [u8]) {
AsyncStreamCipher::encrypt(self, data)
}
fn decrypt(&mut self, data: &mut [u8]) {
AsyncStreamCipher::decrypt(self, data)
}
}
impl<B: BlockEncrypt + BlockCipher<BlockSize = U16>> SymmetricCipher for Ctr128BE<B> {
fn encrypt(&mut self, data: &mut [u8]) {
StreamCipher::apply_keystream(self, data)
}
fn decrypt(&mut self, data: &mut [u8]) {
StreamCipher::apply_keystream(self, data)
}
}
impl SymmetricCipher for ChaCha20 {
fn encrypt(&mut self, data: &mut [u8]) {
StreamCipher::apply_keystream(self, data)
}
fn decrypt(&mut self, data: &mut [u8]) {
StreamCipher::apply_keystream(self, data)
}
}
type Aes128Cfb = Cfb<Aes128>;
type Aes192Cfb = Cfb<Aes192>;
type Aes256Cfb = Cfb<Aes256>;
type Aes128Ctr = Ctr128BE<Aes128>;
type Aes192Ctr = Ctr128BE<Aes192>;
type Aes256Ctr = Ctr128BE<Aes256>;
pub struct Cipher {
pub key: Vec<u8>,
pub key_len: usize,
pub iv: Vec<u8>,
pub iv_len: usize,
pub enc: Option<Box<dyn SymmetricCipher + Send + 'static>>,
pub dec: Option<Box<dyn SymmetricCipher + Send + 'static>>,
cipher_method: CipherMethod,
}
#[derive(Clone, Copy, Debug)]
enum CipherMethod {
Aes128Cfb,
Aes192Cfb,
Aes256Cfb,
Aes128Ctr,
Aes192Ctr,
Aes256Ctr,
ChaCha20,
}
impl Cipher {
pub fn new(method: &str, password: &str) -> Cipher {
let (key_len, cipher_method, iv_len) = match method {
"aes-128-cfb" => (16, CipherMethod::Aes128Cfb, 16),
"aes-192-cfb" => (24, CipherMethod::Aes192Cfb, 16),
"aes-256-cfb" => (32, CipherMethod::Aes256Cfb, 16),
"aes-128-ctr" => (16, CipherMethod::Aes128Ctr, 16),
"aes-192-ctr" => (24, CipherMethod::Aes192Ctr, 16),
"aes-256-ctr" => (32, CipherMethod::Aes256Ctr, 16),
"chacha20" => (32, CipherMethod::ChaCha20, 12),
_ => panic!("method not supported"),
};
let key = generate_key(password.as_bytes(), key_len);
Cipher {
key: Vec::from(&key[..]),
key_len,
iv_len,
iv: vec![0u8; iv_len],
enc: None,
dec: None,
cipher_method,
}
}
pub fn init_encrypt(&mut self) {
if self.iv.is_empty() {
let rng = thread_rng();
self.iv = rng.sample_iter(&Standard).take(self.iv_len).collect();
}
self.enc = Some(self.new_cipher(&self.iv));
}
fn new_cipher(&self, iv: &[u8]) -> Box<dyn SymmetricCipher + Send + 'static> {
let key: &[u8] = &self.key;
match self.cipher_method {
CipherMethod::Aes128Cfb => {
Box::new(Aes128Cfb::new_from_slices(key, iv).expect("init cipher error"))
}
CipherMethod::Aes192Cfb => {
Box::new(Aes192Cfb::new_from_slices(key, iv).expect("init cipher error"))
}
CipherMethod::Aes256Cfb => {
Box::new(Aes256Cfb::new_from_slices(key, iv).expect("init cipher error"))
}
CipherMethod::Aes128Ctr => Box::new(Aes128Ctr::new(key.into(), iv.into())),
CipherMethod::Aes192Ctr => Box::new(Aes192Ctr::new(key.into(), iv.into())),
CipherMethod::Aes256Ctr => Box::new(Aes256Ctr::new(key.into(), iv.into())),
CipherMethod::ChaCha20 => Box::new(ChaCha20::new(key.into(), iv.into())),
}
}
pub fn init_decrypt(&mut self) {
self.dec = Some(self.new_cipher(&self.iv));
}
pub fn encrypt(&mut self, input: &mut [u8]) {
if let Some(enc) = &mut self.enc {
enc.encrypt(input);
}
}
pub fn decrypt(&mut self, input: &mut [u8]) {
if let Some(dec) = &mut self.dec {
dec.decrypt(input)
}
}
pub fn reset(&self) -> Cipher {
Cipher {
key: self.key.clone(),
iv: vec![0u8; self.iv_len],
iv_len: self.iv_len,
key_len: self.key_len,
enc: None,
dec: None,
cipher_method: self.cipher_method,
}
}
}