import utime
from machine import I2C, Pin
import gc

# ── LCD DRIVER ───────────────────────────────────────────────
class LCD_I2C:
    def __init__(self, i2c, addr=0x27):
        self.i2c = i2c
        self.addr = addr
        self._init_lcd()

    def _write_byte(self, data):
        self.i2c.writeto(self.addr, bytes([data]))
        utime.sleep_us(100)

    def _pulse(self, data):
        self._write_byte(data | 0x04)
        utime.sleep_us(500)
        self._write_byte(data & ~0x04)
        utime.sleep_us(100)

    def _write4(self, data):
        self._write_byte(data | 0x08)
        self._pulse(data | 0x08)

    def _send(self, val, mode):
        self._write4((val & 0xF0) | mode | 0x08)
        self._write4(((val << 4) & 0xF0) | mode | 0x08)

    def _init_lcd(self):
        utime.sleep_ms(50)
        for _ in range(3):
            self._write4(0x30)
            utime.sleep_ms(5)
        self._write4(0x20)
        self._send(0x28, 0)
        self._send(0x0C, 0)
        self._send(0x01, 0)
        utime.sleep_ms(2)
        self._send(0x06, 0)

    def clear(self):
        self._send(0x01, 0)
        utime.sleep_ms(2)

    def cursor(self, col, row):
        self._send(0x80 | (col + [0x00, 0x40][row]), 0)

    def write(self, text):
        for ch in text:
            self._send(ord(ch), 1)

    def show(self, l1, l2=""):
        self.clear()
        self.cursor(0, 0)
        self.write(l1[:16])
        if l2:
            self.cursor(0, 1)
            self.write(l2[:16])
        utime.sleep_ms(1800)


# ── HARDWARE SETUP ───────────────────────────────────────────
i2c       = I2C(0, sda=Pin(0), scl=Pin(1), freq=100000)
lcd       = LCD_I2C(i2c)
green_led = Pin(15, Pin.OUT)
red_led   = Pin(14, Pin.OUT)
green_led.off()
red_led.off()


# ═══════════════════════════════════════════════════════════
#  CRYPTOGRAPHIC PRIMITIVES
# ═══════════════════════════════════════════════════════════

P = 251  # Prime for modular arithmetic (paper uses prime p)

# ── PUF SIMULATION ──────────────────────────────────────────
# Paper Sec 6.4: PUF simulated via random number generator
# Same challenge always gives same response (reproducibility)
def PUF(challenge, device_seed=0x1A2B):
    s = ((challenge * device_seed) ^ 0xDEAD) & 0xFFFF
    s = (s * 6364136223846793005 + 1442695040888963407) & 0xFFFF
    s = (s * 6364136223846793005 + 1442695040888963407) & 0xFFFF
    return (s % (P - 1)) + 1

def PUF_B(challenge):
    return PUF(challenge, device_seed=0x3C4D)


# ── HASH FUNCTION ────────────────────────────────────────────
def HASH(a, b=None):
    if b is None:
        b = a ^ 0xABCD
    h = 5381
    for val in [a, b]:
        for _ in range(4):
            h = ((h << 5) + h + (val & 0xFF)) & 0xFFFFFFFF
            val >>= 8
    return h % 65521


# ── XOR ──────────────────────────────────────────────────────
def XOR(a, b):
    return (a ^ b) & 0xFFFF


# ── SECRET SHARING ───────────────────────────────────────────
# Paper Sec 3.2, Eq (1): s_dev = (S + Rand) mod P
#                         s_ver = (S + 2*Rand) mod P
# Eq (2): S = (2*s_dev - s_ver) mod P
def split(secret, rand_val):
    s_dev = (secret + rand_val)     % P
    s_ver = (secret + 2 * rand_val) % P
    return s_dev, s_ver

def reconstruct(s_dev, s_ver):
    return (2 * s_dev - s_ver) % P


# ── NONCE GENERATOR ──────────────────────────────────────────
_nc = 0
def nonce():
    global _nc
    _nc = (_nc + 1) & 0xFF
    t = utime.ticks_us() & 0x1FFF
    return ((t * 97 + _nc * 131 + 7) % 9000) + 1000


# ── PRINT HELPER ─────────────────────────────────────────────
def sep(title):
    print("\n" + "=" * 46)
    print("  " + title)
    print("=" * 46)



# ── PERFORMANCE METRICS ─────────────────────────────
results = {
    "enroll_time": [],
    "dd_time": [],
    "dv_time": [],
    "auth_success": 0,
    "auth_fail": 0,
    "attack_detected": 0,
    "session_match": 0,
    "memory": []
}

# ═══════════════════════════════════════════════════════════
#  PROTOCOL STATE
# ═══════════════════════════════════════════════════════════

# Device A: paper Step EN4 stores <IDA, sA, hcA, CAV>
devA = {"IDA": 0, "sA": 0, "hcA": 0, "CAV": 0}

# Device B: same structure, different PUF seed
devB = {"IDA": 0, "sA": 0, "hcA": 0, "CAV": 0}

# Verifier DB: paper Step EN3 stores <IDA, sAV, hrA, KAV>
# No raw CRP — prevents ML attacks
db_A = {"IDA": 0, "sAV": 0, "hrA": 0, "KAV": 0}
db_B = {"IDA": 0, "sAV": 0, "hrA": 0, "KAV": 0}

IDV = 999  # Verifier identity


# ═══════════════════════════════════════════════════════════
#  PHASE 1: ENROLLMENT  (Paper Sec 4.3, Steps EN1→EN4)
# ═══════════════════════════════════════════════════════════
def enroll(dev_dict, db_dict, puf_fn, label="A"):
    start=utime.ticks_us()
    gc.collect()
    mem_before=gc.mem_free()
    sep("ENROLLMENT: Device {}".format(label))
    lcd.show("Enrolling Dev{}".format(label), "Please wait...")

    # EN1: Verifier generates two random challenges
    C   = (nonce() % (P - 2)) + 1
    CAV = (nonce() % (P - 2)) + 1
    print("[EN1] C_{}={}, CAV_{}={}".format(label, C, label, CAV))

    # EN2: Device applies PUF to both challenges
    R   = puf_fn(C)
    RAV = puf_fn(CAV)
    print("[EN2] R_{}={}, RAV_{}={}".format(label, R, label, RAV))

    # EN3: Verifier splits C, stores only hashes (no raw CRP)
    Rand       = (nonce() % (P - 2)) + 1
    s_dev, s_ver = split(C, Rand)
    hrA        = HASH(R)      # H(R) — not R
    hcA        = HASH(C)      # H(C) — for later verification
    KAV        = RAV           # bootstrap symmetric key
    IDA        = ((C * 13 + R * 7) % 800) + 100

    db_dict["IDA"] = IDA
    db_dict["sAV"] = s_ver
    db_dict["hrA"] = hrA
    db_dict["KAV"] = KAV
    print("[EN3] DB: IDA={} sAV={} hrA={} KAV={}".format(IDA, s_ver, hrA, KAV))
    print("      --> ML attack PREVENTED (no raw CRP)")

    # EN4: Device stores its share
    dev_dict["IDA"] = IDA
    dev_dict["sA"]  = s_dev
    dev_dict["hcA"] = hcA
    dev_dict["CAV"] = CAV
    print("[EN4] Device: IDA={} sA={} hcA={} CAV={}".format(IDA, s_dev, hcA, CAV))
    lcd.show("Dev{} Enrolled".format(label), "IDA={}".format(IDA))
    end=utime.ticks_us()
    gc.collect()
    mem_after=gc.mem_free()
    exec_time=utime.ticks_diff(end,start)
    mem_used=mem_before-mem_after
    results["enroll_time"].append(exec_time)
    results["memory"].append(mem_used)
    print("Enrollment Time:",exec_time,"us")
    print("Memory Used:",mem_used,"bytes")
    return True


# ═══════════════════════════════════════════════════════════
#  PHASE 2: DEVICE-DEVICE AKE  (Paper Sec 4.4.1, DD1→DD4)
# ═══════════════════════════════════════════════════════════
def dd_ake():
    start=utime.ticks_us()
    sep("DD-AKE: Device A <-> Device B")
    lcd.show("DD-AKE Phase", "Dev A <-> Dev B")

    # ── DD1: Device A → Verifier ─────────────────────────
    # KAV = PUF_A(CAV_A)
    # hkAV = H(KAV || NA)
    # TDA  = H(IDA || hkAV),  TDB = H(IDB || hkAV)
    # TDV  = H(IDV || hkAV)
    # SGAV = H(TDA || TDV || TDB || NA || KAV)
    print("\n[DD1] Device A builds auth message...")

    KAV_A  = PUF(devA["CAV"])
    NA     = nonce()
    hkAV_A = HASH(KAV_A, NA)
    TDA    = HASH(devA["IDA"], hkAV_A)
    TDB    = HASH(devB["IDA"], hkAV_A)
    TDV    = HASH(IDV, hkAV_A)
    SGAV   = HASH(HASH(TDA, TDV), HASH(HASH(TDB, NA), KAV_A))

    print("  KAV_A={} NA={} TDA={} SGAV={}".format(KAV_A, NA, TDA, SGAV))
    lcd.show("DD1: A->Verif", "Auth request")

    # ── DD2: Verifier verifies A, prepares shares for A & B
    # Verify: recompute SGAV using KAV from DB and received NA
    # Then generate NV, encrypt shares DA and DB
    # Rp = H(hrA || NV) XOR H(hrB || NV)   (blinded response pair)
    # SGVA, SGVB = signatures for A and B
    print("\n[DD2] Verifier verifying Device A...")

    KAV_V  = db_A["KAV"]
    hkAV_V = HASH(KAV_V, NA)        # use received NA for verification

    TDA_chk = HASH(db_A["IDA"], hkAV_V)
    TDB_chk = HASH(db_B["IDA"], hkAV_V)
    TDV_chk = HASH(IDV, hkAV_V)
    SGAV_chk = HASH(
        HASH(TDA_chk, TDV_chk),
        HASH(HASH(TDB_chk, NA), KAV_V)
    )

    if SGAV_chk != SGAV:
        print("  [FAIL] Device A signature invalid!")
        lcd.show("DD2 FAIL!", "A sig invalid")
        return False
    print("  [OK] Device A authenticated to Verifier")

    KBV_V   = db_B["KAV"]
    NV      = nonce()

    hkAV_NV = HASH(KAV_V, NV)       # temporal key for A using NV
    hkBV_NV = HASH(KBV_V, NV)       # temporal key for B using NV

    TDA_NV  = HASH(db_A["IDA"], hkAV_NV)
    TDV_ANV = HASH(IDV, hkAV_NV)
    TDB_NV  = HASH(db_B["IDA"], hkBV_NV)
    TDV_BNV = HASH(IDV, hkBV_NV)

    # Blinded combined response: H(hrA||NV) XOR H(hrB||NV)
    Rp = XOR(HASH(db_A["hrA"], NV), HASH(db_B["hrA"], NV))

    # Encrypted shares
    DA = XOR(db_A["sAV"], hkAV_NV)
    DB = XOR(db_B["sAV"], hkBV_NV)

    # Signatures
    SGVA = HASH(
        HASH(TDV_ANV, TDA_NV),
        HASH(HASH(DA, Rp), HASH(NV, KAV_V))
    )
    SGVB = HASH(
        HASH(TDV_BNV, TDB_NV),
        HASH(HASH(DB, Rp), HASH(NV, KBV_V))
    )

    print("  NV={} Rp={} DA={} DB={} SGVA={} SGVB={}".format(
        NV, Rp, DA, DB, SGVA, SGVB))
    lcd.show("DD2: Shares", "Sent A and B")

    # ── DD3: Device A → verifies verifier, computes KS ───
    # sAV  = DA XOR hkAV(NV)
    # CA   = reconstruct(sA, sAV)         — Eq (2)
    # verify H(CA) == hcA
    # RA   = PUF_A(CA)                    — controlled PUF
    # RESA = H(H(RA) || NV)
    # RESB = Rp XOR RESA
    # KS   = H(RESA || RESB)
    print("\n[DD3] Device A verifies Verifier and computes session key...")

    hkAV_A_NV = HASH(KAV_A, NV)
    TDA_A_NV  = HASH(devA["IDA"], hkAV_A_NV)
    TDV_A_NV  = HASH(IDV, hkAV_A_NV)
    SGVA_chk  = HASH(
        HASH(TDV_A_NV, TDA_A_NV),
        HASH(HASH(DA, Rp), HASH(NV, KAV_A))
    )

    if SGVA_chk != SGVA:
        print("  [FAIL] Verifier signature invalid for Device A!")
        lcd.show("DD3 FAIL!", "Verif sig bad")
        return False
    print("  [OK] Verifier authenticated to Device A")

    sAV_A = XOR(DA, hkAV_A_NV)
    CA    = reconstruct(devA["sA"], sAV_A)

    if HASH(CA) != devA["hcA"]:
        print("  [FAIL] Challenge hash mismatch in Device A!")
        lcd.show("DD3 FAIL!", "CA hash error")
        return False
    print("  CA={} [verified]".format(CA))

    # PUF runs only after verifier is authenticated (controlled PUF)
    RA   = PUF(CA)
    RESA = HASH(HASH(RA), NV)
    RESB = XOR(Rp, RESA)
    KS_A = HASH(RESA, RESB)

    print("  RA={} RESA={} KS_A={}".format(RA, RESA, KS_A))
    lcd.show("DD3: Dev A OK", "KS={}".format(KS_A % 9999))

    # ── DD4: Device B → verifies verifier, computes KS ───
    # Mirror of DD3 but uses PUF_B
    print("\n[DD4] Device B verifies Verifier and computes session key...")

    KBV_B     = PUF_B(devB["CAV"])
    hkBV_B_NV = HASH(KBV_B, NV)
    TDB_B_NV  = HASH(devB["IDA"], hkBV_B_NV)
    TDV_B_NV  = HASH(IDV, hkBV_B_NV)
    SGVB_chk  = HASH(
        HASH(TDV_B_NV, TDB_B_NV),
        HASH(HASH(DB, Rp), HASH(NV, KBV_B))
    )

    if SGVB_chk != SGVB:
        print("  [FAIL] Verifier signature invalid for Device B!")
        lcd.show("DD4 FAIL!", "Verif sig bad")
        return False
    print("  [OK] Verifier authenticated to Device B")

    sBV_B = XOR(DB, hkBV_B_NV)
    CB    = reconstruct(devB["sA"], sBV_B)

    if HASH(CB) != devB["hcA"]:
        print("  [FAIL] Challenge hash mismatch in Device B!")
        lcd.show("DD4 FAIL!", "CB hash error")
        return False
    print("  CB={} [verified]".format(CB))

    RB     = PUF_B(CB)
    RESB_B = HASH(HASH(RB), NV)
    RESA_B = XOR(Rp, RESB_B)
    KS_B   = HASH(RESA_B, RESB_B)

    print("  RB={} RESB_B={} KS_B={}".format(RB, RESB_B, KS_B))

    if KS_A == KS_B:
        end=utime.ticks_us()
        exec_time=utime.ticks_diff(end,start)
        results["dd_time"].append(exec_time)
        results["auth_success"]+=1
        results["session_match"]+=1
        print("DD-AKE Time:",exec_time,"us")
        print("\n  [SUCCESS] DD-AKE complete — KS = {}".format(KS_A))
        lcd.show("DD-AKE SUCCESS", "KS Match!")
        return True
    else:
        print("  [FAIL] Session key mismatch A={} B={}".format(KS_A, KS_B))
        lcd.show("DD-AKE FAIL!", "KS mismatch")
        return False


# ═══════════════════════════════════════════════════════════
#  PHASE 3: DEVICE-VERIFIER AKE  (Paper Sec 4.4.2, DS1→DS4)
# ═══════════════════════════════════════════════════════════
def dv_ake(inject_attack=False):
    start=utime.ticks_us()
    sep("DV-AKE: Device A <-> Verifier")
    if inject_attack:
        print("  *** ATTACK SIMULATION: Replay/tamper attack ***")
        lcd.show("DV-AKE ATTACK", "Simulating...")
    else:
        lcd.show("DV-AKE Phase", "Dev A <-> Verif")

    # ── DS1: Device A → Verifier ─────────────────────────
    # KAV  = PUF_A(CAV)
    # hkAV = H(KAV || NA)
    # TDA  = H(IDA || hkAV),  TDV = H(IDV || hkAV)
    # SGAV = H(TDA || TDV || NA || KAV)
    print("\n[DS1] Device A building auth message...")

    KAV_dev  = PUF(devA["CAV"])
    NA       = nonce()
    hkAV_dev = HASH(KAV_dev, NA)
    TDA      = HASH(devA["IDA"], hkAV_dev)
    TDV      = HASH(IDV, hkAV_dev)
    SGAV     = HASH(HASH(TDA, TDV), HASH(NA, KAV_dev))

    NA_sent  = (NA + 9999) if inject_attack else NA
    if inject_attack:
        print("  [ATTACKER] Tampered NA: {} -> {}".format(NA, NA_sent))

    print("  KAV={} NA={} TDA={} SGAV={}".format(KAV_dev, NA, TDA, SGAV))
    lcd.show("DS1: A->Verif", "Auth message")

    # ── DS2: Verifier verifies Device A ──────────────────
    # Verify SGAV, then generate CAnew, split into shares
    # Encrypt old sAV, new sAnew, and CAnew for device
    # Cl  = CAnew XOR hkAV(NV) XOR hrA   (challenge concealed)
    # PAV = sAV  XOR hkAV(NV)
    # PA  = sAnew XOR hkAV(NV)
    # SGVA = H(TDV || TDA || PAV || PA || Cl || NV || KAV)
    print("\n[DS2] Verifier checking Device A signature...")

    KAV_ver  = db_A["KAV"]
    hkAV_ver = HASH(KAV_ver, NA_sent)   # uses received (possibly tampered) NA
    TDA_chk  = HASH(db_A["IDA"], hkAV_ver)
    TDV_chk  = HASH(IDV, hkAV_ver)
    SGAV_chk = HASH(HASH(TDA_chk, TDV_chk), HASH(NA_sent, KAV_ver))

    if SGAV_chk != SGAV:
        results["attack_detected"] +=1
        print("  [FAIL] Device A signature INVALID — attack detected!")
        lcd.show("DS2 FAIL!", "Attack Detected")
        return False
    print("  [OK] Device A authenticated to Verifier")

    CAnew    = (nonce() % (P - 2)) + 1
    Randnew  = (nonce() % (P - 2)) + 1
    sAnew, sAVnew = split(CAnew, Randnew)

    NV       = nonce()
    hkAV_NV  = HASH(KAV_ver, NV)

    TDA_NV   = HASH(db_A["IDA"], hkAV_NV)
    TDV_NV   = HASH(IDV, hkAV_NV)

    PAV = XOR(db_A["sAV"], hkAV_NV)    # encrypted old verifier share
    PA  = XOR(sAnew, hkAV_NV)          # encrypted new device share
    Cl  = XOR(XOR(CAnew, hkAV_NV), db_A["hrA"])  # concealed new challenge

    SGVA = HASH(
        HASH(TDV_NV, TDA_NV),
        HASH(HASH(PAV, PA), HASH(Cl, HASH(NV, KAV_ver)))
    )

    print("  NV={} PAV={} PA={} Cl={} SGVA={}".format(NV, PAV, PA, Cl, SGVA))
    lcd.show("DS2: Verif->A", "Shares sent")

    # ── DS3: Device A verifies Verifier, computes KS ─────
    # Verify SGVA
    # sAV  = PAV XOR hkAV(NV)
    # CA   = reconstruct(sA, sAV)         — verify H(CA)==hcA
    # RA   = PUF_A(CA)                    — controlled PUF
    # hrA  = H(RA)
    # CAnew = Cl XOR hkAV(NV) XOR hrA    — recover new challenge
    # RAnew = PUF_A(CAnew)
    # hrAnew = H(RAnew)
    # sAnew = PA XOR hkAV(NV)
    # KS   = hrA XOR hrAnew
    # V1   = hrA   XOR hkAV(NAnew)       — encrypted for verifier
    # V2   = hrAnew XOR hkAV(NAnew)
    # SGAV_fin = H(TDA || TDV || V1 || V2 || NAnew || KAV)
    print("\n[DS3] Device A processing verifier message...")

    hkAV_D3  = HASH(KAV_dev, NV)
    TDA_D3   = HASH(devA["IDA"], hkAV_D3)
    TDV_D3   = HASH(IDV, hkAV_D3)
    SGVA_chk = HASH(
        HASH(TDV_D3, TDA_D3),
        HASH(HASH(PAV, PA), HASH(Cl, HASH(NV, KAV_dev)))
    )

    if SGVA_chk != SGVA:
        print("  [FAIL] Verifier signature INVALID for Device A!")
        lcd.show("DS3 FAIL!", "Verif sig bad")
        return False
    print("  [OK] Verifier authenticated to Device A")

    sAV_D3 = XOR(PAV, hkAV_D3)
    CA     = reconstruct(devA["sA"], sAV_D3)

    if HASH(CA) != devA["hcA"]:
        print("  [FAIL] Challenge hash mismatch!")
        lcd.show("DS3 FAIL!", "CA hash error")
        return False
    print("  CA={} [verified]".format(CA))

    # PUF runs only after verifier is authenticated (controlled PUF)
    RA         = PUF(CA)
    hrA_D3     = HASH(RA)
    CAnew_D3   = XOR(XOR(Cl, hkAV_D3), hrA_D3)   # recover new challenge
    RAnew      = PUF(CAnew_D3)
    hrAnew_D3  = HASH(RAnew)
    sAnew_D3   = XOR(PA, hkAV_D3)                 # recover new share

    KS_dev     = XOR(hrA_D3, hrAnew_D3)           # paper: KS = hrA XOR hrAnew

    NAnew      = nonce()
    hkAV_NA    = HASH(KAV_dev, NAnew)
    V1         = XOR(hrA_D3,   hkAV_NA)
    V2         = XOR(hrAnew_D3, hkAV_NA)

    TDA_NA     = HASH(devA["IDA"], hkAV_NA)
    TDV_NA     = HASH(IDV, hkAV_NA)
    SGAV_fin   = HASH(
        HASH(TDA_NA, TDV_NA),
        HASH(HASH(V1, V2), HASH(NAnew, KAV_dev))
    )

    print("  RA={} CAnew={} KS_dev={}".format(RA, CAnew_D3, KS_dev))
    lcd.show("DS3: Dev Keys", "KS={}".format(KS_dev % 9999))

    # FIX: Per paper DS3 — update device state for next session
    # sA   ← sAnew (new share for next round)
    # hcA  ← H(CAnew)  (hash of new challenge)
    # CAV  ← CA   (old reconstructed challenge becomes new CAV)
    devA["sA"]  = sAnew_D3
    devA["hcA"] = HASH(CAnew_D3)
    devA["CAV"] = CAnew_D3          # paper: CAV = CA (not CAnew)

    # ── DS4: Verifier final verification ─────────────────
    # Verify SGAV_fin
    # hrA    = V1 XOR hkAV(NAnew)
    # hrAnew = V2 XOR hkAV(NAnew)
    # KS     = hrA XOR hrAnew
    # Update DB: sAV←sAVnew, hrA←hrAnew, KAV←RA
    print("\n[DS4] Verifier final verification...")

    hkAV_D4    = HASH(KAV_ver, NAnew)
    TDA_D4     = HASH(db_A["IDA"], hkAV_D4)
    TDV_D4     = HASH(IDV, hkAV_D4)
    SGAV_f_chk = HASH(
        HASH(TDA_D4, TDV_D4),
        HASH(HASH(V1, V2), HASH(NAnew, KAV_ver))
    )

    if SGAV_f_chk != SGAV_fin:
        print("  [FAIL] Final signature check failed!")
        lcd.show("DS4 FAIL!", "Final check err")
        return False

    hrA_D4    = XOR(V1, hkAV_D4)
    hrAnew_D4 = XOR(V2, hkAV_D4)
    KS_ver    = XOR(hrA_D4, hrAnew_D4)

    print("  KS_dev={} KS_ver={}".format(KS_dev, KS_ver))

    if KS_dev == KS_ver:
        end=utime.ticks_us()
        exec_time=utime.ticks_diff(end,start)
        results["dv_time"].append(exec_time)
        results["auth_success"]+=1
        print("DV-AKE Time:",exec_time,"us")
        print("\n  [SUCCESS] DV-AKE complete — KS = {}".format(KS_dev))

        # Update verifier DB for next session (forward secrecy)
        # sAV ← sAVnew,  hrA ← hrAnew,  KAV ← RA (new PUF response)
        db_A["sAV"] = sAVnew
        db_A["hrA"] = hrAnew_D4
        db_A["KAV"] = RAnew      # paper: KAV = RA (new bootstrap key)

        print("  DB updated: sAV={} hrA={} KAV={}".format(
            db_A["sAV"], db_A["hrA"], db_A["KAV"]))
        lcd.show("DV-AKE SUCCESS", "Mutual Auth OK!")
        return True
    else:
        print("  [FAIL] Session key mismatch!")
        lcd.show("DV-AKE FAIL!", "KS mismatch")
        return False


# ═══════════════════════════════════════════════════════════
#  MAIN
# ═══════════════════════════════════════════════════════════
def main():
    utime.sleep_ms(2000)

    # Enrollment for both devices
    enroll(devA, db_A, PUF,   label="A")
    utime.sleep_ms(500)
    enroll(devB, db_B, PUF_B, label="B")
    utime.sleep_ms(1000)

    session = 1
    while True:
        sep("SESSION {}".format(session))
        lcd.show("Session {}".format(session), "Running...")

        # DD-AKE: Device A <-> Device B
        r1 = dd_ake()
        if r1:
            green_led.on(); red_led.off()
            print("[SESSION {}] DD-AKE: SUCCESS".format(session))
            lcd.show("DD-AKE OK!", "Sess {}".format(session))
        else:
            green_led.off(); red_led.on()
            print("[SESSION {}] DD-AKE: FAILED".format(session))
            lcd.show("DD-AKE FAIL!", "Sess {}".format(session))

        utime.sleep_ms(3000)
        green_led.off(); red_led.off()

        # DV-AKE normal
        r2 = dv_ake(inject_attack=False)
        if r2:
            green_led.on(); red_led.off()
            print("[SESSION {}] DV-AKE: SUCCESS".format(session))
            lcd.show("DV-AKE OK!", "Sess {}".format(session))
        else:
            green_led.off(); red_led.on()
            print("[SESSION {}] DV-AKE: FAILED".format(session))
            lcd.show("DV-AKE FAIL!", "Sess {}".format(session))

        utime.sleep_ms(3000)
        green_led.off(); red_led.off()

        # Every 3rd session: simulate replay attack
        if session % 3 == 0:
            print("\n[ATTACK SIM] Injecting replay attack on DV-AKE...")
            lcd.show("Attack Sim!", "Replay attack")
            utime.sleep_ms(1000)
            r3 = dv_ake(inject_attack=True)
            if not r3:
                green_led.off(); red_led.on()
                print("[ATTACK SIM] Attack DETECTED and REJECTED!")
                lcd.show("Attack BLOCKED!", "Protocol secure")
                utime.sleep_ms(2000)
                red_led.off()

        session += 1


        print("\n==========================")
        print("FINAL EXPERIMENT RESULTS")
        print("==========================")
        print("Avg Enrollment:", sum(results["enroll_time"])/len(results["enroll_time"]), "us")
        print("Avg DD:", sum(results["dd_time"])/len(results["dd_time"]), "us")
        print("Avg DV:", sum(results["dv_time"])/len(results["dv_time"]), "us")
        print("Auth Success:", results["auth_success"])
        print("Attack Detection:", results["attack_detected"])
        print("Session Matches:", results["session_match"])

        utime.sleep_ms(5000)


if __name__ == "__main__":
    main()