⚠️ VeridianOS Kernel Documentation - This is low-level kernel code. All functions are unsafe unless explicitly marked otherwise. no_std

veridian_kernel/drivers/wifi/
wpa.rs

1//! WPA2/WPA3 Authentication Implementation
2//!
3//! Provides EAPOL frame handling, PBKDF2-SHA256 key derivation, PRF-SHA256
4//! for PTK expansion, 4-way handshake state machine, CCMP stub, and WPA3-SAE
5//! stubs for WiFi security.
6
7use alloc::vec::Vec;
8
9use crate::{crypto::hash::sha256, net::MacAddress};
10
11// ============================================================================
12// HMAC-SHA256
13// ============================================================================
14
15/// HMAC-SHA256 implementation (RFC 2104)
16///
17/// Computes HMAC using SHA-256 as the underlying hash function.
18/// All integer math, no floating point.
19fn hmac_sha256(key: &[u8], data: &[u8]) -> [u8; 32] {
20    const BLOCK_SIZE: usize = 64;
21
22    // If key is longer than block size, hash it first
23    let hashed_key;
24    let key_ref = if key.len() > BLOCK_SIZE {
25        hashed_key = sha256(key);
26        hashed_key.as_bytes()
27    } else {
28        key
29    };
30
31    // Pad key to block size
32    let mut ipad = [0x36u8; BLOCK_SIZE];
33    let mut opad = [0x5Cu8; BLOCK_SIZE];
34    for (i, &b) in key_ref.iter().enumerate() {
35        ipad[i] ^= b;
36        opad[i] ^= b;
37    }
38
39    // Inner hash: H(K ^ ipad || data)
40    let mut inner_input = Vec::with_capacity(BLOCK_SIZE + data.len());
41    inner_input.extend_from_slice(&ipad);
42    inner_input.extend_from_slice(data);
43    let inner_hash = sha256(&inner_input);
44
45    // Outer hash: H(K ^ opad || inner_hash)
46    let mut outer_input = Vec::with_capacity(BLOCK_SIZE + 32);
47    outer_input.extend_from_slice(&opad);
48    outer_input.extend_from_slice(inner_hash.as_bytes());
49    let outer_hash = sha256(&outer_input);
50
51    let mut result = [0u8; 32];
52    result.copy_from_slice(outer_hash.as_bytes());
53    result
54}
55
56// ============================================================================
57// EAPOL Frame
58// ============================================================================
59
60/// EAPOL (IEEE 802.1X) frame header
61#[derive(Debug, Clone)]
62pub struct EapolFrame {
63    /// Protocol version (1 for 802.1X-2001, 2 for 802.1X-2004)
64    pub protocol_version: u8,
65    /// Packet type (0=EAP, 1=Start, 2=Logoff, 3=Key)
66    pub packet_type: u8,
67    /// Length of packet body
68    pub packet_body_length: u16,
69    /// Packet body
70    pub body: Vec<u8>,
71}
72
73/// EAPOL packet type: Key (used in 4-way handshake)
74pub const EAPOL_KEY: u8 = 3;
75
76impl EapolFrame {
77    /// Parse EAPOL frame from bytes
78    pub fn from_bytes(data: &[u8]) -> Option<Self> {
79        if data.len() < 4 {
80            return None;
81        }
82
83        let protocol_version = data[0];
84        let packet_type = data[1];
85        let packet_body_length = u16::from_be_bytes([data[2], data[3]]);
86
87        let body_end = 4 + packet_body_length as usize;
88        if data.len() < body_end {
89            return None;
90        }
91
92        Some(Self {
93            protocol_version,
94            packet_type,
95            packet_body_length,
96            body: data[4..body_end].to_vec(),
97        })
98    }
99
100    /// Serialize EAPOL frame to bytes
101    pub fn to_bytes(&self) -> Vec<u8> {
102        let mut buf = Vec::with_capacity(4 + self.body.len());
103        buf.push(self.protocol_version);
104        buf.push(self.packet_type);
105        buf.extend_from_slice(&self.packet_body_length.to_be_bytes());
106        buf.extend_from_slice(&self.body);
107        buf
108    }
109}
110
111// ============================================================================
112// EAPOL Key Frame
113// ============================================================================
114
115/// Key information flags (16-bit field in EAPOL-Key frame)
116#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
117pub struct KeyInfo {
118    /// Key descriptor version (1=HMAC-MD5/RC4, 2=HMAC-SHA1/AES, 3=AES-128-CMAC)
119    pub descriptor_version: u8,
120    /// Pairwise key (1) vs Group key (0)
121    pub pairwise: bool,
122    /// Install key flag
123    pub install: bool,
124    /// Key ACK flag (set by authenticator)
125    pub ack: bool,
126    /// MIC flag (MIC is present and valid)
127    pub mic: bool,
128    /// Secure flag (pairwise key installed)
129    pub secure: bool,
130    /// Error flag
131    pub error: bool,
132    /// Request flag
133    pub request: bool,
134    /// Encrypted Key Data flag
135    pub encrypted: bool,
136}
137
138impl KeyInfo {
139    /// Parse from 16-bit value (big-endian)
140    pub fn from_u16(val: u16) -> Self {
141        Self {
142            descriptor_version: (val & 0x0007) as u8,
143            pairwise: (val & 0x0008) != 0,
144            install: (val & 0x0040) != 0,
145            ack: (val & 0x0080) != 0,
146            mic: (val & 0x0100) != 0,
147            secure: (val & 0x0200) != 0,
148            error: (val & 0x0400) != 0,
149            request: (val & 0x0800) != 0,
150            encrypted: (val & 0x1000) != 0,
151        }
152    }
153
154    /// Serialize to 16-bit value (big-endian)
155    pub fn to_u16(&self) -> u16 {
156        (self.descriptor_version as u16 & 0x07)
157            | if self.pairwise { 0x0008 } else { 0 }
158            | if self.install { 0x0040 } else { 0 }
159            | if self.ack { 0x0080 } else { 0 }
160            | if self.mic { 0x0100 } else { 0 }
161            | if self.secure { 0x0200 } else { 0 }
162            | if self.error { 0x0400 } else { 0 }
163            | if self.request { 0x0800 } else { 0 }
164            | if self.encrypted { 0x1000 } else { 0 }
165    }
166}
167
168/// EAPOL-Key frame body (within EAPOL frame)
169#[derive(Debug, Clone)]
170pub struct EapolKeyFrame {
171    /// Descriptor type (2 = RSN/WPA2)
172    pub descriptor_type: u8,
173    /// Key information flags
174    pub key_info: KeyInfo,
175    /// Key length (16 for CCMP)
176    pub key_length: u16,
177    /// Replay counter (monotonically increasing)
178    pub replay_counter: u64,
179    /// Key nonce (32 bytes - ANonce from AP, SNonce from STA)
180    pub key_nonce: [u8; 32],
181    /// Key IV (16 bytes, typically zero for WPA2)
182    pub key_iv: [u8; 16],
183    /// Key MIC (16 bytes, HMAC over frame)
184    pub key_mic: [u8; 16],
185    /// Key data length
186    pub key_data_length: u16,
187    /// Key data (encrypted GTK, PMKID, etc.)
188    pub key_data: Vec<u8>,
189}
190
191impl EapolKeyFrame {
192    /// Minimum key frame body size (without key data):
193    /// 1 + 2 + 2 + 8 + 32 + 16 + 8(RSC) + 8(reserved) + 16 + 2 = 95 bytes
194    pub const MIN_SIZE: usize = 95;
195
196    /// Parse EAPOL-Key frame from body bytes (after EAPOL header)
197    pub fn from_bytes(data: &[u8]) -> Option<Self> {
198        if data.len() < Self::MIN_SIZE {
199            return None;
200        }
201
202        let descriptor_type = data[0];
203        let key_info = KeyInfo::from_u16(u16::from_be_bytes([data[1], data[2]]));
204        let key_length = u16::from_be_bytes([data[3], data[4]]);
205        let replay_counter = u64::from_be_bytes([
206            data[5], data[6], data[7], data[8], data[9], data[10], data[11], data[12],
207        ]);
208
209        let mut key_nonce = [0u8; 32];
210        key_nonce.copy_from_slice(&data[13..45]);
211
212        let mut key_iv = [0u8; 16];
213        key_iv.copy_from_slice(&data[45..61]);
214
215        // Skip RSC (8 bytes at 61..69) and reserved (8 bytes at 69..77)
216        let mut key_mic = [0u8; 16];
217        key_mic.copy_from_slice(&data[77..93]);
218
219        let key_data_length = u16::from_be_bytes([data[93], data[94]]);
220
221        let key_data = if key_data_length > 0 {
222            let end = Self::MIN_SIZE + key_data_length as usize;
223            if data.len() < end {
224                return None;
225            }
226            data[Self::MIN_SIZE..end].to_vec()
227        } else {
228            Vec::new()
229        };
230
231        Some(Self {
232            descriptor_type,
233            key_info,
234            key_length,
235            replay_counter,
236            key_nonce,
237            key_iv,
238            key_mic,
239            key_data_length,
240            key_data,
241        })
242    }
243
244    /// Serialize EAPOL-Key frame body to bytes
245    pub fn to_bytes(&self) -> Vec<u8> {
246        let mut buf = Vec::with_capacity(Self::MIN_SIZE + self.key_data.len());
247        buf.push(self.descriptor_type);
248        buf.extend_from_slice(&self.key_info.to_u16().to_be_bytes());
249        buf.extend_from_slice(&self.key_length.to_be_bytes());
250        buf.extend_from_slice(&self.replay_counter.to_be_bytes());
251        buf.extend_from_slice(&self.key_nonce);
252        buf.extend_from_slice(&self.key_iv);
253        buf.extend_from_slice(&[0u8; 8]); // RSC
254        buf.extend_from_slice(&[0u8; 8]); // Reserved
255        buf.extend_from_slice(&self.key_mic);
256        buf.extend_from_slice(&self.key_data_length.to_be_bytes());
257        buf.extend_from_slice(&self.key_data);
258        buf
259    }
260}
261
262// ============================================================================
263// WPA State Machine
264// ============================================================================
265
266/// WPA handshake state
267#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
268pub enum WpaState {
269    /// Idle, waiting for handshake to begin
270    #[default]
271    Idle,
272    /// PTK start: received message 1, computing PTK
273    PtkStart,
274    /// PTK negotiating: sent message 2, waiting for message 3
275    PtkInitNegotiating,
276    /// PTK done: received message 3, sent message 4
277    PtkInitDone,
278    /// GTK negotiating: group key handshake in progress
279    GtkNegotiating,
280    /// Handshake completed, keys installed
281    Completed,
282    /// Handshake failed
283    Failed,
284}
285
286// ============================================================================
287// Key Derivation Functions
288// ============================================================================
289
290/// PBKDF2-SHA256 key derivation (RFC 8018)
291///
292/// Derives a key from a password and salt using HMAC-SHA256 as the PRF.
293/// Uses 4096 iterations as specified by WPA2.
294pub fn pbkdf2_sha256(password: &[u8], salt: &[u8], iterations: u32, output: &mut [u8]) {
295    let dk_len = output.len();
296    let h_len = 32; // SHA-256 output size
297    let blocks_needed = dk_len.div_ceil(h_len);
298
299    for block_idx in 0..blocks_needed {
300        let block_num = (block_idx as u32) + 1;
301
302        // U_1 = PRF(password, salt || INT(block_num))
303        let mut salt_block = Vec::with_capacity(salt.len() + 4);
304        salt_block.extend_from_slice(salt);
305        salt_block.extend_from_slice(&block_num.to_be_bytes());
306
307        let mut u_prev = hmac_sha256(password, &salt_block);
308        let mut result = u_prev;
309
310        // U_2 .. U_c: PRF(password, U_{i-1}), XOR into result
311        for _ in 1..iterations {
312            let u_curr = hmac_sha256(password, &u_prev);
313            for (r, &u) in result.iter_mut().zip(u_curr.iter()) {
314                *r ^= u;
315            }
316            u_prev = u_curr;
317        }
318
319        // Copy derived block to output
320        let start = block_idx * h_len;
321        let end = core::cmp::min(start + h_len, dk_len);
322        output[start..end].copy_from_slice(&result[..end - start]);
323    }
324}
325
326/// PRF-SHA256: Pseudo-Random Function for key expansion (IEEE 802.11-2012)
327///
328/// Expands a key using HMAC-SHA256 with label and data inputs.
329/// Produces output of the specified bit length.
330pub fn prf_sha256(key: &[u8], label: &[u8], data: &[u8], bits: usize) -> Vec<u8> {
331    let bytes_needed = bits.div_ceil(8);
332    let iterations = bytes_needed.div_ceil(32);
333    let mut result = Vec::with_capacity(iterations * 32);
334
335    for i in 0..iterations {
336        // HMAC-SHA256(key, label || 0x00 || data || i)
337        let mut input = Vec::with_capacity(label.len() + 1 + data.len() + 1);
338        input.extend_from_slice(label);
339        input.push(0x00);
340        input.extend_from_slice(data);
341        input.push(i as u8);
342
343        let block = hmac_sha256(key, &input);
344        result.extend_from_slice(&block);
345    }
346
347    result.truncate(bytes_needed);
348    result
349}
350
351/// Derive PMK (Pairwise Master Key) from passphrase and SSID.
352///
353/// PMK = PBKDF2-SHA256(passphrase, SSID, 4096, 32)
354pub fn derive_pmk(passphrase: &[u8], ssid: &[u8]) -> [u8; 32] {
355    let mut pmk = [0u8; 32];
356    pbkdf2_sha256(passphrase, ssid, 4096, &mut pmk);
357    pmk
358}
359
360/// Derive PTK (Pairwise Transient Key) from PMK, nonces, and addresses.
361///
362/// PTK = PRF-SHA256(PMK, "Pairwise key expansion",
363///                  Min(AA,SPA) || Max(AA,SPA) || Min(ANonce,SNonce) ||
364/// Max(ANonce,SNonce))
365///
366/// PTK is split into: KCK (16) + KEK (16) + TK (16) = 48 bytes
367pub fn derive_ptk(
368    pmk: &[u8; 32],
369    aa: &MacAddress,
370    spa: &MacAddress,
371    anonce: &[u8; 32],
372    snonce: &[u8; 32],
373) -> TemporalKey {
374    // Sort addresses and nonces (smaller first)
375    let (min_addr, max_addr) = if aa.0 < spa.0 {
376        (&aa.0[..], &spa.0[..])
377    } else {
378        (&spa.0[..], &aa.0[..])
379    };
380    let (min_nonce, max_nonce) = if anonce < snonce {
381        (&anonce[..], &snonce[..])
382    } else {
383        (&snonce[..], &anonce[..])
384    };
385
386    let mut data = Vec::with_capacity(6 + 6 + 32 + 32);
387    data.extend_from_slice(min_addr);
388    data.extend_from_slice(max_addr);
389    data.extend_from_slice(min_nonce);
390    data.extend_from_slice(max_nonce);
391
392    let ptk_bytes = prf_sha256(pmk, b"Pairwise key expansion", &data, 384);
393
394    let mut kck = [0u8; 16];
395    let mut kek = [0u8; 16];
396    let mut tk = [0u8; 16];
397    kck.copy_from_slice(&ptk_bytes[0..16]);
398    kek.copy_from_slice(&ptk_bytes[16..32]);
399    tk.copy_from_slice(&ptk_bytes[32..48]);
400
401    TemporalKey { kck, kek, tk }
402}
403
404// ============================================================================
405// Temporal Key
406// ============================================================================
407
408/// Temporal key material derived from the 4-way handshake
409#[derive(Debug, Clone)]
410pub struct TemporalKey {
411    /// Key Confirmation Key (16 bytes) - used for MIC computation
412    pub kck: [u8; 16],
413    /// Key Encryption Key (16 bytes) - used for key data encryption
414    pub kek: [u8; 16],
415    /// Temporal Key (16 bytes) - used for data encryption (CCMP/TKIP)
416    pub tk: [u8; 16],
417}
418
419// ============================================================================
420// 4-Way Handshake
421// ============================================================================
422
423/// WPA2 4-way handshake state machine (supplicant side)
424pub struct FourWayHandshake {
425    /// Handshake state
426    state: WpaState,
427    /// Pairwise Master Key (32 bytes)
428    pmk: [u8; 32],
429    /// Derived Pairwise Transient Key
430    ptk: Option<TemporalKey>,
431    /// Authenticator nonce (from AP, message 1)
432    anonce: [u8; 32],
433    /// Supplicant nonce (generated locally)
434    snonce: [u8; 32],
435    /// Our MAC address (SPA)
436    spa: MacAddress,
437    /// AP MAC address (AA)
438    aa: MacAddress,
439    /// Replay counter from last message
440    replay_counter: u64,
441}
442
443impl FourWayHandshake {
444    /// Create a new handshake instance
445    pub fn new(pmk: [u8; 32], spa: MacAddress, aa: MacAddress) -> Self {
446        Self {
447            state: WpaState::Idle,
448            pmk,
449            ptk: None,
450            anonce: [0u8; 32],
451            snonce: [0u8; 32],
452            spa,
453            aa,
454            replay_counter: 0,
455        }
456    }
457
458    /// Get current handshake state
459    pub fn state(&self) -> WpaState {
460        self.state
461    }
462
463    /// Get derived PTK (available after message 1 processing)
464    pub fn ptk(&self) -> Option<&TemporalKey> {
465        self.ptk.as_ref()
466    }
467
468    /// Process Message 1 from AP: receive ANonce, generate SNonce, derive PTK.
469    ///
470    /// Returns Message 2 EAPOL-Key frame bytes to send back.
471    pub fn process_message_1(&mut self, key_frame: &EapolKeyFrame) -> Option<Vec<u8>> {
472        if self.state != WpaState::Idle {
473            return None;
474        }
475
476        // Verify this is Message 1: ACK set, MIC not set, pairwise set
477        if !key_frame.key_info.ack || key_frame.key_info.mic || !key_frame.key_info.pairwise {
478            return None;
479        }
480
481        // Store ANonce
482        self.anonce = key_frame.key_nonce;
483        self.replay_counter = key_frame.replay_counter;
484
485        // Generate SNonce (in a real implementation, use CSPRNG)
486        // For now, derive from PMK + ANonce as a deterministic placeholder
487        self.snonce = hmac_sha256(&self.pmk, &self.anonce);
488
489        // Derive PTK
490        self.ptk = Some(derive_ptk(
491            &self.pmk,
492            &self.aa,
493            &self.spa,
494            &self.anonce,
495            &self.snonce,
496        ));
497
498        self.state = WpaState::PtkStart;
499
500        // Build Message 2
501        Some(self.build_message_2())
502    }
503
504    /// Process Message 3 from AP: verify MIC, install PTK.
505    ///
506    /// Returns Message 4 EAPOL-Key frame bytes to send back.
507    pub fn process_message_3(&mut self, key_frame: &EapolKeyFrame) -> Option<Vec<u8>> {
508        if self.state != WpaState::PtkInitNegotiating && self.state != WpaState::PtkStart {
509            return None;
510        }
511
512        // Verify this is Message 3: ACK, MIC, pairwise, install, secure all set
513        if !key_frame.key_info.ack
514            || !key_frame.key_info.mic
515            || !key_frame.key_info.pairwise
516            || !key_frame.key_info.install
517        {
518            return None;
519        }
520
521        // Verify replay counter is >= our stored value
522        if key_frame.replay_counter < self.replay_counter {
523            self.state = WpaState::Failed;
524            return None;
525        }
526        self.replay_counter = key_frame.replay_counter;
527
528        // Verify MIC using KCK
529        let ptk = self.ptk.as_ref()?;
530        if !self.verify_mic_internal(&ptk.kck, key_frame) {
531            self.state = WpaState::Failed;
532            return None;
533        }
534
535        // Verify ANonce matches message 1
536        if key_frame.key_nonce != self.anonce {
537            self.state = WpaState::Failed;
538            return None;
539        }
540
541        self.state = WpaState::PtkInitDone;
542
543        // Build Message 4
544        Some(self.build_message_4())
545    }
546
547    /// Process Message 2 (AP/authenticator side): verify MIC from supplicant.
548    ///
549    /// Returns true if Message 2 is valid.
550    pub fn process_message_2(&mut self, key_frame: &EapolKeyFrame) -> bool {
551        // Message 2: MIC set, pairwise set, ACK not set
552        if !key_frame.key_info.mic || key_frame.key_info.ack || !key_frame.key_info.pairwise {
553            return false;
554        }
555
556        // In authenticator role, we need the SNonce from this message to derive PTK
557        self.snonce = key_frame.key_nonce;
558
559        // Derive PTK
560        let ptk = derive_ptk(&self.pmk, &self.aa, &self.spa, &self.anonce, &self.snonce);
561
562        // Verify MIC
563        if !self.verify_mic_internal(&ptk.kck, key_frame) {
564            self.state = WpaState::Failed;
565            return false;
566        }
567
568        self.ptk = Some(ptk);
569        self.state = WpaState::PtkInitNegotiating;
570        true
571    }
572
573    /// Process Message 4 (AP/authenticator side): handshake complete.
574    ///
575    /// Returns true if Message 4 is valid.
576    pub fn process_message_4(&mut self, key_frame: &EapolKeyFrame) -> bool {
577        if self.state != WpaState::PtkInitDone && self.state != WpaState::PtkInitNegotiating {
578            return false;
579        }
580
581        // Message 4: MIC set, pairwise set, ACK not set, secure set
582        if !key_frame.key_info.mic || !key_frame.key_info.pairwise || key_frame.key_info.ack {
583            return false;
584        }
585
586        if let Some(ref ptk) = self.ptk {
587            if !self.verify_mic_internal(&ptk.kck, key_frame) {
588                self.state = WpaState::Failed;
589                return false;
590            }
591        } else {
592            self.state = WpaState::Failed;
593            return false;
594        }
595
596        self.state = WpaState::Completed;
597        true
598    }
599
600    /// Verify MIC on an EAPOL-Key frame using HMAC-SHA256(KCK,
601    /// frame_with_zero_mic)
602    pub fn verify_mic(kck: &[u8; 16], key_frame: &EapolKeyFrame) -> bool {
603        // Reconstruct frame with MIC zeroed for verification
604        let mut frame_copy = key_frame.clone();
605        let original_mic = frame_copy.key_mic;
606        frame_copy.key_mic = [0u8; 16];
607
608        let frame_bytes = frame_copy.to_bytes();
609
610        // Wrap in EAPOL header for MIC computation
611        let mut eapol_bytes = Vec::with_capacity(4 + frame_bytes.len());
612        eapol_bytes.push(2); // version
613        eapol_bytes.push(EAPOL_KEY);
614        eapol_bytes.extend_from_slice(&(frame_bytes.len() as u16).to_be_bytes());
615        eapol_bytes.extend_from_slice(&frame_bytes);
616
617        let computed_mic = hmac_sha256(kck, &eapol_bytes);
618
619        // Compare first 16 bytes of HMAC with stored MIC
620        computed_mic[..16] == original_mic
621    }
622
623    /// Internal MIC verification helper
624    fn verify_mic_internal(&self, kck: &[u8; 16], key_frame: &EapolKeyFrame) -> bool {
625        Self::verify_mic(kck, key_frame)
626    }
627
628    /// Build Message 2 (supplicant to AP)
629    fn build_message_2(&self) -> Vec<u8> {
630        let key_info = KeyInfo {
631            descriptor_version: 2, // HMAC-SHA1-128 / AES-128-CMAC
632            pairwise: true,
633            mic: true,
634            ..Default::default()
635        };
636
637        let mut key_frame = EapolKeyFrame {
638            descriptor_type: 2, // RSN
639            key_info,
640            key_length: 0,
641            replay_counter: self.replay_counter,
642            key_nonce: self.snonce,
643            key_iv: [0u8; 16],
644            key_mic: [0u8; 16],
645            key_data_length: 0,
646            key_data: Vec::new(),
647        };
648
649        // Compute MIC over the frame (with MIC field zeroed)
650        if let Some(ref ptk) = self.ptk {
651            let frame_bytes = key_frame.to_bytes();
652            let mut eapol_bytes = Vec::with_capacity(4 + frame_bytes.len());
653            eapol_bytes.push(2);
654            eapol_bytes.push(EAPOL_KEY);
655            eapol_bytes.extend_from_slice(&(frame_bytes.len() as u16).to_be_bytes());
656            eapol_bytes.extend_from_slice(&frame_bytes);
657
658            let mic = hmac_sha256(&ptk.kck, &eapol_bytes);
659            key_frame.key_mic[..16].copy_from_slice(&mic[..16]);
660        }
661
662        // Wrap in EAPOL frame
663        let body = key_frame.to_bytes();
664        let eapol = EapolFrame {
665            protocol_version: 2,
666            packet_type: EAPOL_KEY,
667            packet_body_length: body.len() as u16,
668            body,
669        };
670        eapol.to_bytes()
671    }
672
673    /// Build Message 4 (supplicant to AP)
674    fn build_message_4(&self) -> Vec<u8> {
675        let key_info = KeyInfo {
676            descriptor_version: 2,
677            pairwise: true,
678            mic: true,
679            secure: true,
680            ..Default::default()
681        };
682
683        let mut key_frame = EapolKeyFrame {
684            descriptor_type: 2,
685            key_info,
686            key_length: 0,
687            replay_counter: self.replay_counter,
688            key_nonce: [0u8; 32],
689            key_iv: [0u8; 16],
690            key_mic: [0u8; 16],
691            key_data_length: 0,
692            key_data: Vec::new(),
693        };
694
695        // Compute MIC
696        if let Some(ref ptk) = self.ptk {
697            let frame_bytes = key_frame.to_bytes();
698            let mut eapol_bytes = Vec::with_capacity(4 + frame_bytes.len());
699            eapol_bytes.push(2);
700            eapol_bytes.push(EAPOL_KEY);
701            eapol_bytes.extend_from_slice(&(frame_bytes.len() as u16).to_be_bytes());
702            eapol_bytes.extend_from_slice(&frame_bytes);
703
704            let mic = hmac_sha256(&ptk.kck, &eapol_bytes);
705            key_frame.key_mic[..16].copy_from_slice(&mic[..16]);
706        }
707
708        let body = key_frame.to_bytes();
709        let eapol = EapolFrame {
710            protocol_version: 2,
711            packet_type: EAPOL_KEY,
712            packet_body_length: body.len() as u16,
713            body,
714        };
715        eapol.to_bytes()
716    }
717}
718
719// ============================================================================
720// CCMP (AES-128-CCM) Stub
721// ============================================================================
722
723/// CCMP encryption parameters
724#[derive(Debug, Clone)]
725pub struct CcmpEncrypt {
726    /// Temporal key (16 bytes, AES-128)
727    tk: [u8; 16],
728    /// Packet number (48-bit, incremented per frame)
729    packet_number: u64,
730}
731
732impl CcmpEncrypt {
733    /// Create a new CCMP encryption context
734    pub fn new(tk: [u8; 16]) -> Self {
735        Self {
736            tk,
737            packet_number: 0,
738        }
739    }
740
741    /// Get the temporal key
742    pub fn temporal_key(&self) -> &[u8; 16] {
743        &self.tk
744    }
745
746    /// Encrypt a frame payload using AES-128-CCM.
747    ///
748    /// Stub: returns data with CCMP header prepended and 8-byte MIC appended.
749    /// A full implementation would use AES-128 in CCM mode.
750    pub fn encrypt(&mut self, _header: &[u8], data: &[u8]) -> Vec<u8> {
751        let pn = self.packet_number;
752        self.packet_number += 1;
753
754        // CCMP header (8 bytes): PN0, PN1, 0, ExtIV|KeyID, PN2, PN3, PN4, PN5
755        let mut result = Vec::with_capacity(8 + data.len() + 8);
756        result.push((pn & 0xFF) as u8);
757        result.push(((pn >> 8) & 0xFF) as u8);
758        result.push(0); // Reserved
759        result.push(0x20); // ExtIV=1, KeyID=0
760        result.push(((pn >> 16) & 0xFF) as u8);
761        result.push(((pn >> 24) & 0xFF) as u8);
762        result.push(((pn >> 32) & 0xFF) as u8);
763        result.push(((pn >> 40) & 0xFF) as u8);
764
765        // Stub: copy plaintext (real impl would AES-CCM encrypt)
766        result.extend_from_slice(data);
767
768        // Stub MIC (8 bytes, real impl computes via AES-CCM)
769        result.extend_from_slice(&[0u8; 8]);
770
771        result
772    }
773
774    /// Decrypt a frame payload using AES-128-CCM.
775    ///
776    /// Stub: strips CCMP header and MIC, returns payload.
777    /// A full implementation would verify MIC and decrypt via AES-128-CCM.
778    pub fn decrypt(&self, data: &[u8]) -> Option<Vec<u8>> {
779        // Need at least CCMP header (8) + MIC (8)
780        if data.len() < 16 {
781            return None;
782        }
783        // Strip header and MIC, return payload
784        Some(data[8..data.len() - 8].to_vec())
785    }
786}
787
788// ============================================================================
789// WPA3-SAE Stubs
790// ============================================================================
791
792/// SAE (Simultaneous Authentication of Equals) state for WPA3
793#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
794pub enum SaeState {
795    /// Idle, not started
796    #[default]
797    Idle,
798    /// Commit exchange: sending/receiving commit messages
799    Committed,
800    /// Confirm exchange: sending/receiving confirm messages
801    Confirmed,
802    /// SAE authentication completed
803    Accepted,
804    /// SAE authentication failed
805    Failed,
806}
807
808/// WPA3-SAE authentication handler (stub)
809///
810/// Full implementation requires elliptic curve operations (hunting-and-pecking
811/// for password element, Diffie-Hellman exchange). This stub provides the
812/// state machine and message framing.
813pub struct SaeAuth {
814    /// Current SAE state
815    state: SaeState,
816    /// Our scalar (32 bytes, stub)
817    _scalar: [u8; 32],
818    /// Our element (64 bytes, stub - EC point)
819    _element: [u8; 64],
820}
821
822impl SaeAuth {
823    /// Create a new SAE authentication instance
824    pub fn new() -> Self {
825        Self {
826            state: SaeState::Idle,
827            _scalar: [0u8; 32],
828            _element: [0u8; 64],
829        }
830    }
831
832    /// Get current SAE state
833    pub fn state(&self) -> SaeState {
834        self.state
835    }
836
837    /// Generate SAE commit message.
838    ///
839    /// Stub: returns placeholder commit frame. Full implementation requires
840    /// hunting-and-pecking to derive password element on NIST P-256, then
841    /// scalar/element generation via Diffie-Hellman.
842    pub fn generate_commit(&mut self, _password: &[u8], _own_addr: &MacAddress) -> Vec<u8> {
843        self.state = SaeState::Committed;
844
845        // Stub commit frame: group_id(2) + scalar(32) + element(64)
846        let mut frame = Vec::with_capacity(98);
847        frame.extend_from_slice(&19u16.to_le_bytes()); // Group 19 = NIST P-256
848        frame.extend_from_slice(&self._scalar);
849        frame.extend_from_slice(&self._element);
850        frame
851    }
852
853    /// Process received SAE commit message.
854    ///
855    /// Stub: validates message length and transitions state.
856    pub fn process_commit(&mut self, data: &[u8]) -> bool {
857        if self.state != SaeState::Committed {
858            return false;
859        }
860        // Minimum commit: group_id(2) + scalar(32) + element(64)
861        if data.len() < 98 {
862            self.state = SaeState::Failed;
863            return false;
864        }
865        self.state = SaeState::Confirmed;
866        true
867    }
868
869    /// Generate SAE confirm message.
870    ///
871    /// Stub: returns placeholder confirm frame.
872    pub fn generate_confirm(&mut self) -> Vec<u8> {
873        // Stub confirm frame: send_confirm(2) + confirm(32)
874        let mut frame = Vec::with_capacity(34);
875        frame.extend_from_slice(&1u16.to_le_bytes()); // send_confirm counter
876        frame.extend_from_slice(&[0u8; 32]); // Stub confirm value
877        frame
878    }
879
880    /// Process received SAE confirm message.
881    ///
882    /// Stub: validates message length and completes authentication.
883    pub fn process_confirm(&mut self, data: &[u8]) -> bool {
884        if self.state != SaeState::Confirmed {
885            return false;
886        }
887        if data.len() < 34 {
888            self.state = SaeState::Failed;
889            return false;
890        }
891        self.state = SaeState::Accepted;
892        true
893    }
894}
895
896impl Default for SaeAuth {
897    fn default() -> Self {
898        Self::new()
899    }
900}
901
902// ============================================================================
903// Tests
904// ============================================================================
905
906#[cfg(test)]
907mod tests {
908    #[allow(unused_imports)]
909    use alloc::vec;
910
911    use super::*;
912
913    #[test]
914    fn test_hmac_sha256_known_vector() {
915        // RFC 4231 Test Case 2: HMAC-SHA256 with "Jefe" key and "what do ya want for
916        // nothing?" data
917        let key = b"Jefe";
918        let data = b"what do ya want for nothing?";
919        let result = hmac_sha256(key, data);
920        // Expected: 5bdcc146bf60754e6a042426089575c75a003f089d2739839dec58b964ec3843
921        let expected: [u8; 32] = [
922            0x5b, 0xdc, 0xc1, 0x46, 0xbf, 0x60, 0x75, 0x4e, 0x6a, 0x04, 0x24, 0x26, 0x08, 0x95,
923            0x75, 0xc7, 0x5a, 0x00, 0x3f, 0x08, 0x9d, 0x27, 0x39, 0x83, 0x9d, 0xec, 0x58, 0xb9,
924            0x64, 0xec, 0x38, 0x43,
925        ];
926        assert_eq!(result, expected);
927    }
928
929    #[test]
930    fn test_hmac_sha256_empty() {
931        let key = &[];
932        let data = &[];
933        let result = hmac_sha256(key, data);
934        assert_eq!(result.len(), 32);
935    }
936
937    #[test]
938    fn test_pbkdf2_sha256_basic() {
939        let mut output = [0u8; 32];
940        pbkdf2_sha256(b"password", b"salt", 1, &mut output);
941        assert_ne!(output, [0u8; 32]);
942    }
943
944    #[test]
945    fn test_derive_pmk() {
946        let pmk = derive_pmk(b"testpassword", b"TestSSID");
947        assert_ne!(pmk, [0u8; 32]);
948        assert_eq!(pmk.len(), 32);
949
950        // Same inputs should produce same PMK
951        let pmk2 = derive_pmk(b"testpassword", b"TestSSID");
952        assert_eq!(pmk, pmk2);
953    }
954
955    #[test]
956    fn test_derive_ptk() {
957        let pmk = [0x42u8; 32];
958        let aa = MacAddress::new([0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF]);
959        let spa = MacAddress::new([0x11, 0x22, 0x33, 0x44, 0x55, 0x66]);
960        let anonce = [0x01u8; 32];
961        let snonce = [0x02u8; 32];
962
963        let ptk = derive_ptk(&pmk, &aa, &spa, &anonce, &snonce);
964        assert_eq!(ptk.kck.len(), 16);
965        assert_eq!(ptk.kek.len(), 16);
966        assert_eq!(ptk.tk.len(), 16);
967
968        // Deterministic: same inputs -> same output
969        let ptk2 = derive_ptk(&pmk, &aa, &spa, &anonce, &snonce);
970        assert_eq!(ptk.kck, ptk2.kck);
971        assert_eq!(ptk.kek, ptk2.kek);
972        assert_eq!(ptk.tk, ptk2.tk);
973    }
974
975    #[test]
976    fn test_prf_sha256_output_length() {
977        let key = [0x42u8; 32];
978        let result = prf_sha256(&key, b"test label", b"test data", 384);
979        assert_eq!(result.len(), 48); // 384 bits = 48 bytes
980    }
981
982    #[test]
983    fn test_eapol_frame_roundtrip() {
984        let frame = EapolFrame {
985            protocol_version: 2,
986            packet_type: EAPOL_KEY,
987            packet_body_length: 4,
988            body: vec![0x01, 0x02, 0x03, 0x04],
989        };
990        let bytes = frame.to_bytes();
991        let parsed = EapolFrame::from_bytes(&bytes).unwrap();
992        assert_eq!(parsed.protocol_version, 2);
993        assert_eq!(parsed.packet_type, EAPOL_KEY);
994        assert_eq!(parsed.body, vec![0x01, 0x02, 0x03, 0x04]);
995    }
996
997    #[test]
998    fn test_eapol_frame_too_short() {
999        assert!(EapolFrame::from_bytes(&[0, 3]).is_none());
1000    }
1001
1002    #[test]
1003    fn test_key_info_roundtrip() {
1004        let ki = KeyInfo {
1005            descriptor_version: 2,
1006            pairwise: true,
1007            install: true,
1008            ack: true,
1009            mic: true,
1010            secure: false,
1011            error: false,
1012            request: false,
1013            encrypted: false,
1014        };
1015        let val = ki.to_u16();
1016        let parsed = KeyInfo::from_u16(val);
1017        assert_eq!(parsed.descriptor_version, 2);
1018        assert!(parsed.pairwise);
1019        assert!(parsed.install);
1020        assert!(parsed.ack);
1021        assert!(parsed.mic);
1022        assert!(!parsed.secure);
1023    }
1024
1025    #[test]
1026    fn test_eapol_key_frame_roundtrip() {
1027        let kf = EapolKeyFrame {
1028            descriptor_type: 2,
1029            key_info: KeyInfo {
1030                descriptor_version: 2,
1031                pairwise: true,
1032                ack: true,
1033                ..Default::default()
1034            },
1035            key_length: 16,
1036            replay_counter: 1,
1037            key_nonce: [0xAB; 32],
1038            key_iv: [0; 16],
1039            key_mic: [0; 16],
1040            key_data_length: 0,
1041            key_data: Vec::new(),
1042        };
1043        let bytes = kf.to_bytes();
1044        let parsed = EapolKeyFrame::from_bytes(&bytes).unwrap();
1045        assert_eq!(parsed.descriptor_type, 2);
1046        assert_eq!(parsed.key_nonce, [0xAB; 32]);
1047        assert_eq!(parsed.replay_counter, 1);
1048        assert_eq!(parsed.key_length, 16);
1049    }
1050
1051    #[test]
1052    fn test_ccmp_encrypt_decrypt() {
1053        let mut ccmp = CcmpEncrypt::new([0x42u8; 16]);
1054        let header = [0u8; 24];
1055        let plaintext = b"Hello WiFi";
1056
1057        let encrypted = ccmp.encrypt(&header, plaintext);
1058        // Should have CCMP header (8) + data (10) + MIC (8) = 26
1059        assert_eq!(encrypted.len(), 8 + plaintext.len() + 8);
1060
1061        let decrypted = ccmp.decrypt(&encrypted).unwrap();
1062        assert_eq!(decrypted, plaintext);
1063    }
1064
1065    #[test]
1066    fn test_ccmp_decrypt_too_short() {
1067        let ccmp = CcmpEncrypt::new([0u8; 16]);
1068        assert!(ccmp.decrypt(&[0u8; 10]).is_none());
1069    }
1070
1071    #[test]
1072    fn test_sae_state_machine() {
1073        let mut sae = SaeAuth::new();
1074        assert_eq!(sae.state(), SaeState::Idle);
1075
1076        let _commit = sae.generate_commit(b"password", &MacAddress::ZERO);
1077        assert_eq!(sae.state(), SaeState::Committed);
1078
1079        // Simulate receiving peer commit (98 bytes minimum)
1080        let peer_commit = vec![0u8; 98];
1081        assert!(sae.process_commit(&peer_commit));
1082        assert_eq!(sae.state(), SaeState::Confirmed);
1083
1084        let _confirm = sae.generate_confirm();
1085
1086        let peer_confirm = vec![0u8; 34];
1087        assert!(sae.process_confirm(&peer_confirm));
1088        assert_eq!(sae.state(), SaeState::Accepted);
1089    }
1090}