#!/usr/bin/env python3
"""Drives the real MLME through a complete join and validates every frame it
puts on the air, plus the 802.11 <-> Ethernet translation in both directions.

The access point side is written independently here (hashlib / cryptography for
the handshake, hand-decoded 802.11 for the frames), so this is a check of what
the driver actually emits rather than a restatement of it.

Copyright (c) 2026 Daniel Hammer
"""
import hashlib, hmac, os, struct, subprocess, sys
from cryptography.hazmat.primitives.keywrap import aes_key_wrap

HARNESS = sys.argv[1]
AP  = bytes.fromhex('001122334455')
STA = bytes.fromhex('aabbccddeeff')
SSID = "MontaukTest"
PASS = "supersecret123"
CHANNEL = 6

fails = []
def state_of(result):
    f = result.split()
    return int(f[f.index('STATE') + 1]) if 'STATE' in f else -1

def link_of(result):
    f = result.split()
    return int(f[f.index('LINK') + 1]) if 'LINK' in f else -1

def check(cond, what):
    print(f"  {'PASS' if cond else 'FAIL'}  {what}")
    if not cond:
        fails.append(what)

# --------------------------------------------------------------- harness I/O
class Driver:
    def __init__(self):
        self.p = subprocess.Popen([HARNESS], stdin=subprocess.PIPE,
                                  stdout=subprocess.PIPE, text=True, bufsize=1)
    def cmd(self, line):
        """Send a command; collect emitted events until the terminator."""
        self.p.stdin.write(line + '\n'); self.p.stdin.flush()
        ev = {'TX': [], 'KEY': [], 'KEYDEL': [], 'ETH': [], 'CMD': [], 'TXQ': []}
        while True:
            l = self.p.stdout.readline()
            if not l:
                raise RuntimeError("harness died")
            l = l.strip()
            k = l.split()[0] if l else ''
            if k == 'TX':
                f = l.split()
                ev['TX'].append({'enc': f[1] == 'enc=1', 'rate': f[2] == 'rate=1',
                                 'hdr': bytes.fromhex(f[3]),
                                 'body': bytes.fromhex(f[4]) if len(f) > 4 else b''})
            elif k == 'KEY':
                f = l.split()
                ev['KEY'].append({'pairwise': f[1] == 'pairwise=1',
                                  'idx': int(f[2].split('=')[1]),
                                  'cipher': int(f[3].split('=')[1]),
                                  'key': bytes.fromhex(f[4]),
                                  'rsc': bytes.fromhex(f[5]) if len(f) > 5 else b''})
            elif k == 'ETH':
                ev['ETH'].append(bytes.fromhex(l.split()[1]))
            elif k == 'CMD':
                f = l.split()
                ev['CMD'].append((int(f[1]),
                                  bytes.fromhex(f[2]) if len(f) > 2 else b''))
            elif k == 'KEY-REMOVE':
                ev['KEYDEL'].append(l)
            elif k in ('TXQ-UP', 'TXQ-DOWN'):
                ev['TXQ'].append(l)
            else:
                ev['result'] = l
                return ev

# --------------------------------------------------------------- 802.11 bits
def mgmt(subtype, dst, src, bssid, body, seq=0):
    return bytes([0x00 | subtype, 0x00]) + b'\x00\x00' + dst + src + bssid \
           + struct.pack('<H', seq << 4) + body

def data_from_ds(da, bssid, sa, ethertype, payload, protected=False):
    fc1 = 0x02 | (0x40 if protected else 0)
    hdr = bytes([0x08, fc1]) + b'\x00\x00' + da + bssid + sa + b'\x00\x00'
    # The firmware decrypts in place and strips the MIC, but leaves the
    # 8-byte CCMP header, so model that when the protected bit is set.
    iv = b'\x11\x22\x00\x20\x00\x00\x00\x00' if protected else b''
    snap = bytes([0xaa, 0xaa, 0x03, 0, 0, 0]) + struct.pack('>H', ethertype)
    return hdr + iv + snap + payload

def parse_ies(b):
    ies, off = {}, 0
    while off + 2 <= len(b):
        i, l = b[off], b[off + 1]
        if off + 2 + l > len(b): break
        ies[i] = b[off + 2:off + 2 + l]
        off += 2 + l
    return ies

def suite(t): return bytes([0x00, 0x0f, 0xac, t])
AP_RSN = (struct.pack('<H', 1) + suite(4)
          + struct.pack('<H', 1) + suite(4)
          + struct.pack('<H', 1) + suite(2)
          + struct.pack('<H', 0))

# ------------------------------------------------------------------- EAPOL
def prf(key, label, data, n):
    r, i = b'', 0
    while len(r) < n:
        r += hmac.new(key, label.encode() + b'\x00' + data + bytes([i]), hashlib.sha1).digest(); i += 1
    return r[:n]
def derive_ptk(pmk, an, sn):
    return prf(pmk, "Pairwise key expansion",
               min(AP, STA) + max(AP, STA) + min(an, sn) + max(an, sn), 48)
def emic(kck, f): return hmac.new(kck, f, hashlib.sha1).digest()[:16]
HDR = 99
def eapol(ki, replay, nonce, rsc, kd, kck=None):
    b = bytearray(HDR + len(kd)); b[0] = 2; b[1] = 3
    struct.pack_into('>H', b, 2, HDR + len(kd) - 4); b[4] = 2
    struct.pack_into('>H', b, 5, ki); struct.pack_into('>H', b, 7, 16)
    b[9:17] = replay; b[17:49] = nonce; b[65:73] = rsc
    struct.pack_into('>H', b, 97, len(kd)); b[HDR:] = kd
    if ki & 0x0100: b[81:97] = emic(kck, bytes(b))
    return bytes(b)
def check_emic(f, kck):
    x = bytearray(f); x[81:97] = b'\x00' * 16
    return emic(kck, bytes(x)) == f[81:97]
def gtk_kde(gtk, idx):
    kd = bytes([0xdd, 6 + len(gtk), 0x00, 0x0f, 0xac, 0x01, idx, 0x00]) + gtk
    if len(kd) % 8: kd += b'\xdd' + b'\x00' * (7 - len(kd) % 8)
    return kd

# =============================================================================
d = Driver()
d.cmd("MAC " + STA.hex())

print("=== authentication ===")
ev = d.cmd(f"CONNECT {AP.hex()} {CHANNEL} 0 {SSID} {PASS} {AP_RSN.hex()}")
check(ev['result'].startswith('CONNECT-OK 1'), "connect starts")
check(len(ev['TXQ']) == 1 and ev['TXQ'][0].startswith('TXQ-UP'),
      "a transmit queue is opened before any frame is sent")
check(len(ev['TX']) == 1, "exactly one frame is sent (the authentication request)")
auth = ev['TX'][0]
h, b = auth['hdr'], auth['body']
check(h[0] == 0xb0, "authentication frame: type management, subtype auth")
check(h[4:10] == AP and h[10:16] == STA and h[16:22] == AP,
      "authentication frame: addressed to the AP, from us, BSSID correct")
check(len(h) == 24, "authentication frame: 24-byte header")
check(struct.unpack('<H', b[0:2])[0] == 0, "authentication: open system algorithm")
check(struct.unpack('<H', b[2:4])[0] == 1, "authentication: transaction sequence 1")
check(struct.unpack('<H', b[4:6])[0] == 0, "authentication: status 0")
check(not auth['enc'] and auth['rate'],
      "authentication frame sent in the clear at a fixed rate")

print("\n=== association ===")
d.cmd("RXMGMT " + mgmt(0xb0, STA, AP, AP, struct.pack('<HHH', 0, 2, 0)).hex())
ev = d.cmd("SERVICE")
check(len(ev['TX']) == 1, "authentication success triggers one association request")
ar = ev['TX'][0]
check(ar['hdr'][0] == 0x00, "association request: subtype assoc-req")
check(ar['hdr'][4:10] == AP and ar['hdr'][10:16] == STA,
      "association request: addressing correct")
caps = struct.unpack('<H', ar['body'][0:2])[0]
check(caps & 0x0001, "association request: ESS capability set")
check(caps & 0x0010, "association request: privacy bit set for an encrypted network")
li = struct.unpack('<H', ar['body'][2:4])[0]
check(li == 10, "association request: listen interval present")
ies = parse_ies(ar['body'][4:])
check(ies.get(0) == SSID.encode(), "association request: SSID element matches")
check(1 in ies and len(ies[1]) >= 4, "association request: supported rates present")
check(any(r & 0x80 for r in ies[1]), "association request: at least one basic rate")
check(48 in ies, "association request: RSN element present")
sta_rsn = ies[48]
check(sta_rsn[0:2] == struct.pack('<H', 1), "RSN element: version 1")
check(sta_rsn[2:6] == suite(4), "RSN element: CCMP group cipher")
check(sta_rsn[8:12] == suite(4), "RSN element: CCMP pairwise cipher")
check(sta_rsn[14:18] == suite(2), "RSN element: PSK AKM")

aid = 0x0007
d.cmd("RXMGMT " + mgmt(0x10, STA, AP, AP,
                       struct.pack('<HHH', 0x0431, 0, aid | 0xc000)).hex())
ev = d.cmd("SERVICE")
check(state_of(ev['result']) == 6, "association moves the link into the handshake")
check(len(ev['CMD']) >= 2, "post-association context updates are sent")

print("\n=== 4-way handshake carried over 802.11 data frames ===")
pmk = hashlib.pbkdf2_hmac('sha1', PASS.encode(), SSID.encode(), 4096, 32)
anonce = os.urandom(32)
m1 = eapol(0x0002 | 0x0008 | 0x0080, (1).to_bytes(8, 'big'), anonce, b'\x00' * 8, b'')
d.cmd("RXDATA " + data_from_ds(STA, AP, AP, 0x888e, m1).hex())
ev = d.cmd("SERVICE")
check(len(ev['TX']) == 1, "EAPOL message 1 produces exactly one reply")
m2f = ev['TX'][0]
h = m2f['hdr']
check(h[0] == 0x08 and (h[1] & 0x01), "EAPOL reply: data frame with to-DS set")
check(not (h[1] & 0x40), "EAPOL reply: not marked protected (no key yet)")
check(not m2f['enc'], "EAPOL reply: firmware told not to encrypt")
check(h[4:10] == AP and h[10:16] == STA and h[16:22] == AP,
      "EAPOL reply: addr1 the AP, addr2 us, addr3 the AP")
check(m2f['body'][0:6] == bytes([0xaa, 0xaa, 0x03, 0, 0, 0]),
      "EAPOL reply: RFC 1042 LLC/SNAP shim")
check(m2f['body'][6:8] == struct.pack('>H', 0x888e),
      "EAPOL reply: EtherType 0x888e")
m2 = m2f['body'][8:]
snonce = m2[17:49]
ptk = derive_ptk(pmk, anonce, snonce)
kck, kek, tk = ptk[:16], ptk[16:32], ptk[32:48]
check(check_emic(m2, kck), "EAPOL message 2 MIC verifies under the AP's own PTK")
check(m2[HDR:] == bytes([48, len(sta_rsn)]) + sta_rsn,
      "EAPOL message 2 carries the same RSN element as the association request")

gtk = os.urandom(16)
rsc = bytes.fromhex('0a0b0c0d0e0f0000')
m3 = eapol(0x0002 | 0x0008 | 0x0040 | 0x0080 | 0x0100 | 0x0200 | 0x1000,
           (2).to_bytes(8, 'big'), anonce, rsc, aes_key_wrap(kek, gtk_kde(gtk, 1)), kck)
d.cmd("RXDATA " + data_from_ds(STA, AP, AP, 0x888e, m3).hex())
ev = d.cmd("SERVICE")
check(len(ev['TX']) == 1, "EAPOL message 3 produces message 4")
m4 = ev['TX'][0]['body'][8:]
check(check_emic(m4, kck), "EAPOL message 4 MIC verifies")
keys = {('pairwise' if k['pairwise'] else 'group'): k for k in ev['KEY']}
check('pairwise' in keys and keys['pairwise']['key'] == tk,
      "pairwise key installed matches the AP's temporal key")
check(keys.get('pairwise', {}).get('cipher') == 4, "pairwise key installed as CCMP")
check('group' in keys and keys['group']['key'] == gtk, "group key installed matches")
check(keys.get('group', {}).get('rsc') == rsc[:6],
      "group key installed with the AP's receive sequence counter")
check(state_of(ev['result']) == 7 and link_of(ev['result']) == 1,
      "link reports up once keyed")

print("\n=== data path ===")
ip = bytes.fromhex('4500002800010000401100000a0000010a000002') + b'payload-here'
peer = bytes.fromhex('665544332211')
eth_out = peer + STA + struct.pack('>H', 0x0800) + ip
ev = d.cmd("TXETH " + eth_out.hex())
check(ev['result'] == 'TXETH-OK 1', "an Ethernet frame is accepted for transmit")
f = ev['TX'][0]
h = f['hdr']
check(h[0] == 0x08 and (h[1] & 0x01), "outbound data: to-DS data frame")
check(h[1] & 0x40, "outbound data: protected bit set now that keys are installed")
check(f['enc'], "outbound data: firmware asked to encrypt")
check(not f['rate'], "outbound data: rate control left to the firmware")
check(h[4:10] == AP, "outbound data: addr1 is the AP")
check(h[10:16] == STA, "outbound data: addr2 is us")
check(h[16:22] == peer, "outbound data: addr3 is the final destination")
check(f['body'][0:8] == bytes([0xaa, 0xaa, 0x03, 0, 0, 0]) + struct.pack('>H', 0x0800),
      "outbound data: LLC/SNAP carries the EtherType")
check(f['body'][8:] == ip, "outbound data: IP payload preserved byte for byte")
seq1 = struct.unpack('<H', h[22:24])[0]
ev2 = d.cmd("TXETH " + eth_out.hex())
seq2 = struct.unpack('<H', ev2['TX'][0]['hdr'][22:24])[0]
check(seq2 != seq1, "outbound data: the sequence number advances between frames")

reply = bytes.fromhex('450000300002000040110000' '0a000002' '0a000001') + b'inbound-payload'
ev = d.cmd("RXDATA " + data_from_ds(STA, AP, peer, 0x0800, reply, protected=True).hex())
check(len(ev['ETH']) == 1, "an encrypted inbound data frame yields one Ethernet frame")
e = ev['ETH'][0]
check(e[0:6] == STA, "inbound data: Ethernet destination is addr1")
check(e[6:12] == peer, "inbound data: Ethernet source is addr3")
check(e[12:14] == struct.pack('>H', 0x0800), "inbound data: EtherType recovered")
check(e[14:] == reply, "inbound data: payload recovered past the CCMP header")

bcast = bytes.fromhex('ffffffffffff')
ev = d.cmd("RXDATA " + data_from_ds(bcast, AP, peer, 0x0806, b'arp-request-body').hex())
check(len(ev['ETH']) == 1 and ev['ETH'][0][0:6] == bcast,
      "inbound broadcast (ARP) is delivered")

ev = d.cmd("RXDATA " + data_from_ds(STA, bytes.fromhex('aa0000000001'),
                                    peer, 0x0800, b'from-a-stranger').hex())
check(len(ev['ETH']) == 0, "a data frame from a different BSSID is ignored")

ev = d.cmd("RXDATA " + (bytes([0x48, 0x02]) + b'\x00\x00' + STA + AP + peer
                        + b'\x00\x00').hex())
check(len(ev['ETH']) == 0, "a null-data keepalive produces no Ethernet frame")

print("\n=== teardown ===")
ev = d.cmd("ABORT")
check(any(x.startswith('TXQ-DOWN') for x in ev['TXQ']), "abort closes the transmit queue")
check(len(ev['KEYDEL']) == 2, "abort removes both hardware keys before the station")
check(state_of(ev['result']) == 0, "abort returns to idle")
ev = d.cmd("TXETH " + eth_out.hex())
check(ev['result'] == 'TXETH-OK 0', "transmit is refused once the link is down")

# =============================================================================
# Other paths through the state machine, each on a fresh harness.
# =============================================================================

print("\n=== open network ===")
d2 = Driver()
d2.cmd("MAC " + STA.hex())
ev = d2.cmd(f"CONNECT {AP.hex()} 11 0 OpenNet - -")
check(ev['result'].startswith('CONNECT-OK 1'), "open network: connect starts")
caps_seen = []
d2.cmd("RXMGMT " + mgmt(0xb0, STA, AP, AP, struct.pack('<HHH', 0, 2, 0)).hex())
ev = d2.cmd("SERVICE")
ar = ev['TX'][0]
caps = struct.unpack('<H', ar['body'][0:2])[0]
check(not (caps & 0x0010), "open network: privacy bit clear")
ies = parse_ies(ar['body'][4:])
check(48 not in ies, "open network: no RSN element in the association request")
d2.cmd("RXMGMT " + mgmt(0x10, STA, AP, AP,
                        struct.pack('<HHH', 0x0421, 0, 3 | 0xc000)).hex())
ev = d2.cmd("SERVICE")
check(state_of(ev['result']) == 7 and link_of(ev['result']) == 1,
      "open network: link comes up straight after association, no handshake")
check(len(ev['KEY']) == 0, "open network: no keys installed")
ev = d2.cmd("TXETH " + (peer + STA + struct.pack('>H', 0x0800) + ip).hex())
f = ev['TX'][0]
check(not (f['hdr'][1] & 0x40) and not f['enc'],
      "open network: data frames are sent unprotected")

print("\n=== the access point never answers ===")
d3 = Driver()
d3.cmd("MAC " + STA.hex())
d3.cmd(f"CONNECT {AP.hex()} 6 0 Nowhere - -")
retries = 0
for _ in range(6):
    d3.cmd("TICK 500")
    ev = d3.cmd("SERVICE")
    retries += len(ev['TX'])
    if state_of(ev['result']) == 8:
        break
check(retries >= 3, f"authentication is retransmitted ({retries} retries) before giving up")
check(state_of(ev['result']) == 8, "the attempt eventually fails rather than hanging")
check(any(x.startswith('TXQ-DOWN') for x in ev['TXQ']),
      "giving up tears the transmit queue back down")

print("\n=== the access point rejects us ===")
d4 = Driver()
d4.cmd("MAC " + STA.hex())
d4.cmd(f"CONNECT {AP.hex()} 6 0 Rejects - -")
ev = d4.cmd("RXMGMT " + mgmt(0xb0, STA, AP, AP, struct.pack('<HHH', 0, 2, 1)).hex())
check(state_of(ev['result']) == 8, "an authentication rejection fails the attempt")

d5 = Driver()
d5.cmd("MAC " + STA.hex())
d5.cmd(f"CONNECT {AP.hex()} 6 0 Rejects - -")
d5.cmd("RXMGMT " + mgmt(0xb0, STA, AP, AP, struct.pack('<HHH', 0, 2, 0)).hex())
d5.cmd("SERVICE")
ev = d5.cmd("RXMGMT " + mgmt(0x10, STA, AP, AP,
                             struct.pack('<HHH', 0x0421, 17, 0)).hex())
check(state_of(ev['result']) == 8, "an association rejection fails the attempt")

print("\n=== the access point drops us ===")
d6 = Driver()
d6.cmd("MAC " + STA.hex())
d6.cmd(f"CONNECT {AP.hex()} 6 0 OpenNet - -")
d6.cmd("RXMGMT " + mgmt(0xb0, STA, AP, AP, struct.pack('<HHH', 0, 2, 0)).hex())
d6.cmd("SERVICE")
d6.cmd("RXMGMT " + mgmt(0x10, STA, AP, AP,
                        struct.pack('<HHH', 0x0421, 0, 3 | 0xc000)).hex())
ev = d6.cmd("SERVICE")
check(link_of(ev['result']) == 1, "link up before the deauthentication")
d6.cmd("RXMGMT " + mgmt(0xc0, STA, AP, AP, struct.pack('<H', 3)).hex())
ev = d6.cmd("SERVICE")
check(state_of(ev['result']) == 0 and link_of(ev['result']) == 0,
      "a deauthentication from the AP brings the link down")
ev = d6.cmd("TXETH " + (peer + STA + struct.pack('>H', 0x0800) + ip).hex())
check(ev['result'] == 'TXETH-OK 0', "transmit is refused after being dropped")

# =============================================================================
# PHY context command layout.
#
# Regression test for a lockup: firmware 89 advertises ULTRA_HB_CHANNELS, which
# selects an 8-byte channel-info sub-structure (32-bit channel first).  Sending
# the older 4-byte form shifted every field after it, asserted the firmware, and
# the driver then spun forever waiting for a reply that would never come.
# =============================================================================

print("\n=== PHY context command layout ===")
PHY_CONTEXT_CMD = 0x08
ULTRA_HB = 48

def phy_cmd_for(uhb):
    dd = Driver()
    dd.cmd("MAC " + STA.hex())
    dd.cmd(f"CAPA {ULTRA_HB} {1 if uhb else 0}")
    ev = dd.cmd(f"CONNECT {AP.hex()} {CHANNEL} 0 OpenNet - -")
    dd.p.stdin.close()
    for cid, payload in ev['CMD']:
        if cid == PHY_CONTEXT_CMD:
            return payload
    return None

p_uhb = phy_cmd_for(True)
check(p_uhb is not None and len(p_uhb) == 32,
      f"ultra-high-band firmware gets a 32-byte PHY context ({len(p_uhb) if p_uhb else 0})")
if p_uhb and len(p_uhb) == 32:
    chan = struct.unpack_from('<I', p_uhb, 8)[0]
    band, width, ctrl = p_uhb[12], p_uhb[13], p_uhb[14]
    lmac = struct.unpack_from('<I', p_uhb, 16)[0]
    check(chan == CHANNEL, "UHB layout: channel is the 32-bit field at offset 8")
    check(band == 1, "UHB layout: band follows the channel (2.4 GHz = 1)")
    check(width == 0 and ctrl == 0, "UHB layout: 20 MHz, control position below")
    check(lmac == 0, "UHB layout: lmac_id lands at offset 16, not shifted")

p_v1 = phy_cmd_for(False)
check(p_v1 is not None and len(p_v1) == 28,
      f"older firmware still gets the 28-byte form ({len(p_v1) if p_v1 else 0})")
if p_v1 and len(p_v1) == 28:
    check(p_v1[8] == 1 and p_v1[9] == CHANNEL,
          "legacy layout: band then channel, both single bytes")
    check(struct.unpack_from('<I', p_v1, 12)[0] == 0,
          "legacy layout: lmac_id at offset 12")

# =============================================================================
# MLD command sizes and layout.
#
# This firmware implements the MLD API, so the connect path uses MAC_CONFIG /
# LINK_CONFIG / STA_CONFIG rather than the legacy MAC_CONTEXT / BINDING /
# ADD_STA.  Sizes and field offsets below are exactly what Linux puts on the
# wire against this same firmware (tests/wifi/decode_iwl_trace.py).
# =============================================================================

print("\n=== MLD command sizes ===")
MAC_CONFIG  = (3 << 8) | 0x08
LINK_CONFIG = (3 << 8) | 0x09
STA_CONFIG  = (3 << 8) | 0x0a
STA_REMOVE  = (3 << 8) | 0x0c
PHY_CONTEXT = 0x08

dd = Driver()
dd.cmd("MAC " + STA.hex())
dd.cmd(f"CAPA {ULTRA_HB} 1")
ev = dd.cmd(f"CONNECT {AP.hex()} {CHANNEL} 0 OpenNet - -")
sizes, first = {}, {}
for cid, payload in ev['CMD']:
    sizes.setdefault(cid, len(payload))
    first.setdefault(cid, payload)

check(sizes.get(MAC_CONFIG) == 52,  f"MAC_CONFIG_CMD is 52 bytes (got {sizes.get(MAC_CONFIG)})")
check(sizes.get(LINK_CONFIG) == 208, f"LINK_CONFIG_CMD is 208 bytes (got {sizes.get(LINK_CONFIG)})")
check(sizes.get(STA_CONFIG) == 96,  f"STA_CONFIG_CMD is 96 bytes (got {sizes.get(STA_CONFIG)})")
check(sizes.get(PHY_CONTEXT) == 32, f"PHY_CONTEXT_CMD is 32 bytes (got {sizes.get(PHY_CONTEXT)})")
check(0x28 not in sizes and 0x2b not in sizes and 0x18 not in sizes,
      "no legacy MAC_CONTEXT / BINDING / ADD_STA is sent")

m = first.get(MAC_CONFIG)
if m:
    check(struct.unpack_from('<I', m, 8)[0] == 5, "MAC_CONFIG: mac_type is BSS_STA")
    check(m[12:18] == STA, "MAC_CONFIG: local_mld_addr is our MAC")
    check(struct.unpack_from('<I', m, 20)[0] == 0x0c,
          "MAC_CONFIG: accepts group + beacon frames before association")
    check(m[36] == 0, "MAC_CONFIG: is_assoc clear on the initial add")

l = first.get(LINK_CONFIG)
if l:
    check(struct.unpack_from('<I', l, 0)[0] == 1, "LINK_CONFIG: first one is an ADD")
    check(struct.unpack_from('<I', l, 12)[0] == 0xffffffff,
          "LINK_CONFIG: phy_id is INVALID until the link is bound")
    check(l[16:22] == STA, "LINK_CONFIG: local_link_addr is our MAC")
    check(struct.unpack_from('<I', l, 136)[0] == 0,
          "LINK_CONFIG: no beacon timing before association")

# The PHY binding and the activation must be two separate MODIFYs, in that
# order, after the PHY context exists.  Folding them into one command asserts
# fw 89 (UMAC error 0x2010330F, seen on the AX211); Linux sends the binding
# first with mask 0 (__iwl_mvm_mld_assign_vif_chanctx).
links = [(i, p) for i, (c, p) in enumerate(ev['CMD']) if c == LINK_CONFIG]
phy_add = next((i for i, (c, p) in enumerate(ev['CMD'])
                if c == PHY_CONTEXT and struct.unpack_from('<I', p, 4)[0] == 1), None)
check(len(links) >= 3, f"link add, PHY binding and activation are three commands "
      f"(got {len(links)})")
if len(links) >= 3 and phy_add is not None:
    bind_i, bind = links[1]
    act_i, act = links[2]
    check(phy_add < bind_i, "the PHY context exists before the link binds to it")
    check(struct.unpack_from('<I', bind, 0)[0] == 2
          and struct.unpack_from('<I', bind, 12)[0] == 0
          and struct.unpack_from('<I', bind, 24)[0] == 0
          and struct.unpack_from('<I', bind, 28)[0] == 0,
          "PHY binding is its own MODIFY: phy_id set, mask 0, still inactive")
    check(struct.unpack_from('<I', act, 0)[0] == 2
          and struct.unpack_from('<I', act, 24)[0] == 0x03
          and struct.unpack_from('<I', act, 28)[0] == 1,
          "activation is a later MODIFY with exactly ACTIVE|RATES_INFO")
    # Field values the Linux trace carries at this stage (fw 89 asserted with
    # UMAC error 0x2010303E when they differed).
    check(struct.unpack_from('<I', act, 136)[0] == 100
          and struct.unpack_from('<I', act, 140)[0] == 0,
          "activation carries the beacon interval but no DTIM interval yet")
    check(struct.unpack_from('<I', act, 56)[0] == 0x2,
          "pre-assoc qos_flags is TGN without UPDATE_EDCA")
    check(struct.unpack_from('<H', act, 82)[0] == 0
          and struct.unpack_from('<H', act, 62)[0] == 1023,
          "pre-assoc EDCA is the default contention set, not WMM")
    check(struct.unpack_from('<H', act, 146)[0] == 0,
          "no RTS threshold before association")

st = first.get(STA_CONFIG)
if st:
    check(struct.unpack_from('<I', st, 0)[0] == 0, "STA_CONFIG: station id 0")
    check(st[8:14] == AP and st[16:22] == AP,
          "STA_CONFIG: both peer addresses are the BSSID")
    check(struct.unpack_from('<I', st, 28)[0] == 0,
          "STA_CONFIG: association id 0 before association")
dd.p.stdin.close()

# =============================================================================
# Beacon timing arrives with association, not before.
# =============================================================================

print("\n=== link beacon timing ===")
dd = Driver()
dd.cmd("MAC " + STA.hex())
dd.cmd(f"CAPA {ULTRA_HB} 1")
dd.cmd(f"CONNECT {AP.hex()} {CHANNEL} 0 OpenNet - -")
dd.cmd("RXMGMT " + mgmt(0xb0, STA, AP, AP, struct.pack('<HHH', 0, 2, 0)).hex())
dd.cmd("SERVICE")
dd.cmd("RXMGMT " + mgmt(0x10, STA, AP, AP,
                        struct.pack('<HHH', 0x0421, 0, 7 | 0xc000)).hex())
ev = dd.cmd("SERVICE")
link_i = next((i for i, (c, p) in enumerate(ev['CMD']) if c == LINK_CONFIG), None)
mac_i  = next((i for i, (c, p) in enumerate(ev['CMD']) if c == MAC_CONFIG), None)
link = ev['CMD'][link_i][1] if link_i is not None else None
mac  = ev['CMD'][mac_i][1] if mac_i is not None else None
check(link is not None and mac is not None,
      "association updates the link and the MAC")
check(link_i is not None and mac_i is not None and link_i < mac_i,
      "one-shot link timing is programmed before the MAC is marked associated")
if link:
    mask = struct.unpack_from('<I', link, 24)[0]
    check(struct.unpack_from('<I', link, 0)[0] == 2, "post-assoc link is a MODIFY")
    check(mask & 0x10, "post-assoc link sets the BEACON_TIMING modify bit")
    check(mask == 0x1e,
          f"legacy post-assoc link changes exactly rates/protection/QoS/timing (got {mask:#x})")
    check(not (mask & 0x01),
          "post-assoc link does not re-assert ACTIVE on an active link")
    check(struct.unpack_from('<I', link, 136)[0] == 100,
          "post-assoc link carries the beacon interval")
    check(struct.unpack_from('<I', link, 140)[0] == 200,
          "post-assoc link carries bi * dtim period")
    check(struct.unpack_from('<I', link, 56)[0] == 0x3,
          "post-assoc qos_flags is TGN plus UPDATE_EDCA")
    check(struct.unpack_from('<H', link, 82)[0] == 3008,
          "post-assoc EDCA switches to the WMM set")
if mac:
    check(mac[36] == 1, "post-assoc MAC has is_assoc set")
    check(struct.unpack_from('<H', mac, 40)[0] == 7, "post-assoc MAC carries the AID")
    check(struct.unpack_from('<I', mac, 20)[0] == 0x04,
          "post-assoc MAC stops asking for beacons")
check(not any(c == STA_CONFIG for c, p in ev['CMD']),
      "client AID is not incorrectly written to STA_CONFIG's GO-only assoc_id")
dd.p.stdin.close()

# =============================================================================
# Deadlines are not allowed to expire the moment they are set.
#
# The service loop sampled the clock once at the top of the pass, then spent
# real milliseconds inside the post-association context commands, each of which
# stamps a *newer* timestamp.  Comparing the stale `now` against those stamps
# underflowed the unsigned subtraction, so a join that had just succeeded
# reported "timed out while joining the network" straight after "Associated",
# tore the contexts down, and the access point's first EAPOL frame arrived to
# find the station already gone.
#
# The harness advances its clock on every host command (CMD_ROUND_TRIP_MS), so
# a pass that sends commands and then checks a deadline reproduces this.
# =============================================================================

print("\n=== deadlines survive the time spent sending commands ===")
d7 = Driver()
d7.cmd("MAC " + STA.hex())
d7.cmd(f"CONNECT {AP.hex()} {CHANNEL} 0 {SSID} {PASS} {AP_RSN.hex()}")
d7.cmd("RXMGMT " + mgmt(0xb0, STA, AP, AP, struct.pack('<HHH', 0, 2, 0)).hex())
ev = d7.cmd("SERVICE")
check(state_of(ev['result']) == 4,
      "sending the association request does not immediately expire its own retry timer")
check(len(ev['TX']) == 1,
      "the association request is sent once per pass, not retransmitted on the spot")

d7.cmd("RXMGMT " + mgmt(0x10, STA, AP, AP,
                        struct.pack('<HHH', 0x0421, 0, 7 | 0xc000)).hex())
ev = d7.cmd("SERVICE")
check(state_of(ev['result']) == 6,
      "association reaches the handshake instead of timing out in the same pass")
check(not any(x.startswith('TXQ-DOWN') for x in ev['TXQ']),
      "a successful association does not tear the transmit queue back down")
check(not any(c == STA_REMOVE for c, p in ev['CMD']),
      "a successful association does not remove the station it just added")

# The EAPOL exchange arrives a moment later, exactly as it does on the air.
m1_late = eapol(0x0002 | 0x0008 | 0x0080, (1).to_bytes(8, 'big'),
                bytes(range(32)), b'\x00' * 8, b'')
ev = d7.cmd("RXDATA " + data_from_ds(STA, AP, AP, 0x888e, m1_late).hex())
ev = d7.cmd("SERVICE")
check(len(ev['TX']) == 1,
      "an EAPOL message 1 that arrives after the context updates is still answered")
check(state_of(ev['result']) == 6,
      "the station is still in the handshake when message 1 is answered")

# =============================================================================
# Teardown order: the MAC gives up its association first.
#
# While MAC_CONFIG.is_assoc is set, the firmware's MAC context owns the link
# carrying the BSS.  Deactivating that link underneath it asserts the UMAC
# (0x2000320F on firmware 89) -- the failure the log above ends in.
# =============================================================================

print("\n=== teardown clears the association before the link ===")
ev = d7.cmd("ABORT")
seq = [(c, p) for c, p in ev['CMD']]


def index_of(pred):
    return next((i for i, (c, p) in enumerate(seq) if pred(c, p)), None)


deassoc = index_of(lambda c, p: c == MAC_CONFIG
                   and struct.unpack_from('<I', p, 4)[0] == 2 and p[36] == 0)
sta_rm = index_of(lambda c, p: c == STA_REMOVE)
deact = index_of(lambda c, p: c == LINK_CONFIG
                 and struct.unpack_from('<I', p, 0)[0] == 2
                 and struct.unpack_from('<I', p, 24)[0] & 0x01
                 and struct.unpack_from('<I', p, 28)[0] == 0)
link_rm = index_of(lambda c, p: c == LINK_CONFIG
                   and struct.unpack_from('<I', p, 0)[0] == 3)
mac_rm = index_of(lambda c, p: c == MAC_CONFIG
                  and struct.unpack_from('<I', p, 4)[0] == 3)

check(deassoc is not None,
      "teardown sends a MAC_CONFIG MODIFY clearing is_assoc")
check(deact is not None,
      "teardown deactivates the link before removing it")
check(deassoc is not None and deact is not None and deassoc < deact,
      "the association is cleared before the link is deactivated")
check(deassoc is not None and sta_rm is not None and deassoc < sta_rm,
      "the association is cleared before the station is removed")
check(sta_rm is not None and deact is not None and sta_rm < deact,
      "the station is removed before the link it sits on is deactivated")
check(deact is not None and link_rm is not None and deact < link_rm,
      "the link is deactivated before it is removed")
check(link_rm is not None and mac_rm is not None and link_rm < mac_rm,
      "the link is removed before the MAC that owns it")
d7.p.stdin.close()

# =============================================================================
# An access point that stops acknowledging brings the link down.
#
# IwxLinkUp() used to be nothing but "the state machine reached Connected", and
# only an explicit deauthentication frame moved it off that state.  A hotspot
# that simply went away -- slept, changed channel, dropped the station without
# saying so -- therefore left the link reported as up forever: NetIf kept
# choosing wlan0, every packet vanished, and the desktop showed a healthy
# connection while nothing worked.  Beacons cannot be used to notice this
# (MacConfigCmd stops asking for them once associated), so the driver watches
# its own frames going unacknowledged instead.
# =============================================================================

print("\n=== a silent access point takes the link down ===")
d8 = Driver()
d8.cmd("MAC " + STA.hex())
d8.cmd(f"CONNECT {AP.hex()} {CHANNEL} 0 OpenNet - -")
d8.cmd("RXMGMT " + mgmt(0xb0, STA, AP, AP, struct.pack('<HHH', 0, 2, 0)).hex())
d8.cmd("SERVICE")
d8.cmd("RXMGMT " + mgmt(0x10, STA, AP, AP,
                        struct.pack('<HHH', 0x0421, 0, 7 | 0xc000)).hex())
ev = d8.cmd("SERVICE")
check(link_of(ev['result']) == 1, "the open network is associated and the link is up")

# Well short of the threshold: a few unacknowledged frames are ordinary.
ev = d8.cmd("TXSTATUS 0 15")
check(link_of(ev['result']) == 1,
      "a handful of unacknowledged frames does not drop the link")
ev = d8.cmd("SERVICE")
check(link_of(ev['result']) == 1, "and the service pass leaves it alone")

# One acknowledgement means the access point is still there; the count restarts.
d8.cmd("TXSTATUS 1 1")
ev = d8.cmd("TXSTATUS 0 15")
check(link_of(ev['result']) == 1,
      "an acknowledgement in between restarts the count")

# Anything received from the BSS is equally good proof, and also restarts it.
d8.cmd("RXDATA " + data_from_ds(STA, AP, AP, 0x0800, b'\x45' * 20).hex())
ev = d8.cmd("TXSTATUS 0 15")
check(link_of(ev['result']) == 1,
      "a frame received from the BSS restarts the count too")

# Now let it run past the threshold with nothing coming back.
ev = d8.cmd("TXSTATUS 0 16")
check(link_of(ev['result']) == 1,
      "the transmit path itself does not tear anything down")
check(not ev['CMD'],
      "no firmware command is sent from the completion path")

ev = d8.cmd("SERVICE")
check(state_of(ev['result']) == 0 and link_of(ev['result']) == 0,
      "the next service pass drops the link")
check(any(c == MAC_CONFIG and struct.unpack_from('<I', p, 4)[0] == 3
          for c, p in ev['CMD']),
      "and unwinds the firmware contexts")
d8.p.stdin.close()

print()
if fails:
    for f in fails: print("FAILURE:", f)
    sys.exit(1)
print("All MLME and data-path cases passed.")