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

//! Stateful encryptor object.

use crate::{Cipher, Error, Result};
use cipher::{Block, BlockCipher, BlockCipherEncrypt, KeyIvInit};

#[cfg(feature = "aes-ctr")]
use {
    crate::Ctr128BE,
    cipher::{array::sizes::U16, StreamCipherCore},
};

#[cfg(feature = "tdes")]
use des::TdesEde3;

#[cfg(any(feature = "aes-cbc", feature = "aes-ctr"))]
use aes::{Aes128, Aes192, Aes256};

#[cfg(any(feature = "aes-cbc", feature = "tdes"))]
use cipher::BlockModeEncrypt;

/// Stateful encryptor object for unauthenticated SSH symmetric ciphers.
///
/// Note that this deliberately does not support AEAD modes such as AES-GCM and ChaCha20Poly1305,
/// which are one-shot by design.
pub struct Encryptor {
    /// Inner enum over possible encryption ciphers.
    inner: Inner,
}

/// Inner encryptor enum which is deliberately kept out of the public API.
enum Inner {
    #[cfg(feature = "aes-cbc")]
    Aes128Cbc(cbc::Encryptor<Aes128>),
    #[cfg(feature = "aes-cbc")]
    Aes192Cbc(cbc::Encryptor<Aes192>),
    #[cfg(feature = "aes-cbc")]
    Aes256Cbc(cbc::Encryptor<Aes256>),
    #[cfg(feature = "aes-ctr")]
    Aes128Ctr(Ctr128BE<Aes128>),
    #[cfg(feature = "aes-ctr")]
    Aes192Ctr(Ctr128BE<Aes192>),
    #[cfg(feature = "aes-ctr")]
    Aes256Ctr(Ctr128BE<Aes256>),
    #[cfg(feature = "tdes")]
    TDesCbc(cbc::Encryptor<TdesEde3>),
}

impl Encryptor {
    /// Create a new encryptor object with the given [`Cipher`], key, and IV.
    pub fn new(cipher: Cipher, key: &[u8], iv: &[u8]) -> Result<Self> {
        cipher.check_key_and_iv(key, iv)?;

        let inner = match cipher {
            #[cfg(feature = "aes-cbc")]
            Cipher::Aes128Cbc => cbc::Encryptor::new_from_slices(key, iv).map(Inner::Aes128Cbc),
            #[cfg(feature = "aes-cbc")]
            Cipher::Aes192Cbc => cbc::Encryptor::new_from_slices(key, iv).map(Inner::Aes192Cbc),
            #[cfg(feature = "aes-cbc")]
            Cipher::Aes256Cbc => cbc::Encryptor::new_from_slices(key, iv).map(Inner::Aes256Cbc),
            #[cfg(feature = "aes-ctr")]
            Cipher::Aes128Ctr => Ctr128BE::new_from_slices(key, iv).map(Inner::Aes128Ctr),
            #[cfg(feature = "aes-ctr")]
            Cipher::Aes192Ctr => Ctr128BE::new_from_slices(key, iv).map(Inner::Aes192Ctr),
            #[cfg(feature = "aes-ctr")]
            Cipher::Aes256Ctr => Ctr128BE::new_from_slices(key, iv).map(Inner::Aes256Ctr),
            #[cfg(feature = "tdes")]
            Cipher::TDesCbc => cbc::Encryptor::new_from_slices(key, iv).map(Inner::TDesCbc),
            _ => return Err(cipher.unsupported()),
        }
        .map_err(|_| Error::Length)?;

        Ok(Self { inner })
    }

    /// Get the cipher for this encryptor.
    pub fn cipher(&self) -> Cipher {
        match &self.inner {
            #[cfg(feature = "aes-cbc")]
            Inner::Aes128Cbc(_) => Cipher::Aes128Cbc,
            #[cfg(feature = "aes-cbc")]
            Inner::Aes192Cbc(_) => Cipher::Aes192Cbc,
            #[cfg(feature = "aes-cbc")]
            Inner::Aes256Cbc(_) => Cipher::Aes256Cbc,
            #[cfg(feature = "aes-ctr")]
            Inner::Aes128Ctr(_) => Cipher::Aes128Ctr,
            #[cfg(feature = "aes-ctr")]
            Inner::Aes192Ctr(_) => Cipher::Aes192Ctr,
            #[cfg(feature = "aes-ctr")]
            Inner::Aes256Ctr(_) => Cipher::Aes256Ctr,
            #[cfg(feature = "tdes")]
            Inner::TDesCbc(_) => Cipher::TDesCbc,
        }
    }

    /// Encrypt the given buffer in place.
    ///
    /// Returns [`Error::Length`] in the event that `buffer` is not a multiple of the cipher's
    /// block size.
    pub fn encrypt(&mut self, buffer: &mut [u8]) -> Result<()> {
        match &mut self.inner {
            #[cfg(feature = "aes-cbc")]
            Inner::Aes128Cbc(cipher) => cbc_encrypt(cipher, buffer)?,
            #[cfg(feature = "aes-cbc")]
            Inner::Aes192Cbc(cipher) => cbc_encrypt(cipher, buffer)?,
            #[cfg(feature = "aes-cbc")]
            Inner::Aes256Cbc(cipher) => cbc_encrypt(cipher, buffer)?,
            #[cfg(feature = "aes-ctr")]
            Inner::Aes128Ctr(cipher) => ctr_encrypt(cipher, buffer)?,
            #[cfg(feature = "aes-ctr")]
            Inner::Aes192Ctr(cipher) => ctr_encrypt(cipher, buffer)?,
            #[cfg(feature = "aes-ctr")]
            Inner::Aes256Ctr(cipher) => ctr_encrypt(cipher, buffer)?,
            #[cfg(feature = "tdes")]
            Inner::TDesCbc(cipher) => cbc_encrypt(cipher, buffer)?,
        }

        Ok(())
    }
}

/// CBC mode encryption helper which assumes the input is unpadded and block-aligned.
#[cfg(any(feature = "aes-cbc", feature = "tdes"))]
fn cbc_encrypt<C>(encryptor: &mut cbc::Encryptor<C>, buffer: &mut [u8]) -> Result<()>
where
    C: BlockCipher + BlockCipherEncrypt,
{
    let (blocks, remaining) = Block::<C>::slice_as_chunks_mut(buffer);

    // Ensure input is block-aligned.
    if !remaining.is_empty() {
        return Err(Error::Length);
    }

    encryptor.encrypt_blocks(blocks);
    Ok(())
}

/// CTR mode encryption helper which assumes the input is unpadded and block-aligned.
#[cfg(feature = "aes-ctr")]
pub(crate) fn ctr_encrypt<C>(encryptor: &mut Ctr128BE<C>, buffer: &mut [u8]) -> Result<()>
where
    C: BlockCipher<BlockSize = U16> + BlockCipherEncrypt,
{
    let (blocks, remaining) = Block::<C>::slice_as_chunks_mut(buffer);

    // Ensure input is block-aligned.
    if !remaining.is_empty() {
        return Err(Error::Length);
    }

    encryptor.apply_keystream_blocks(blocks);
    Ok(())
}