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

use crate::util::generate_key;

use openssl::symm;
use rand::distributions::Standard;
use rand::{thread_rng, Rng};

pub struct Cipher {
    pub cipher: symm::Cipher,
    pub key: Vec<u8>,
    pub iv: Vec<u8>,
    pub iv_len: usize,
    pub enc: Option<symm::Crypter>,
    pub dec: Option<symm::Crypter>,
}

impl Cipher {
    pub fn new(method: &str, password: &str) -> Cipher {
        let cipher = match method {
            "aes-256-cfb" => symm::Cipher::aes_256_cfb128(),
            "aes-128-cfb" => symm::Cipher::aes_128_cfb128(),
            "aes-128-cfb1" => symm::Cipher::aes_128_cfb1(),
            "aes-128-cfb8" => symm::Cipher::aes_128_cfb8(),
            "rc4-md5" => symm::Cipher::rc4(),
            _ => panic!("method not supported"),
        };

        let key = generate_key(password.as_bytes(), cipher.key_len());
        Cipher {
            cipher,
            key: Vec::from(&key[..]),
            iv: vec![],
            iv_len: cipher.iv_len().unwrap_or_default(),
            enc: None,
            dec: None,
        }
    }

    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(
            symm::Crypter::new(self.cipher, symm::Mode::Encrypt, &self.key, Some(&self.iv))
                .expect("init enc error"),
        );
    }

    pub fn init_decrypt(&mut self, iv: &[u8]) {
        self.dec = Some(
            symm::Crypter::new(self.cipher, symm::Mode::Decrypt, &self.key, Some(iv))
                .expect("init enc error"),
        );
    }

    pub fn encrypt(&mut self, input: &[u8]) -> Option<Vec<u8>> {
        let reserve_len = input.len() + self.cipher.block_size();
        let mut out = vec![0u8; reserve_len];
        if let Some(ref mut enc) = self.enc {
            if enc.update(input, &mut out).is_ok() {
                return Some(out);
            }
        }
        None
    }

    pub fn decrypt(&mut self, input: &[u8]) -> Option<Vec<u8>> {
        let reserve_len = input.len() + self.cipher.block_size();
        let mut out = vec![0u8; reserve_len];

        if let Some(ref mut dec) = self.dec {
            if dec.update(input, &mut out).is_ok() {
                return Some(out);
            }
        }
        None
    }

    pub fn reset(&self) -> Cipher {
        Cipher {
            cipher: self.cipher,
            key: self.key.clone(),
            iv: vec![],
            iv_len: self.iv_len,
            enc: None,
            dec: None,
        }
    }
}