1use crate::{prf, Error};
6use anyhow::ensure;
7use ieee80211::MacAddr;
8use mundane::hash::Sha256;
9use std::cmp::{max, min};
10use wlan_common::ie::rsn::akm::{self, Akm};
11use wlan_common::ie::rsn::cipher::Cipher;
12use wlan_sae::hmac_utils;
13
14#[derive(Debug, Clone, PartialEq)]
18pub struct Ptk {
19 pub ptk: Vec<u8>,
20 kck_len: usize,
21 kek_len: usize,
22 tk_len: usize,
23 pub cipher: Cipher,
24 }
26
27impl Ptk {
28 pub fn from_ptk(ptk: Vec<u8>, akm: &Akm, cipher: Cipher) -> Result<Self, anyhow::Error> {
29 let kck_len = akm.kck_bytes().ok_or(Error::PtkHierarchyUnsupportedAkmError)? as usize;
30 let kek_len = akm.kek_bytes().ok_or(Error::PtkHierarchyUnsupportedAkmError)? as usize;
31 let tk_len: usize =
32 cipher.tk_bytes().ok_or(Error::PtkHierarchyUnsupportedCipherError)?.into();
33 ensure!(kck_len + kek_len + tk_len == ptk.len(), "invalid ptk length");
34 Ok(Ptk { ptk, kck_len, kek_len, tk_len, cipher })
35 }
36
37 #[allow(deprecated)]
39 pub fn new(
41 pmk: &[u8],
42 aa: &MacAddr,
43 spa: &MacAddr,
44 anonce: &[u8],
45 snonce: &[u8],
46 akm: &Akm,
47 cipher: Cipher,
48 ) -> Result<Ptk, anyhow::Error> {
49 ensure!(anonce.len() == 32 && snonce.len() == 32, Error::InvalidNonceSize(anonce.len()));
50
51 let pmk_len = akm
52 .pmk_bits()
53 .map(|bits| (bits / 8) as usize)
54 .ok_or(Error::PtkHierarchyUnsupportedAkmError)?;
55 ensure!(pmk.len() == pmk_len, Error::PtkHierarchyInvalidPmkError);
56
57 let kck_bits = akm.kck_bits().ok_or(Error::PtkHierarchyUnsupportedAkmError)?;
58 let kek_bits = akm.kek_bits().ok_or(Error::PtkHierarchyUnsupportedAkmError)?;
59 let tk_bits = cipher.tk_bits().ok_or(Error::PtkHierarchyUnsupportedCipherError)?;
60 let prf_bits = kck_bits + kek_bits + tk_bits;
61
62 let mut data: [u8; 76] = [0; 76];
64 data[0..6].copy_from_slice(&min(aa.as_slice(), spa.as_slice())[..]);
65 data[6..12].copy_from_slice(&max(aa.as_slice(), spa.as_slice())[..]);
66 data[12..44].copy_from_slice(&min(anonce, snonce)[..]);
67 data[44..].copy_from_slice(&max(anonce, snonce)[..]);
68
69 let ptk_bytes = match akm.suite_type {
72 akm::SAE => hmac_utils::kdf_hash_length::<Sha256>(
74 pmk,
75 "Pairwise key expansion",
76 &data,
77 prf_bits as usize,
78 ),
79 _ => prf::prf(pmk, "Pairwise key expansion", &data, prf_bits as usize)?,
80 };
81 let ptk = Ptk {
82 ptk: ptk_bytes,
83 kck_len: (kck_bits / 8) as usize,
84 kek_len: (kek_bits / 8) as usize,
85 tk_len: (tk_bits / 8) as usize,
86 cipher,
87 };
88 Ok(ptk)
89 }
90
91 pub fn kck(&self) -> &[u8] {
92 &self.ptk[0..self.kck_len]
93 }
94
95 pub fn kek(&self) -> &[u8] {
96 let start = self.kck_len;
97 &self.ptk[start..start + self.kek_len]
98 }
99
100 pub fn tk(&self) -> &[u8] {
101 let start = self.kck_len + self.kek_len;
102 &self.ptk[start..start + self.tk_len]
103 }
104}
105
106#[cfg(test)]
107mod tests {
108 use super::*;
109 use hex::FromHex;
110 use wlan_common::ie::rsn::akm::PSK;
111 use wlan_common::ie::rsn::cipher::{CCMP_128, TKIP};
112
113 struct TestData {
114 pmk: Vec<u8>,
115 aa: MacAddr,
116 spa: MacAddr,
117 anonce: [u8; 32],
118 snonce: [u8; 32],
119 }
120
121 fn ieee_test_data() -> TestData {
123 let pmk = Vec::from_hex("0dc0d6eb90555ed6419756b9a15ec3e3209b63df707dd508d14581f8982721af")
124 .unwrap();
125 let aa = MacAddr::from(<[u8; 6]>::from_hex("a0a1a1a3a4a5").unwrap());
126 let spa = MacAddr::from(<[u8; 6]>::from_hex("b0b1b2b3b4b5").unwrap());
127 let anonce = <[u8; 32]>::from_hex(
128 "e0e1e2e3e4e5e6e7e8e9f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff000102030405",
129 )
130 .unwrap();
131 let snonce = <[u8; 32]>::from_hex(
132 "c0c1c2c3c4c5c6c7c8c9d0d1d2d3d4d5d6d7d8d9dadbdcdddedfe0e1e2e3e4e5",
133 )
134 .unwrap();
135 TestData { pmk, aa, spa, anonce, snonce }
136 }
137
138 fn new_ptk(data: &TestData, akm_suite: u8, cipher_suite: u8) -> Result<Ptk, anyhow::Error> {
139 let akm = Akm::new_dot11(akm_suite);
140 let cipher = Cipher::new_dot11(cipher_suite);
141 Ptk::new(&data.pmk[..], &data.aa, &data.spa, &data.anonce, &data.snonce, &akm, cipher)
142 }
143
144 #[test]
146 fn test_pairwise_key_hierarchy_ccmp() {
147 let data = ieee_test_data();
148 let ptk_result = new_ptk(&data, PSK, CCMP_128);
149 assert_eq!(ptk_result.is_ok(), true);
150
151 let expected_kck = Vec::from_hex("379f9852d0199236b94e407ce4c00ec8").unwrap();
153 let expected_kek = Vec::from_hex("47c9edc01c2c6e5b4910caddfb3e51a7").unwrap();
154 let expected_tk = Vec::from_hex("b2360c79e9710fdd58bea93deaf06599").unwrap();
155 let ptk = ptk_result.unwrap();
156 assert_eq!(ptk.kck(), &expected_kck[..]);
157 assert_eq!(ptk.kek(), &expected_kek[..]);
158 assert_eq!(ptk.tk(), &expected_tk[..]);
159 }
160
161 #[test]
163 fn test_pairwise_key_hierarchy_tkip() {
164 let data = ieee_test_data();
165 let ptk_result = new_ptk(&data, PSK, TKIP);
166 assert_eq!(ptk_result.is_ok(), true);
167
168 let expected_kck = Vec::from_hex("379f9852d0199236b94e407ce4c00ec8").unwrap();
170 let expected_kek = Vec::from_hex("47c9edc01c2c6e5b4910caddfb3e51a7").unwrap();
171 let expected_tk =
172 Vec::from_hex("b2360c79e9710fdd58bea93deaf06599db980afbc29c152855740a6ce5ae3827")
173 .unwrap();
174 let ptk = ptk_result.unwrap();
175 assert_eq!(ptk.kck(), &expected_kck[..]);
176 assert_eq!(ptk.kek(), &expected_kek[..]);
177 assert_eq!(ptk.tk(), &expected_tk[..]);
178 }
179
180 #[test]
181 fn test_pairwise_key_hierarchy_invalid_pmk() {
182 let mut data = ieee_test_data();
183 data.pmk.remove(0); let ptk_result = new_ptk(&data, PSK, CCMP_128);
185 assert_eq!(ptk_result.is_err(), true);
186 }
187
188 #[test]
189 fn test_pairwise_key_hierarchy_unsupported_akm() {
190 let data = ieee_test_data();
191 let ptk_result = new_ptk(&data, 200, CCMP_128);
192 assert_eq!(ptk_result.is_err(), true);
193 }
194
195 #[test]
196 fn test_pairwise_key_hierarchy_unsupported_cipher() {
197 let data = ieee_test_data();
198 let ptk_result = new_ptk(&data, PSK, 200);
199 assert_eq!(ptk_result.is_err(), true);
200 }
201}