#!/usr/bin/env python3
"""Independent WPA2 authenticator driving the kernel supplicant through a
4-way handshake.  All crypto here comes from hashlib / the `cryptography`
package, so it is a genuine oracle for the C++ implementation."""
import hashlib, hmac, os, subprocess, sys, struct
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.keywrap import aes_key_wrap, aes_key_unwrap

HARNESS = sys.argv[1]

def prf(key, label, data, nbytes):
    r = b''
    i = 0
    while len(r) < nbytes:
        r += hmac.new(key, label.encode() + b'\x00' + data + bytes([i]), hashlib.sha1).digest()
        i += 1
    return r[:nbytes]

def kdf_sha256(key, label, data, nbytes):
    r = b''
    i = 1
    bits = nbytes * 8
    while len(r) < nbytes:
        r += hmac.new(key, struct.pack('<H', i) + label.encode() + data + struct.pack('<H', bits),
                      hashlib.sha256).digest()
        i += 1
    return r[:nbytes]

def derive_ptk(pmk, aa, spa, anonce, snonce, nbytes, sha256=False):
    data = min(aa, spa) + max(aa, spa) + min(anonce, snonce) + max(anonce, snonce)
    f = kdf_sha256 if sha256 else prf
    return f(pmk, "Pairwise key expansion", data, nbytes)

def mic(kck, frame, ver):
    if ver == 3:
        from cryptography.hazmat.primitives.cmac import CMAC
        c = CMAC(algorithms.AES(kck)); c.update(frame); return c.finalize()[:16]
    return hmac.new(kck, frame, hashlib.sha1).digest()[:16]

HDR = 99
def build(ver, desc, key_info, replay, nonce, rsc, key_data, kck=None, mic_ver=2):
    body = bytearray(HDR + len(key_data))
    body[0] = ver
    body[1] = 3                                   # EAPOL-Key
    struct.pack_into('>H', body, 2, HDR + len(key_data) - 4)
    body[4] = desc
    struct.pack_into('>H', body, 5, key_info)
    struct.pack_into('>H', body, 7, 16)           # key length (CCMP)
    body[9:17] = replay
    body[17:49] = nonce
    body[65:73] = rsc
    struct.pack_into('>H', body, 97, len(key_data))
    body[HDR:] = key_data
    if key_info & 0x0100:                         # MIC bit
        body[81:97] = mic(kck, bytes(body), mic_ver)
    return bytes(body)

def parse(frame):
    return {
        'key_info': struct.unpack_from('>H', frame, 5)[0],
        'replay': frame[9:17],
        'nonce': frame[17:49],
        'mic': frame[81:97],
        'kdlen': struct.unpack_from('>H', frame, 97)[0],
        'kd': frame[HDR:],
    }

def check_mic(frame, kck, ver):
    f = bytearray(frame); f[81:97] = b'\x00' * 16
    return mic(kck, bytes(f), ver) == frame[81:97]

def run_case(name, ssid, passphrase, akm, pcipher, gcipher, desc_ver, corrupt_pass=False):
    print(f"\n=== {name} ===")
    aa  = bytes.fromhex('001122334455')          # AP
    spa = bytes.fromhex('aabbccddeeff')          # station
    p = subprocess.Popen([HARNESS], stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True, bufsize=1)

    def send(line):
        p.stdin.write(line + '\n'); p.stdin.flush()
    def readline():
        return p.stdout.readline().strip()

    sta_pass = passphrase + ('X' if corrupt_pass else '')
    send(f"START {spa.hex()} {aa.hex()} {ssid} {sta_pass} {akm} {pcipher} {gcipher}")
    ok = readline()
    rsnie = readline()
    assert ok.startswith('START-OK 1'), ok
    sta_rsn = bytes.fromhex(rsnie.split()[1])
    print(f"  supplicant RSN IE: {sta_rsn.hex()}")

    pmk = hashlib.pbkdf2_hmac('sha1', passphrase.encode(), ssid.encode(), 4096, 32)
    anonce = os.urandom(32)
    replay = (1).to_bytes(8, 'big')

    # --- message 1 ---
    m1 = build(2, 2, desc_ver | 0x0008 | 0x0080, replay, anonce, b'\x00'*8, b'')
    send("RX " + m1.hex())
    tx = readline()
    assert tx.startswith('TX '), f"expected msg2, got {tx}"
    m2 = bytes.fromhex(tx.split()[1])
    status = readline()
    print(f"  msg2 {len(m2)} bytes, {status}")

    i2 = parse(m2)
    snonce = i2['nonce']
    ptk = derive_ptk(pmk, aa, spa, anonce, snonce, 48, sha256=(akm == 6))
    kck, kek, tk = ptk[:16], ptk[16:32], ptk[32:48]

    m2_ok = check_mic(m2, kck, desc_ver)
    print(f"  msg2 MIC verifies against independently derived PTK: {m2_ok}")
    if corrupt_pass:
        assert not m2_ok, "a wrong passphrase must not produce a valid MIC"
        print("  (expected: wrong passphrase -> MIC mismatch)")
        p.stdin.close(); return
    assert m2_ok
    assert i2['replay'] == replay, "msg2 must echo the replay counter"
    assert i2['kd'] == sta_rsn, "msg2 key data must be the supplicant's RSN IE"
    assert i2['key_info'] & 0x0008, "msg2 must set the pairwise bit"
    assert i2['key_info'] & 0x0100, "msg2 must set the MIC bit"

    # --- message 3, carrying an AES-wrapped GTK KDE ---
    gtk = os.urandom(16)
    gtk_kde = bytes([0xdd, 6 + len(gtk), 0x00, 0x0f, 0xac, 0x01, 0x02, 0x00]) + gtk
    kd = gtk_kde
    if len(kd) % 8:
        kd += b'\xdd' + b'\x00' * (7 - len(kd) % 8)
    wrapped = aes_key_wrap(kek, kd)
    replay3 = (2).to_bytes(8, 'big')
    rsc = bytes.fromhex('0102030405060000')
    m3 = build(2, 2, desc_ver | 0x0008 | 0x0040 | 0x0080 | 0x0100 | 0x0200 | 0x1000,
               replay3, anonce, rsc, wrapped, kck, desc_ver)
    send("RX " + m3.hex())

    got = {}
    while True:
        l = readline()
        if l.startswith('RX-OK'):
            print(f"  {l}")
            break
        k, _, v = l.partition(' ')
        got[k] = v
        if k == 'TX':
            m4 = bytes.fromhex(v)

    assert 'TX' in got, "supplicant must answer message 3 with message 4"
    m4_ok = check_mic(m4, kck, desc_ver)
    i4 = parse(m4)
    print(f"  msg4 MIC verifies: {m4_ok}")
    assert m4_ok
    assert i4['replay'] == replay3, "msg4 must echo message 3's replay counter"
    assert i4['key_info'] & 0x0200, "msg4 must set the secure bit"
    assert i4['kdlen'] == 0, "msg4 must carry no key data"

    assert got['PTK'] == tk.hex(), f"installed TK {got['PTK']} != expected {tk.hex()}"
    print(f"  installed TK matches the AP's derivation: True")
    assert got['GTK'] == gtk.hex(), f"installed GTK {got['GTK']} != {gtk.hex()}"
    print(f"  installed GTK matches: True  (key index {got['GTK-IDX']})")
    assert got['GTK-RSC'].startswith('010203040506')
    assert 'STATE 3' in l, f"supplicant should be Complete, got {l}"
    print("  PASS")
    p.stdin.close()

run_case("WPA2-PSK / CCMP / key descriptor v2", "MontaukTest", "supersecret123", 2, 4, 4, 2)
run_case("WPA2-PSK-SHA256 / CCMP / key descriptor v3", "MontaukTest", "supersecret123", 6, 4, 4, 3)
run_case("wrong passphrase is rejected", "MontaukTest", "supersecret123", 2, 4, 4, 2, corrupt_pass=True)
print("\nAll supplicant cases passed.")