1use aes::cipher::{KeyIvInit, StreamCipher};
2use anyhow::Result;
3use hmac::{Hmac, Mac};
4use sha2::Sha256;
5
6type HmacSha256 = Hmac<Sha256>;
7
8fn hmac_ad(key: &[u8], data: &[u8], ad: &[u8]) -> [u8; 32] {
10 let mut mac = <HmacSha256 as hmac::Mac>::new_from_slice(key).expect("HMAC key length error");
11 mac.update(ad);
12 mac.update(data);
13 mac.update(&(ad.len() as u64).to_le_bytes());
14
15 mac.finalize().into_bytes().into()
16}
17
18fn hmac_sha256(key: &[u8], data: &[u8]) -> [u8; 32] {
21 let mut mac = <HmacSha256 as hmac::Mac>::new_from_slice(key).expect("HMAC key length error");
22 mac.update(data);
23
24 mac.finalize().into_bytes().into()
25}
26
27pub fn verify_padding(key: &[u8], ciphertext: &[u8], expected_hmac: &[u8]) -> Result<Vec<u8>> {
28 let calculated_hmac = hmac_sha256(key, ciphertext);
29 if calculated_hmac != expected_hmac {
30 anyhow::bail!("Integrity check failed: HMAC mismatch");
31 }
32
33 Ok(ciphertext.to_vec())
34}
35
36pub fn aes_aead_decrypt(
38 key_80b: &[u8],
39 ciphertext: &[u8],
40 ad: &[u8],
41 expected_hmac: &[u8],
42) -> Result<Vec<u8>> {
43 let hmac_key = &key_80b[0..32];
44 let aes_key = &key_80b[32..64];
45 let aes_iv = &key_80b[64..80];
46
47 let calculated_hmac = hmac_ad(hmac_key, ciphertext, ad);
48 if calculated_hmac != expected_hmac {
49 anyhow::bail!("Integrity check failed: HMAC_AD mismatch");
50 }
51
52 let mut output = ciphertext.to_vec();
53 let mut cipher = ctr::Ctr128BE::<aes::Aes256>::new(aes_key.into(), aes_iv.into());
54 cipher.apply_keystream(&mut output);
55
56 Ok(output)
57}
58
59pub fn derive_main_key(salt: &[u8], external_key: &[u8], profile_id: &[u8]) -> Result<[u8; 32]> {
61 let mut amk_info = b"AEA_AMK".to_vec();
62 amk_info.extend_from_slice(profile_id);
63 amk_info.push(0);
64
65 let hk_amk = hkdf::Hkdf::<Sha256>::new(Some(salt), external_key);
66 let mut amk = [0u8; 32];
67 hk_amk
68 .expand(&amk_info, &mut amk)
69 .map_err(|_| anyhow::anyhow!("AMK expand fail"))?;
70
71 Ok(amk)
72}
73
74pub fn derive_cluster_key(amk: &[u8; 32], cluster_index: u32) -> anyhow::Result<[u8; 32]> {
76 let mut ck_info = b"AEA_CK".to_vec();
77 ck_info.extend_from_slice(&cluster_index.to_le_bytes());
78 let hk_ck = hkdf::Hkdf::<sha2::Sha256>::new(None, amk);
79 let mut ck = [0u8; 32];
80 hk_ck
81 .expand(&ck_info, &mut ck)
82 .map_err(|_| anyhow::anyhow!("CK expand fail"))?;
83
84 Ok(ck)
85}
86
87pub fn derive_cluster_header_encryption_key(ck: &[u8; 32]) -> [u8; 80] {
89 let chek_info = b"AEA_CHEK".to_vec();
90 let hk_chek = hkdf::Hkdf::<sha2::Sha256>::new(None, ck);
91 let mut chek = [0u8; 80];
92 hk_chek
93 .expand(&chek_info, &mut chek)
94 .expect("CHEK expand fail");
95
96 chek
97}
98
99pub fn derive_segment_key(ck: &[u8; 32], segment_index: u32) -> [u8; 80] {
101 let mut sk_info = b"AEA_SK".to_vec();
102 sk_info.extend_from_slice(&segment_index.to_le_bytes());
103
104 let hk_sk = hkdf::Hkdf::<sha2::Sha256>::new(None, ck);
105 let mut sk = [0u8; 80];
106 hk_sk.expand(&sk_info, &mut sk).expect("SK expand fail");
107
108 sk
109}
110
111pub fn derive_padding_authentication_key(amk: &[u8; 32]) -> [u8; 32] {
113 let pak_info = b"AEA_PAK".to_vec();
114 let hk_pak = hkdf::Hkdf::<sha2::Sha256>::new(None, amk);
115 let mut pak = [0u8; 32];
116 hk_pak.expand(&pak_info, &mut pak).expect("PAK expand fail");
117
118 pak
119}