#!/usr/bin/env python3
"""glagla — polynomial ROS forgery against a 2-of-3 FROST cosigner.

Flaw: the per-signer binding factor rho = H(id, msg) does NOT commit to the
round-1 commitments, and the player's commitment (d2,e2) is attacker-chosen and
unchecked. So the group commitment R = A_i + W_i is fully attacker-shiftable
(W_i under our control) while the server's response z1 = a_i + c_i*x uses the
commitment-independent scalar a_i = dlog(A_i). This decouples the Fiat-Shamir
challenge c_i = H(R_i, Y, msg) from the server nonce -> a Schnorr signing oracle
on x = lambda_1*share_1 where we may freely sample the challenge. With 256
parallel sessions the Benhamouda et al. polynomial ROS attack forges a signature
on the forbidden TARGET message.
"""
import re
import sys
import ec
from ec import N, G, jmul, jadd, decode_point, encode_point, binding, challenge

HOST, PORT = "ctf.blackalps.ch", 1337
L = 256  # bit length -> number of parallel sessions

from pwn import remote, context
context.log_level = "warning"


def inv(a):
    return pow(a % N, N - 2, N)


def main():
    io = remote(HOST, PORT)

    # ---- params ----
    io.recvuntil(b"glagla>")
    io.sendline(b"params")
    pr = io.recvuntil(b"glagla>").decode()

    def grab(key):
        m = re.search(key + r"\s+([0-9a-fA-F]+)", pr)
        return m.group(1)

    Y = decode_point(grab("group_key"))
    verify1 = decode_point(grab("verify_1"))
    verify2 = decode_point(grab("verify_2"))
    share2 = int(grab("your_share"), 16)
    lam1 = int(grab("lambda_1"), 16)
    lam2 = int(grab("lambda_2"), 16)
    target = bytes.fromhex(grab("target"))
    print(f"[+] target = {target!r}")
    S2 = (lam2 * share2) % N  # known player contribution lambda_2*share_2

    # ---- open L sessions ----
    io.sendline(f"open {L}".encode())
    ob = io.recvuntil(b"glagla>").decode()
    rows = re.findall(r"^\s*(\d+)\s+([0-9a-fA-F]{66})\s+([0-9a-fA-F]{66})\s*$",
                      ob, re.M)
    rows = {int(s): (d, e) for s, d, e in rows}
    assert len(rows) >= L, f"only got {len(rows)} sessions"
    print(f"[+] opened {len(rows)} sessions")

    sess = []
    for i in range(L):
        D1 = decode_point(rows[i][0])
        E1 = decode_point(rows[i][1])
        msg = b"forge-%d" % i
        assert msg != target
        rho_s = binding(1, msg)
        rho_p = binding(2, msg)
        A = jadd(D1, jmul(rho_s, E1))          # server effective commitment
        # two candidate committed points (dlog irrelevant)
        P0 = jmul((i * 2654435761 + 12345) % N or 1, G)
        P1 = jmul((i * 40503 + 99991) % N or 1, G)
        c0 = challenge(P0, Y, msg)
        c1 = challenge(P1, Y, msg)
        assert c0 != c1
        delta = (c1 - c0) % N
        rho = (pow(2, i, N) * inv(delta)) % N   # coefficient 2^i / (c1-c0)
        sess.append(dict(msg=msg, rho_p=rho_p, A=A, P0=P0, P1=P1,
                         c0=c0, c1=c1, rho=rho))

    # ---- build the forgery target commitment R* and the ROS bit target ----
    Rstar = ec.INF
    K = 0
    for s in sess:
        Rstar = jadd(Rstar, jmul(s["rho"], s["A"]))
        K = (K + s["rho"] * s["c0"]) % N
    estar = challenge(Rstar, Y, target)
    V = (estar - K) % N
    assert V < (1 << L)
    for i, s in enumerate(sess):
        s["b"] = (V >> i) & 1

    # ---- round two: one sign per session with the selected challenge ----
    zsum = 0
    for i, s in enumerate(sess):
        P = s["P1"] if s["b"] else s["P0"]
        cexp = s["c1"] if s["b"] else s["c0"]
        # realise committed point P = A + D2 + rho_p*G  with E2 = G
        D2 = jadd(P, jadd(jmul(N - 1, s["A"]), jmul((N - s["rho_p"]) % N, G)))
        io.sendline(b"sign %d %s %s %s" % (
            i, encode_point(D2).hex().encode(), encode_point(G).hex().encode(),
            s["msg"].hex().encode()))
        resp = io.recvuntil(b"glagla>").decode()
        R_h = re.search(r"\bR\s+([0-9a-fA-F]{66})", resp).group(1)
        c_h = re.search(r"\bc\s+([0-9a-fA-F]{64})", resp).group(1)
        z_h = re.search(r"\bz1\s+([0-9a-fA-F]{64})", resp).group(1)
        assert ec.to_affine(decode_point(R_h)) == ec.to_affine(P), f"R mismatch session {i}"
        assert int(c_h, 16) == cexp, f"c mismatch session {i}"
        zsum = (zsum + s["rho"] * int(z_h, 16)) % N
        if i % 64 == 0:
            print(f"[+] signed session {i}")

    # s* = sum rho_i z_i satisfies s*G = R* + e*X (X = lam1*verify1)
    # full forgery: z* = s* + e**S2  ->  z*G = R* + e*Y
    zstar = (zsum + estar * S2) % N

    # local verification: z*G == R* + e*Y
    lhs = jmul(zstar, G)
    rhs = jadd(Rstar, jmul(estar, Y))
    assert ec.to_affine(lhs) == ec.to_affine(rhs), "local verify failed!"
    print("[+] forgery verifies locally")

    io.sendline(b"submit %s %s" % (
        encode_point(Rstar).hex().encode(), zstar.to_bytes(32, "big").hex().encode()))
    out = io.recvall(timeout=5).decode(errors="replace")
    print(out)
    m = re.search(r"BA26\{[^}]*\}", out)
    if m:
        print("FLAG:", m.group(0))
    io.close()


if __name__ == "__main__":
    main()