#!/usr/bin/env python3
"""
REWIND_REVENGE — AES-GCM nonce-reuse forgery (Joux "Forbidden Attack").

Usage:
  python3 solve.py HOST PORT        # TLS on by default
  python3 solve.py HOST PORT --no-ssl
"""

from __future__ import annotations

import argparse
import sys

try:
    from pwn import remote
except ImportError as exc:
    raise SystemExit("pwntools required: pip install pwntools") from exc

ONE = 1 << 127


def gf_mult(x: int, y: int) -> int:
    r = 0xE1000000000000000000000000000000
    z = 0
    v = x
    for i in range(127, -1, -1):
        if (y >> i) & 1:
            z ^= v
        v = (v >> 1) ^ r if v & 1 else v >> 1
    return z


def gf_pow(x: int, n: int) -> int:
    result = ONE
    base = x
    while n:
        if n & 1:
            result = gf_mult(result, base)
        base = gf_mult(base, base)
        n >>= 1
    return result


def gf_inv(x: int) -> int:
    return gf_pow(x, (1 << 128) - 2)


def gf_sqrt(x: int) -> int:
    return gf_pow(x, 1 << 127)


def b2i(b: bytes) -> int:
    return int.from_bytes(b, "big")


def i2b(i: int) -> bytes:
    return (i & ((1 << 128) - 1)).to_bytes(16, "big")


def ghash_single_block(h: int, c_int: int, aad_len_bits: int = 0, ct_len_bits: int = 128) -> int:
    length_block = b2i(
        aad_len_bits.to_bytes(8, "big") + ct_len_bits.to_bytes(8, "big")
    )
    return gf_mult(gf_mult(c_int, h) ^ length_block, h)


def seal(io, pt_bytes: bytes) -> tuple[bytes, bytes]:
    io.sendlineafter(b"> ", b"1")
    io.sendlineafter(b"> ", pt_bytes.hex().encode())
    out = io.recvuntil(b"> ", timeout=5).decode()
    ct_hex = out.split("ciphertext = ")[1].split("\n")[0].strip()
    tag_hex = out.split("tag = ")[1].split("\n")[0].strip()
    return bytes.fromhex(ct_hex), bytes.fromhex(tag_hex)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("host")
    parser.add_argument("port", type=int)
    parser.add_argument("--no-ssl", action="store_true")
    args = parser.parse_args()

    io = remote(args.host, args.port, ssl=not args.no_ssl)

    p1 = bytes([0x33] * 16)
    p2 = bytes([0x44] * 16)
    c1, t1 = seal(io, p1)
    c2, t2 = seal(io, p2)

    c1i, t1i = b2i(c1), b2i(t1)
    c2i, t2i = b2i(c2), b2i(t2)

    h2 = gf_mult(t1i ^ t2i, gf_inv(c1i ^ c2i))
    h = gf_sqrt(h2)
    e = t1i ^ ghash_single_block(h, c1i)
    assert e == t2i ^ ghash_single_block(h, c2i), "H/E derivation inconsistent"

    ks = b2i(p1) ^ c1i
    target = b"print_the_flag!!"
    c_forge_i = b2i(target) ^ ks
    t_forge_i = e ^ ghash_single_block(h, c_forge_i)
    c_forge, t_forge = i2b(c_forge_i), i2b(t_forge_i)

    print("forged ct :", c_forge.hex())
    print("forged tag:", t_forge.hex())

    io.sendlineafter(b"> ", b"2")
    io.sendlineafter(b">", c_forge.hex().encode())
    io.sendlineafter(b">", t_forge.hex().encode())
    print(io.recvall(timeout=5).decode(errors="replace"))


if __name__ == "__main__":
    main()