#!/usr/bin/env python3
"""Second-round supplicant tests: GTK rekey, message-3 retransmission, and RSN
negotiation against realistic access-point information elements."""
import hashlib, hmac, os, subprocess, sys, struct
from cryptography.hazmat.primitives.keywrap import aes_key_wrap
HARNESS = sys.argv[1]
fails = []

def prf(key, label, data, n):
    r=b''; i=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, aa, spa, an, sn):
    return prf(pmk, "Pairwise key expansion", min(aa,spa)+max(aa,spa)+min(an,sn)+max(an,sn), 48)
def mic(kck, f): return hmac.new(kck, f, hashlib.sha1).digest()[:16]
HDR=99
def build(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]=mic(kck, bytes(b))
    return bytes(b)
def check_mic(f,kck):
    x=bytearray(f); x[81:97]=b'\x00'*16
    return mic(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

def start(p, ssid, passphrase, spa, aa):
    p.stdin.write(f"START {spa.hex()} {aa.hex()} {ssid} {passphrase} 2 4 4\n"); p.stdin.flush()
    p.stdout.readline(); p.stdout.readline()

def rd(p): return p.stdout.readline().strip()

# ---------------------------------------------------------------- handshake +
print("=== GTK rekey and message-3 retransmission ===")
aa=bytes.fromhex('001122334455'); spa=bytes.fromhex('aabbccddeeff')
ssid="MontaukTest"; pw="supersecret123"
p=subprocess.Popen([HARNESS],stdin=subprocess.PIPE,stdout=subprocess.PIPE,text=True,bufsize=1)
start(p, ssid, pw, spa, aa)
pmk=hashlib.pbkdf2_hmac('sha1',pw.encode(),ssid.encode(),4096,32)
an=os.urandom(32)
p.stdin.write("RX "+build(0x0002|0x0008|0x0080,(1).to_bytes(8,'big'),an,b'\x00'*8,b'').hex()+"\n"); p.stdin.flush()
m2=bytes.fromhex(rd(p).split()[1]); rd(p)
sn=m2[17:49]
ptk=derive_ptk(pmk,aa,spa,an,sn); kck,kek,tk=ptk[:16],ptk[16:32],ptk[32:48]
gtk1=os.urandom(16)
m3=build(0x0002|0x0008|0x0040|0x0080|0x0100|0x0200|0x1000,(2).to_bytes(8,'big'),an,
         b'\x00'*8, aes_key_wrap(kek, gtk_kde(gtk1,1)), kck)
p.stdin.write("RX "+m3.hex()+"\n"); p.stdin.flush()
got={}
while True:
    l=rd(p)
    if l.startswith('RX-OK'): break
    k,_,v=l.partition(' '); got[k]=v
assert got['GTK']==gtk1.hex() and got['PTK']==tk.hex()
print(f"  initial handshake complete, state {l}")

# message 3 retransmitted (AP missed message 4)
m3b=build(0x0002|0x0008|0x0040|0x0080|0x0100|0x0200|0x1000,(3).to_bytes(8,'big'),an,
          b'\x00'*8, aes_key_wrap(kek, gtk_kde(gtk1,1)), kck)
p.stdin.write("RX "+m3b.hex()+"\n"); p.stdin.flush()
got2={}
while True:
    l=rd(p)
    if l.startswith('RX-OK'): break
    k,_,v=l.partition(' '); got2[k]=v
if 'TX' not in got2:
    fails.append("no message 4 in response to a retransmitted message 3")
else:
    m4=bytes.fromhex(got2['TX'])
    ok = check_mic(m4,kck) and m4[9:17]==(3).to_bytes(8,'big')
    print(f"  answered retransmitted message 3 with a fresh message 4: {ok}")
    if not ok: fails.append("message 4 replay answer malformed")

# group key rekey (2-way, no pairwise bit)
gtk2=os.urandom(16)
gk=build(0x0002|0x0080|0x0100|0x0200|0x1000,(4).to_bytes(8,'big'),b'\x00'*32,
         bytes.fromhex('0a0b0c000000 0000'.replace(' ','')), aes_key_wrap(kek, gtk_kde(gtk2,2)), kck)
p.stdin.write("RX "+gk.hex()+"\n"); p.stdin.flush()
got3={}
while True:
    l=rd(p)
    if l.startswith('RX-OK'): break
    k,_,v=l.partition(' '); got3[k]=v
if got3.get('GTK')!=gtk2.hex():
    fails.append(f"group rekey installed {got3.get('GTK')} instead of {gtk2.hex()}")
else:
    print(f"  group rekey installed the new GTK (index {got3['GTK-IDX']})")
if 'TX' not in got3:
    fails.append("no acknowledgement to the group key message")
else:
    m=bytes.fromhex(got3['TX'])
    ki=struct.unpack_from('>H',m,5)[0]
    ok = bool(check_mic(m, kck) and not (ki & 0x0008) and (ki & 0x0200))
    print(f"  group key acknowledged correctly (no pairwise bit, secure set): {ok}")
    if not ok: fails.append("group key acknowledgement malformed")
p.stdin.close()

# ---------------------------------------------------------------- RSN parsing
print("\n=== RSN negotiation against real-world information elements ===")
def suite(t): return bytes([0x00,0x0f,0xac,t])
def ie(group, pairwise, akms, caps=0):
    b = struct.pack('<H',1) + suite(group)
    b += struct.pack('<H',len(pairwise)) + b''.join(suite(x) for x in pairwise)
    b += struct.pack('<H',len(akms)) + b''.join(suite(x) for x in akms)
    b += struct.pack('<H',caps)
    return b

cases = [
    ("WPA2-PSK, CCMP only",              ie(4,[4],[2]),            True, 2, 4),
    ("WPA2 mixed TKIP+CCMP (TKIP group)",ie(2,[4,2],[2]),          False,0, 0),
    ("WPA2 mixed pairwise, CCMP group",  ie(4,[4,2],[2]),          True, 2, 4),
    ("WPA2/WPA3 transition (PSK+SAE)",   ie(4,[4],[2,8],0x0080),   True, 2, 4),
    ("WPA3-only SAE, MFP required",      ie(4,[4],[8],0x00c0),     False,0, 0),
    ("TKIP-only",                        ie(2,[2],[2]),            False,0, 0),
    ("enterprise 802.1X",                ie(4,[4],[1]),            False,0, 0),
    ("PSK-SHA256 only",                  ie(4,[4],[6]),            True, 6, 4),
    ("GCMP-256",                         ie(9,[9],[2]),            True, 2, 9),
]
p=subprocess.Popen([HARNESS],stdin=subprocess.PIPE,stdout=subprocess.PIPE,text=True,bufsize=1)
for name, body, expect_ok, expect_akm, expect_pc in cases:
    p.stdin.write("PARSE "+body.hex()+"\n"); p.stdin.flush()
    r=rd(p).split()
    ok = r[1]=='1'; akm=int(r[3]); pc=int(r[5])
    good = (ok==expect_ok) and (not expect_ok or (akm==expect_akm and pc==expect_pc))
    print(f"  {'PASS' if good else 'FAIL'}  {name:36} -> accepted={ok} akm={akm} pairwise={pc}")
    if not good: fails.append(f"RSN parse: {name}")
p.stdin.close()

print()
if fails:
    for f in fails: print("FAILURE:", f)
    sys.exit(1)
print("All second-round supplicant cases passed.")