---

#!/usr/bin/env python3
import urllib.request
import json
import base64
import struct
import sys

BASE = sys.argv[1] if len(sys.argv) > 1 else "http://localhost:5000"

R = 0xE1000000000000000000000000000000
MASK128 = (1 << 128) - 1

# --- GF(2^128) ---
def gf128_mul(x, y):
    z = 0
    v = x & MASK128
    for i in range(128):
        if (y >> i) & 1:
            z ^= v
        msb = (v >> 127) & 1
        v = (v << 1) & MASK128
        if msb:
            v ^= R
    return z & MASK128

def gf128_inv(x):
    """Инверсия в GF(2^128) через расширенный алгоритм Евклида."""
    if x == 0:
        raise ZeroDivisionError
    MOD = (1 << 128) | 0x87
    def poly_deg(a):
        return a.bit_length() - 1 if a else -1
    def poly_divmod(a, b):
        if b == 0:
            raise ZeroDivisionError
        q = 0
        while a and poly_deg(a) >= poly_deg(b):
            shift = poly_deg(a) - poly_deg(b)
            q ^= (1 << shift)
            a ^= (b << shift)
        return q, a
    def poly_mul(a, b):
        result = 0
        while b:
            if b & 1:
                result ^= a
            a <<= 1
            b >>= 1
        return result
    def poly_mod(a, m):
        while a and poly_deg(a) >= poly_deg(m):
            shift = poly_deg(a) - poly_deg(m)
            a ^= (m << shift)
        return a
    r0, r1 = MOD, x
    s0, s1 = 0, 1
    steps = 0
    while r1 != 0:
        steps += 1
        if steps > 1000:
            raise RuntimeError("EGCD зациклился")
        q, r = poly_divmod(r0, r1)
        r0, r1 = r1, r
        s0, s1 = s1, s0 ^ poly_mod(poly_mul(q, s1), MOD)
    return s0 & MASK128

def ghash(H_bytes, aad, ct):
    H = int.from_bytes(H_bytes, "big")
    len_block = struct.pack(">QQ", len(aad) * 8, len(ct) * 8)
    data = aad + ct + len_block
    Y = 0
    for i in range(0, len(data), 16):
        block = data[i:i+16]
        if len(block) < 16:
            block = block + b"\x00" * (16 - len(block))
        X = int.from_bytes(block, "big")
        Y = gf128_mul(Y ^ X, H)
    return Y.to_bytes(16, "big")

# --- HTTP ---
def post(path, payload):
    req = urllib.request.Request(
        BASE + path,
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
        method="POST"
    )
    with urllib.request.urlopen(req, timeout=15) as resp:
        return json.loads(resp.read().decode())

def get(path):
    with urllib.request.urlopen(BASE + path, timeout=15) as resp:
        return json.loads(resp.read().decode())

def main():
    # --- ТЕСТ МАТЕМАТИКИ ---
    print("=== ТЕСТ МАТЕМАТИКИ ===")
    for t in [1, 2, 3, 0x80, 0x123456789abcdef0123456789abcdef0, MASK128]:
        inv_t = gf128_inv(t)
        p = gf128_mul(t, inv_t)
        status = "OK" if p == 1 else "FAIL"
        print(f"  inv({t:032x}) * {t:032x} = {p:032x} {status}")
        if p != 1:
            print("!!! МАТЕМАТИКА СЛОМАНА, ДАЛЬШЕ НЕТ СМЫСЛА !!!")
            return
    print("=== МАТЕМАТИКА OK ===\n")

    # 1. FLAG_BLOB
    r = get("/flag")
    flag_blob = base64.b64decode(r["flag_blob"])
    FLAG_IV  = flag_blob[:12]
    FLAG_CT  = flag_blob[12:-16]
    FLAG_TAG = flag_blob[-16:]
    print(f"[+] FLAG_IV  = {FLAG_IV.hex()}")
    print(f"[+] FLAG_CT  = {FLAG_CT.hex()} ({len(FLAG_CT)} bytes)")
    print(f"[+] FLAG_TAG = {FLAG_TAG.hex()}")

    # 2. GLOBAL_H через /aes_ecb_service (0^16)
    print("\n[*] Получаем GLOBAL_H через /aes_ecb_service...")
    r = post("/aes_ecb_service", {
        "block_b64": base64.b64encode(b"\x00" * 16).decode()
    })
    GLOBAL_H = base64.b64decode(r["cipher_b64"])
    print(f"[+] GLOBAL_H = {GLOBAL_H.hex()}")

    # 3. E_K(FLAG_J0) через /aes_ecb_tenant (who=victim)
    print("\n[*] Получаем E_K(FLAG_J0) через /aes_ecb_tenant...")
    FLAG_J0 = FLAG_IV + b"\x00\x00\x00\x01"
    r = post("/aes_ecb_tenant", {
        "who": "victim",
        "block_b64": base64.b64encode(FLAG_J0).decode()
    })
    E_K_J0 = base64.b64decode(r["cipher_b64"])
    print(f"[+] E_K(J0) = {E_K_J0.hex()}")

    # 4. known = GHASH(H, AAD_VICTIM, FLAG_CT) = FLAG_TAG XOR E_K(J0)
    known = bytes(a ^ b for a, b in zip(FLAG_TAG, E_K_J0))
    print(f"[+] known = {known.hex()}")

    # 5. G0 = GHASH(H, AAD_PREFIX || 0^16, FLAG_CT)
    AAD_PREFIX = b"user-aad-prefix:"
    G0 = ghash(GLOBAL_H, AAD_PREFIX + b"\x00" * 16, FLAG_CT)
    print(f"[+] G0 = {G0.hex()}")

    # 6. Разница = known XOR G0 = SECRET_VICTIM * H^k
    diff = bytes(a ^ b for a, b in zip(known, G0))
    diff_int = int.from_bytes(diff, "big")
    print(f"[+] diff = {diff.hex()}")

    # 7. Находим k (степень H для позиции SECRET_VICTIM)
    # data = AAD_PREFIX (16) || SECRET_VICTIM (16) || FLAG_CT || len_block
    # Общее число блоков N = 2 + ceil(len(FLAG_CT)/16) + 1
    N = 2 + (len(FLAG_CT) + 15) // 16 + 1
    # SECRET_VICTIM — блок с индексом 1 (0-индексация)
    # Его вклад в GHASH = SECRET_VICTIM * H^(N-1)
    k = N - 1
    print(f"[+] N = {N}, k = {k}")

    # 8. Считаем H^k
    H = int.from_bytes(GLOBAL_H, "big")
    H_k = 1
    for _ in range(k):
        H_k = gf128_mul(H_k, H)
    print(f"[+] H^k = {H_k:032x}")

    # 9. SECRET_VICTIM = diff * inv(H^k)
    inv_H_k = gf128_inv(H_k)
    SECRET_VICTIM = gf128_mul(diff_int, inv_H_k).to_bytes(16, "big")
    print(f"[+] SECRET_VICTIM = {SECRET_VICTIM.hex()}")

    # 10. Проверка: GHASH(H, AAD_PREFIX || SECRET_VICTIM, FLAG_CT) == known?
    check = ghash(GLOBAL_H, AAD_PREFIX + SECRET_VICTIM, FLAG_CT)
    print(f"[+] check = {check.hex()}")
    print(f"[+] known = {known.hex()}")
    if check == known:
        print("[+] ПРОВЕРКА ПРОЙДЕНА!")
    else:
        print("[-] ПРОВЕРКА НЕ ПРОЙДЕНА")

    # 11. /verify
    print("\n[*] Отправляем на /verify...")
    r = post("/verify", {
        "who": "victim",
        "blob": base64.b64encode(flag_blob).decode(),
        "aad_hex": SECRET_VICTIM.hex()
    })
    print(f"[+] /verify response: {r}")
    if "plaintext_b64" in r:
        flag = base64.b64decode(r["plaintext_b64"]).decode()
        print(f"\n[***] FLAG: {flag}")
        return flag
    return None

if __name__ == "__main__":
    main()