1
use std::{
2
    ops::{Mul, Rem, Shr},
3
    sync::LazyLock,
4
};
5

            
6
use cbc::cipher::{
7
    BlockDecryptMut, BlockEncryptMut, BlockSizeUser, IvSizeUser, KeyIvInit, KeySizeUser,
8
    block_padding::{NoPadding, Pkcs7},
9
};
10
use hkdf::Hkdf;
11
use md5::Md5;
12
use num::{FromPrimitive, Integer, One, Zero};
13
use num_bigint_dig::BigUint;
14
use pbkdf2::hmac::{
15
    Mac,
16
    digest::{Digest, FixedOutput, Output, OutputSizeUser},
17
};
18
use sha2::Sha256;
19
use subtle::ConstantTimeEq;
20
use zeroize::{Zeroize, Zeroizing};
21

            
22
use crate::{Key, file};
23

            
24
type EncAlg = cbc::Encryptor<aes::Aes128>;
25
type DecAlg = cbc::Decryptor<aes::Aes128>;
26
type MacAlg = pbkdf2::hmac::Hmac<sha2::Sha256>;
27

            
28
34
pub fn encrypt(
29
    data: impl AsRef<[u8]>,
30
    key: &Key,
31
    iv: impl AsRef<[u8]>,
32
) -> Result<Vec<u8>, super::Error> {
33
68
    let mut blob = vec![0; data.as_ref().len() + EncAlg::block_size()];
34

            
35
    // Unwrapping since adding `CIPHER_BLOCK_SIZE` to array is enough space for
36
    // PKCS7
37
137
    let encrypted_len = EncAlg::new_from_slices(key.as_ref(), iv.as_ref())
38
        .expect("Invalid key length")
39
68
        .encrypt_padded_b2b_mut::<Pkcs7>(data.as_ref(), &mut blob)?
40
        .len();
41

            
42
35
    blob.truncate(encrypted_len);
43

            
44
36
    Ok(blob)
45
}
46

            
47
18
pub fn decrypt(
48
    blob: impl AsRef<[u8]>,
49
    key: &Key,
50
    iv: impl AsRef<[u8]>,
51
) -> Result<Zeroizing<Vec<u8>>, super::Error> {
52
39
    let mut data = blob.as_ref().to_vec();
53

            
54
99
    Ok(DecAlg::new_from_slices(key.as_ref(), iv.as_ref())
55
21
        .expect("Invalid key length")
56
37
        .decrypt_padded_mut::<Pkcs7>(&mut data)?
57
20
        .to_vec()
58
21
        .into())
59
}
60

            
61
4
pub(crate) fn decrypt_no_padding(
62
    blob: impl AsRef<[u8]>,
63
    key: &Key,
64
    iv: impl AsRef<[u8]>,
65
) -> Result<Zeroizing<Vec<u8>>, super::Error> {
66
8
    let mut data = blob.as_ref().to_vec();
67

            
68
20
    Ok(DecAlg::new_from_slices(key.as_ref(), iv.as_ref())
69
4
        .expect("Invalid key length")
70
8
        .decrypt_padded_mut::<NoPadding>(&mut data)?
71
4
        .to_vec()
72
4
        .into())
73
}
74

            
75
9
pub(crate) fn iv_len() -> usize {
76
9
    DecAlg::iv_size()
77
}
78

            
79
12
pub(crate) fn generate_private_key() -> Result<Zeroizing<Vec<u8>>, super::Error> {
80
10
    let mut key = vec![0u8; EncAlg::key_size()];
81
22
    getrandom::fill(&mut key)?;
82
12
    Ok(Zeroizing::new(key))
83
}
84

            
85
11
pub(crate) fn generate_public_key(private_key: impl AsRef<[u8]>) -> Result<Vec<u8>, super::Error> {
86
22
    let private_key_uint = BigUint::from_bytes_be(private_key.as_ref());
87
16
    static DH_GENERATOR: LazyLock<BigUint> = LazyLock::new(|| BigUint::from_u64(0x2).unwrap());
88
22
    let public_key_uint = powm(&DH_GENERATOR, private_key_uint);
89

            
90
17
    Ok(public_key_uint.to_bytes_be())
91
}
92

            
93
10
pub(crate) fn generate_aes_key(
94
    private_key: impl AsRef<[u8]>,
95
    server_public_key: impl AsRef<[u8]>,
96
) -> Result<Zeroizing<Vec<u8>>, super::Error> {
97
18
    let server_public_key_uint = BigUint::from_bytes_be(server_public_key.as_ref());
98
18
    let private_key_uint = BigUint::from_bytes_be(private_key.as_ref());
99
8
    let common_secret = powm(&server_public_key_uint, private_key_uint);
100

            
101
9
    let mut common_secret_bytes = common_secret.to_bytes_be();
102
17
    let mut common_secret_padded = vec![0; 128 - common_secret_bytes.len()];
103
    // inefficient, but ok for now
104
10
    common_secret_padded.append(&mut common_secret_bytes);
105

            
106
    // hkdf
107
    // input_keying_material
108
8
    let ikm = common_secret_padded;
109
10
    let salt = None;
110
    let info = [];
111

            
112
    // output keying material
113
18
    let mut okm = Zeroizing::new(vec![0; 16]);
114

            
115
18
    let (_, hk) = Hkdf::<Sha256>::extract(salt, &ikm);
116
8
    hk.expand(&info, okm.as_mut())
117
        .expect("hkdf expand should never fail");
118

            
119
9
    Ok(okm)
120
}
121

            
122
17
pub fn generate_iv() -> Result<Vec<u8>, super::Error> {
123
16
    let mut iv = vec![0u8; EncAlg::iv_size()];
124
31
    getrandom::fill(&mut iv)?;
125
12
    Ok(iv)
126
}
127

            
128
9
pub(crate) fn mac_len() -> usize {
129
9
    MacAlg::output_size()
130
}
131

            
132
26
pub(crate) fn compute_mac(data: impl AsRef<[u8]>, key: &Key) -> Result<crate::Mac, super::Error> {
133
49
    let mut mac = MacAlg::new_from_slice(key.as_ref()).unwrap();
134
25
    mac.update(data.as_ref());
135
23
    Ok(crate::Mac::new(mac.finalize().into_bytes().to_vec()))
136
}
137

            
138
18
pub(crate) fn verify_mac(
139
    data: impl AsRef<[u8]>,
140
    key: &Key,
141
    expected_mac: impl AsRef<[u8]>,
142
) -> Result<bool, super::Error> {
143
39
    let mut mac = MacAlg::new_from_slice(key.as_ref()).unwrap();
144
20
    mac.update(data.as_ref());
145
18
    Ok(mac.verify_slice(expected_mac.as_ref()).is_ok())
146
}
147

            
148
4
pub(crate) fn verify_checksum_md5(digest: impl AsRef<[u8]>, content: impl AsRef<[u8]>) -> bool {
149
4
    let mut hasher = Md5::new();
150
4
    hasher.update(content.as_ref());
151
4
    hasher.finalize_fixed().ct_eq(digest.as_ref()).into()
152
}
153

            
154
18
pub(crate) fn derive_key(
155
    secret: impl AsRef<[u8]>,
156
    key_strength: Result<(), file::WeakKeyError>,
157
    salt: impl AsRef<[u8]>,
158
    iteration_count: usize,
159
) -> Result<Key, super::Error> {
160
34
    let mut key = Key::new_with_strength(vec![0; EncAlg::block_size()], key_strength);
161

            
162
    pbkdf2::pbkdf2::<pbkdf2::hmac::Hmac<sha2::Sha256>>(
163
17
        secret.as_ref(),
164
18
        salt.as_ref(),
165
17
        iteration_count.try_into().unwrap(),
166
18
        key.as_mut(),
167
    )
168
    .expect("HMAC can be initialized with any key length");
169

            
170
17
    Ok(key)
171
}
172

            
173
5
pub(crate) fn legacy_derive_key_and_iv(
174
    secret: impl AsRef<[u8]>,
175
    key_strength: Result<(), file::WeakKeyError>,
176
    salt: impl AsRef<[u8]>,
177
    iteration_count: usize,
178
) -> Result<(Key, Vec<u8>), super::Error> {
179
10
    let mut buffer = vec![0; EncAlg::key_size() + EncAlg::iv_size()];
180
5
    let mut hasher = Sha256::new();
181
5
    let mut digest_buffer = vec![0; <Sha256 as Digest>::output_size()];
182
    #[allow(deprecated)]
183
10
    let digest = Output::<Sha256>::from_mut_slice(digest_buffer.as_mut_slice());
184

            
185
5
    let mut pos = 0usize;
186

            
187
    loop {
188
5
        hasher.update(secret.as_ref());
189
5
        hasher.update(salt.as_ref());
190
5
        hasher.finalize_into_reset(digest);
191

            
192
5
        for _ in 1..iteration_count {
193
            // We can't pass an instance, the borrow checker
194
            // would complain about digest being dropped at the end of
195
            // for block
196
            #[allow(clippy::needless_borrows_for_generic_args)]
197
5
            hasher.update(&digest);
198
5
            hasher.finalize_into_reset(digest);
199
        }
200

            
201
5
        let to_read = usize::min(digest.len(), buffer.len() - pos);
202
5
        buffer[pos..].copy_from_slice(&digest[..to_read]);
203
5
        pos += to_read;
204

            
205
10
        if pos == buffer.len() {
206
            break;
207
        }
208

            
209
        // We can't pass an instance, the borrow checker
210
        // would complain about digest being dropped at the end of
211
        // loop block
212
        #[allow(clippy::needless_borrows_for_generic_args)]
213
        hasher.update(&digest);
214
    }
215

            
216
10
    let iv = buffer.split_off(EncAlg::key_size());
217
10
    Ok((Key::new_with_strength(buffer, key_strength), iv))
218
}
219

            
220
/// from https://github.com/plietar/librespot/blob/master/core/src/util/mod.rs#L53
221
11
fn powm(base: &BigUint, mut exp: BigUint) -> BigUint {
222
    // for key exchange
223
8
    static DH_PRIME: LazyLock<BigUint> = LazyLock::new(|| {
224
8
        BigUint::from_bytes_be(&[
225
            0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xC9, 0x0F, 0xDA, 0xA2, 0x21, 0x68,
226
            0xC2, 0x34, 0xC4, 0xC6, 0x62, 0x8B, 0x80, 0xDC, 0x1C, 0xD1, 0x29, 0x02, 0x4E, 0x08,
227
            0x8A, 0x67, 0xCC, 0x74, 0x02, 0x0B, 0xBE, 0xA6, 0x3B, 0x13, 0x9B, 0x22, 0x51, 0x4A,
228
            0x08, 0x79, 0x8E, 0x34, 0x04, 0xDD, 0xEF, 0x95, 0x19, 0xB3, 0xCD, 0x3A, 0x43, 0x1B,
229
            0x30, 0x2B, 0x0A, 0x6D, 0xF2, 0x5F, 0x14, 0x37, 0x4F, 0xE1, 0x35, 0x6D, 0x6D, 0x51,
230
            0xC2, 0x45, 0xE4, 0x85, 0xB5, 0x76, 0x62, 0x5E, 0x7E, 0xC6, 0xF4, 0x4C, 0x42, 0xE9,
231
            0xA6, 0x37, 0xED, 0x6B, 0x0B, 0xFF, 0x5C, 0xB6, 0xF4, 0x06, 0xB7, 0xED, 0xEE, 0x38,
232
            0x6B, 0xFB, 0x5A, 0x89, 0x9F, 0xA5, 0xAE, 0x9F, 0x24, 0x11, 0x7C, 0x4B, 0x1F, 0xE6,
233
            0x49, 0x28, 0x66, 0x51, 0xEC, 0xE6, 0x53, 0x81, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
234
            0xFF, 0xFF,
235
        ])
236
    });
237

            
238
11
    let mut base = base.clone();
239
11
    let mut result: BigUint = One::one();
240

            
241
31
    while !exp.is_zero() {
242
31
        if exp.is_odd() {
243
9
            result = result.mul(&base).rem(&*DH_PRIME);
244
        }
245
18
        exp = exp.shr(1);
246
9
        base = (&base).mul(&base).rem(&*DH_PRIME);
247
    }
248
8
    exp.zeroize();
249

            
250
10
    result
251
}