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
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

use crate::util::generate_key;

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};

pub trait SymmetricCipher {
    /// Encrypt data in place.
    fn encrypt(&mut self, data: &mut [u8]);

    /// Decrypt data in place.
    fn decrypt(&mut self, data: &mut [u8]);
}

impl<C: BlockCipher + BlockEncrypt> SymmetricCipher for Cfb<C> {
    /// Encrypt data in place.
    fn encrypt(&mut self, data: &mut [u8]) {
        AsyncStreamCipher::encrypt(self, data)
    }

    /// Decrypt data in place.
    fn decrypt(&mut self, data: &mut [u8]) {
        AsyncStreamCipher::decrypt(self, data)
    }
}

impl<B: BlockEncrypt + BlockCipher<BlockSize = U16>> SymmetricCipher for Ctr128BE<B> {
    /// Encrypt data in place.
    fn encrypt(&mut self, data: &mut [u8]) {
        StreamCipher::apply_keystream(self, data)
    }

    /// Decrypt data in place.
    fn decrypt(&mut self, data: &mut [u8]) {
        StreamCipher::apply_keystream(self, data)
    }
}

impl SymmetricCipher for ChaCha20 {
    /// Encrypt data in place.
    fn encrypt(&mut self, data: &mut [u8]) {
        StreamCipher::apply_keystream(self, data)
    }

    /// Decrypt data in place.
    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, iv: &[u8]) {
        self.dec = Some(self.new_cipher(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![],
            key_len: self.key_len,
            iv_len: self.iv_len,
            enc: None,
            dec: None,
            cipher_method: self.cipher_method,
        }
    }
}