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

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

            
19
use crate::{Key, file};
20

            
21
type EncAlg = cbc::Encryptor<aes::Aes128>;
22
type DecAlg = cbc::Decryptor<aes::Aes128>;
23
type MacAlg = hmac::Hmac<sha2::Sha256>;
24

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

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

            
39
6
    blob.truncate(encrypted_len);
40

            
41
6
    Ok(blob)
42
}
43

            
44
5
pub fn decrypt(
45
    blob: impl AsRef<[u8]>,
46
    key: &Key,
47
    iv: impl AsRef<[u8]>,
48
) -> Result<Zeroizing<Vec<u8>>, super::Error> {
49
10
    let mut data = blob.as_ref().to_vec();
50

            
51
25
    Ok(DecAlg::new_from_slices(key.as_ref(), iv.as_ref())
52
5
        .expect("Invalid key length")
53
10
        .decrypt_padded_mut::<Pkcs7>(&mut data)?
54
5
        .to_vec()
55
5
        .into())
56
}
57

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

            
65
10
    Ok(DecAlg::new_from_slices(key.as_ref(), iv.as_ref())
66
2
        .expect("Invalid key length")
67
4
        .decrypt_padded_mut::<NoPadding>(&mut data)?
68
2
        .to_vec()
69
2
        .into())
70
}
71

            
72
2
pub(crate) fn iv_len() -> usize {
73
2
    DecAlg::iv_size()
74
}
75

            
76
1
pub(crate) fn generate_private_key() -> Result<Zeroizing<Vec<u8>>, super::Error> {
77
1
    let mut key = vec![0u8; EncAlg::key_size()];
78
2
    getrandom::fill(&mut key)?;
79
1
    Ok(Zeroizing::new(key))
80
}
81

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

            
87
4
    Ok(public_key_uint.to_bytes_be())
88
}
89

            
90
2
pub(crate) fn generate_aes_key(
91
    private_key: impl AsRef<[u8]>,
92
    server_public_key: impl AsRef<[u8]>,
93
) -> Result<Zeroizing<Vec<u8>>, super::Error> {
94
4
    let server_public_key_uint = BigUint::from_bytes_be(server_public_key.as_ref());
95
4
    let private_key_uint = BigUint::from_bytes_be(private_key.as_ref());
96
2
    let common_secret = powm(&server_public_key_uint, private_key_uint);
97

            
98
2
    let mut common_secret_bytes = common_secret.to_bytes_be();
99
4
    let mut common_secret_padded = vec![0; 128 - common_secret_bytes.len()];
100
    // inefficient, but ok for now
101
2
    common_secret_padded.append(&mut common_secret_bytes);
102

            
103
    // hkdf
104
    // input_keying_material
105
2
    let ikm = common_secret_padded;
106
2
    let salt = None;
107
    let info = [];
108

            
109
    // output keying material
110
4
    let mut okm = Zeroizing::new(vec![0; 16]);
111

            
112
4
    let (_, hk) = Hkdf::<Sha256>::extract(salt, &ikm);
113
2
    hk.expand(&info, okm.as_mut())
114
        .expect("hkdf expand should never fail");
115

            
116
2
    Ok(okm)
117
}
118

            
119
2
pub fn generate_iv() -> Result<Vec<u8>, super::Error> {
120
2
    let mut iv = vec![0u8; EncAlg::iv_size()];
121
4
    getrandom::fill(&mut iv)?;
122
2
    Ok(iv)
123
}
124

            
125
2
pub(crate) fn mac_len() -> usize {
126
2
    MacAlg::output_size()
127
}
128

            
129
4
pub(crate) fn compute_mac(data: impl AsRef<[u8]>, key: &Key) -> Result<crate::Mac, super::Error> {
130
8
    let mut mac = MacAlg::new_from_slice(key.as_ref()).unwrap();
131
4
    mac.update(data.as_ref());
132
6
    Ok(crate::Mac::new(mac.finalize().into_bytes().to_vec()))
133
}
134

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

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

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

            
159
    pbkdf2::pbkdf2::<hmac::Hmac<sha2::Sha256>>(
160
3
        secret.as_ref(),
161
3
        salt.as_ref(),
162
3
        iteration_count.try_into().unwrap(),
163
3
        key.as_mut(),
164
    )
165
    .expect("HMAC can be initialized with any key length");
166

            
167
2
    Ok(key)
168
}
169

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

            
182
3
    let mut pos = 0usize;
183

            
184
    loop {
185
3
        hasher.update(secret.as_ref());
186
3
        hasher.update(salt.as_ref());
187
3
        hasher.finalize_into_reset(digest);
188

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

            
198
3
        let to_read = usize::min(digest.len(), buffer.len() - pos);
199
3
        buffer[pos..].copy_from_slice(&digest[..to_read]);
200
3
        pos += to_read;
201

            
202
6
        if pos == buffer.len() {
203
            break;
204
        }
205

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

            
213
6
    let iv = buffer.split_off(EncAlg::key_size());
214
6
    Ok((Key::new_with_strength(buffer, key_strength), iv))
215
}
216

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

            
235
2
    let mut base = base.clone();
236
2
    let mut result: BigUint = One::one();
237

            
238
6
    while !exp.is_zero() {
239
6
        if exp.is_odd() {
240
2
            result = result.mul(&base).rem(&*DH_PRIME);
241
        }
242
4
        exp = exp.shr(1);
243
2
        base = (&base).mul(&base).rem(&*DH_PRIME);
244
    }
245
2
    exp.zeroize();
246

            
247
2
    result
248
}