#!/usr/bin/env python3
import re, sys, json, hashlib, os
from pwn import remote, context
context.log_level = 'error'

HOST, PORT = "ctf.blackalps.ch", 1339

# secp256r1
p  = 0xffffffff00000001000000000000000000000000ffffffffffffffffffffffff
a  = p-3
b  = 0x5ac635d8aa3a93e7b3ebbd55769886bc651d06b0cc53b0f63bce3c3e27d2604b
n  = 0xffffffff00000000ffffffffffffffffbce6faada7179e84f3b9cac2fc632551
Gx = 0x6b17d1f2e12c4247f8bce6e563a440f277037d812deb33a0f4a13945d898c296
Gy = 0x4fe342e2fe1a7f9b8ee7eb4a7c0f9e162bce33576b315ececbb6406837bf51f5

def inv(x,m): return pow(x%m,-1,m)
def pt_add(P,Q):
    if P is None: return Q
    if Q is None: return P
    x1,y1=P; x2,y2=Q
    if x1==x2 and (y1+y2)%p==0: return None
    if P==Q: l=(3*x1*x1+a)*inv(2*y1,p)%p
    else:    l=(y2-y1)*inv(x2-x1,p)%p
    x3=(l*l-x1-x2)%p; y3=(l*(x1-x3)-y1)%p
    return (x3,y3)
def pt_mul(k,P):
    R=None
    while k>0:
        if k&1: R=pt_add(R,P)
        P=pt_add(P,P); k>>=1
    return R
G=(Gx,Gy)

# Precompute kG for k in 1..255 -> map x mod n -> k
small={}
acc=None
for k in range(1,256):
    acc=pt_add(acc,G)
    small[acc[0]%n]=(k,acc)

pk=json.load(open(os.path.join(os.path.dirname(__file__),"pubkeys.json")))
QX=[int(h,16) for h in pk["QX"]]; QY=[int(h,16) for h in pk["QY"]]
PUB=[(QX[i],QY[i]) for i in range(22)]

NAMES=["MAX VERSTAPPEN","ISACK HADJAR","CHARLES LECLERC","LEWIS HAMILTON",
"GEORGE RUSSELL","KIMI ANTONELLI","LANDO NORRIS","OSCAR PIASTRI",
"FERNANDO ALONSO","LANCE STROLL","PIERRE GASLY","FRANCO COLAPINTO",
"CARLOS SAINZ","ALEXANDER ALBON","LIAM LAWSON","ARVID LINDBLAD",
"NICO HULKENBERG","GABRIEL BORTOLETO","OLIVER BEARMAN","ESTEBAN OCON",
"SERGIO PEREZ","VALTTERI BOTTAS"]
PREFIX=[("["+nm+"] ").encode() for nm in NAMES]

# body templates: (fixed_prefix, suffix)
TEMPLATES=[
 (b"Activate ECU mode ", b""),
 (b"Go plan ",           b""),
 (b"Pit window lap ",    b""),
 (b"Fuel mode ",         b" laps"),
 (b"Front wing damage ", b" pct"),
 (b"DRS enabled, gap ",  b" sec"),
 (b"Brake bias ",        b""),
 (b"Copy that, stay out",b""),
]

def e_of(msg: bytes) -> int:
    return int.from_bytes(hashlib.sha256(msg).digest(),'big')

def recover_d(idx, msg, r, s):
    e=e_of(msg)
    if r not in small: return None
    k,_=small[r]
    # d = r^-1 (s*k - e) mod n
    d=(inv(r,n)*((s*k - e)%n))%n
    if pt_mul(d,G)==PUB[idx]:
        return d
    return None

def main():
    io=remote(HOST,PORT)
    def spy(pilot1):  # pilot1 is 1-based
        io.recvuntil(b"> ")
        io.sendline(b"1")
        io.recvuntil(b"Pilot [1-22]: ")
        io.sendline(str(pilot1).encode())
        io.recvuntil(b"Ciphertext : ")
        ct=bytes.fromhex(io.recvline().strip().decode())
        io.recvuntil(b"r          : ")
        r=int(io.recvline().strip(),16)
        io.recvuntil(b"s          : ")
        s=int(io.recvline().strip(),16)
        return ct,r,s

    # collect several messages per pilot
    msgs=[]  # (idx0, ct, r, s)
    REPS=6
    for rep in range(REPS):
        for pilot1 in range(1,23):
            ct,r,s=spy(pilot1)
            msgs.append((pilot1-1,ct,r,s))
    print(f"[+] collected {len(msgs)} messages")

    # ---- recover global keystream KS via crib propagation ----
    KS={}
    def set_ks(pos,ks_byte):
        if pos in KS and KS[pos]!=ks_byte:
            return False
        KS[pos]=ks_byte; return True
    def add_crib(ct, off, known):
        for i,bchar in enumerate(known):
            if off+i < len(ct):
                set_ks(off+i, ct[off+i]^bchar)
    # seed with prefixes
    for idx,ct,r,s in msgs:
        add_crib(ct,0,PREFIX[idx])
    # iterate
    for _ in range(30):
        changed=False
        before=len(KS)
        for idx,ct,r,s in msgs:
            P=PREFIX[idx]; boff=len(P)
            # decrypt known body bytes contiguously from boff
            body_known=bytearray()
            j=boff
            while j<len(ct) and j in KS:
                body_known.append(ct[j]^KS[j]); j+=1
            body_known=bytes(body_known)
            # find matching templates
            cands=[]
            for ti,(fp,suf) in enumerate(TEMPLATES):
                m=min(len(fp),len(body_known))
                if body_known[:m]==fp[:m]:
                    cands.append(ti)
            if len(cands)==1:
                fp,suf=TEMPLATES[cands[0]]
                add_crib(ct,boff,fp)
                if suf:
                    add_crib(ct,len(ct)-len(suf),suf)
        if len(KS)!=before: changed=True
        if not changed: break
    print(f"[+] KS bytes recovered: {len(KS)} (max pos {max(KS) if KS else -1})")

    # ---- recover private key per pilot ----
    maxlen=max(len(ct) for _,ct,_,_ in msgs)
    def decrypt_full(ct):
        out=bytearray()
        for i in range(len(ct)):
            if i not in KS: return None
            out.append(ct[i]^KS[i])
        return bytes(out)

    keys={}
    for idx,ct,r,s in msgs:
        if idx in keys: continue
        pt=decrypt_full(ct)
        if pt is None: continue
        d=recover_d(idx,pt,r,s)
        if d is not None:
            keys[idx]=d
    print(f"[+] keys recovered from full-decrypt: {len(keys)}/22 -> {sorted(keys)}")

    # fallback: for missing pilots, brute small unknown tail/number bytes
    if len(keys)<22:
        import itertools
        for idx,ct,r,s in msgs:
            if idx in keys: continue
            unknown=[i for i in range(len(ct)) if i not in KS]
            if len(unknown)>3:  # too many unknowns, skip
                continue
            base=bytearray()
            ok=True
            for i in range(len(ct)):
                base.append(ct[i]^KS[i] if i in KS else 0)
            found=False
            for combo in itertools.product(range(32,127),repeat=len(unknown)):
                for u,c in zip(unknown,combo): base[u]=c
                d=recover_d(idx,bytes(base),r,s)
                if d is not None:
                    keys[idx]=d; found=True; break
            if found:
                print(f"[+] pilot {idx} key via brute ({len(unknown)} unk)")
    print(f"[+] total keys: {len(keys)}/22")
    if len(keys)<22:
        missing=[i for i in range(22) if i not in keys]
        print("[!] MISSING pilots:",missing)
        json.dump({str(i):hex(d) for i,d in keys.items()},
                  open('keys_partial.json','w'))
        io.close(); return

    # persist keys (constant across sessions: loaded from private_keys.bin)
    json.dump({str(i):hex(d) for i,d in keys.items()},
              open(os.path.join(os.path.dirname(__file__),"keys.json"),'w'))

    # ---- forge Box Box Box! and send to all 22 ----
    MSG=b"Box Box Box!"
    e2=e_of(MSG)
    def sign(d):
        k=1; R=pt_mul(k,G); r=R[0]%n
        s=(inv(k,n)*((e2+r*d)%n))%n
        return r,s
    for idx in range(22):
        d=keys[idx]
        r,s=sign(d)
        io.recvuntil(b"> "); io.sendline(b"2")
        io.recvuntil(b"Pilot [1-22]: "); io.sendline(str(idx+1).encode())
        io.recvuntil(b"Message: "); io.sendline(MSG)
        io.recvuntil(b"r (64 hex chars): "); io.sendline(f"{r:064x}".encode())
        io.recvuntil(b"s (64 hex chars): "); io.sendline(f"{s:064x}".encode())
    # dump everything remaining (RADIO lines + FULL GRID + FLAG)
    import time; time.sleep(1)
    data=io.recvrepeat(4)
    txt=data.decode('latin1',errors='replace')
    print(txt)
    m=re.search(r"BA26\{[^}]*\}", txt)
    if m: print("\n[FLAG]", m.group(0))
    io.close()

if __name__=="__main__":
    main()