---

from flask import Flask, request, jsonify
from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
import base64, struct, os

app = Flask(__name__)
SERVICE_KEY = os.urandom(16)
KEY_VICTIM = os.urandom(16)
KEY_ATTACKER = os.urandom(16)
SECRET_VICTIM = os.urandom(16)
AAD_PREFIX = b"user-aad-prefix:"
FLAG = os.environ.get("FLAG") or "DUCKERZ{default_flag}"
FLAG_IV = get_random_bytes(12)

def aes_ecb_encrypt_block(key: bytes, block16: bytes) -> bytes:
    cipher = AES.new(key, AES.MODE_ECB)
    return cipher.encrypt(block16)

def aes_ctr_keystream(key: bytes, j0: bytes, length: int) -> bytes:
    out = bytearray()
    ctr = int.from_bytes(j0, "big")
    ctr += 1
    cipher = AES.new(key, AES.MODE_ECB)
    while len(out) < length:
        block = (ctr & ((1 << 128) - 1)).to_bytes(16, "big")
        out_block = cipher.encrypt(block)
        take = min(16, length - len(out))
        out.extend(out_block[:take])
        ctr += 1
    return bytes(out)

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

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

def ghash(H_bytes: bytes, aad: bytes, ct: bytes) -> bytes:
    H = int.from_bytes(H_bytes, "big")
    len_a_bits = len(aad) * 8
    len_c_bits = len(ct) * 8
    len_block = struct.pack(">QQ", len_a_bits, len_c_bits)
    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")

def compute_j0_from_iv(iv12: bytes) -> bytes:
    return iv12 + b"\x00\x00\x00\x01"

def buggy_encrypt(tenant_key: bytes, iv12: bytes, pt: bytes, aad: bytes, global_H: bytes):
    j0 = compute_j0_from_iv(iv12)
    ks = aes_ctr_keystream(tenant_key, j0, len(pt))
    ct = bytes(x ^ y for x, y in zip(pt, ks))
    e_j0 = aes_ecb_encrypt_block(tenant_key, j0)
    g = ghash(global_H, aad, ct)
    tag = bytes(x ^ y for x, y in zip(e_j0, g))
    return ct, tag

def buggy_decrypt_and_verify(tenant_key: bytes, iv12: bytes, ct: bytes, tag: bytes, aad: bytes, global_H: bytes):
    j0 = compute_j0_from_iv(iv12)
    e_j0 = aes_ecb_encrypt_block(tenant_key, j0)
    g = ghash(global_H, aad, ct)
    expected_tag = bytes(x ^ y for x, y in zip(e_j0, g))
    if expected_tag != tag:
        return None, False
    ks = aes_ctr_keystream(tenant_key, j0, len(ct))
    pt = bytes(x ^ y for x, y in zip(ct, ks))
    return pt, True

GLOBAL_H = aes_ecb_encrypt_block(SERVICE_KEY, b"\x00" * 16)
AAD_VICTIM = AAD_PREFIX + SECRET_VICTIM
FLAG_CT, FLAG_TAG = buggy_encrypt(KEY_VICTIM, FLAG_IV, FLAG.encode("latin1"), AAD_VICTIM, GLOBAL_H)
FLAG_BLOB = base64.b64encode(FLAG_IV + FLAG_CT + FLAG_TAG).decode()

@app.route("/", methods=["GET"])
def index():
    return jsonify({
        "note": "Use /flag, /aes_ecb_service, /aes_ecb_tenant, /ghash_probe, /verify.",
        "AAD_PREFIX_hex": AAD_PREFIX.hex(),
        "source_code": "Available at /source"
    })

@app.route("/flag", methods=["GET"])
def get_flag():
    return jsonify({"flag_blob": FLAG_BLOB})

@app.route("/source", methods=["GET"])
def get_source():
    import os
    current_file = os.path.abspath(__file__)
    try:
        with open(current_file, 'r', encoding='utf-8') as f:
            source_code = f.read()
        return source_code, 200, {'Content-Type': 'text/plain; charset=utf-8'}
    except Exception as e:
        return f"Error reading source code: {e}", 500

@app.route("/aes_ecb_service", methods=["POST"])
def aes_ecb_service():
    data = request.get_json(force=True)
    block_b64 = data.get("block_b64", "")
    try:
        blk = base64.b64decode(block_b64)
    except Exception:
        return jsonify({"error": "block_b64 must be base64"}), 400
    if len(blk) != 16:
        return jsonify({"error": "block must be 16 bytes"}), 400
    c = aes_ecb_encrypt_block(SERVICE_KEY, blk)
    return jsonify({"cipher_b64": base64.b64encode(c).decode()})

@app.route("/aes_ecb_tenant", methods=["POST"])
def aes_ecb_tenant():
    data = request.get_json(force=True)
    who = data.get("who", "")
    block_b64 = data.get("block_b64", "")
    if who not in ("victim", "attacker"):
        return jsonify({"error": "who must be 'victim' or 'attacker'"}), 400
    try:
        blk = base64.b64decode(block_b64)
    except Exception:
        return jsonify({"error": "block_b64 must be base64"}), 400
    if len(blk) != 16:
        return jsonify({"error": "block must be 16 bytes"}), 400
    key = KEY_VICTIM if who == "victim" else KEY_ATTACKER
    c = aes_ecb_encrypt_block(key, blk)
    return jsonify({"cipher_b64": base64.b64encode(c).decode()})

@app.route("/encrypt", methods=["POST"])
def encrypt_endpoint():
    data = request.get_json(force=True)
    who = data.get("who", "")
    if who not in ("victim", "attacker"):
        return jsonify({"error": "who must be 'victim' or 'attacker'"}), 400
    aad_hex = data.get("aad_hex", "")
    msg_b64 = data.get("msg_b64", "")
    try:
        aad_from_client = bytes.fromhex(aad_hex) if aad_hex != "" else b""
        pt = base64.b64decode(msg_b64) if msg_b64 != "" else b""
    except Exception:
        return jsonify({"error": "invalid inputs"}), 400
    iv = get_random_bytes(12)
    aad_actual = AAD_PREFIX + aad_from_client
    tenant_key = KEY_VICTIM if who == "victim" else KEY_ATTACKER
    ct, tag = buggy_encrypt(tenant_key, iv, pt, aad_actual, GLOBAL_H)
    blob = base64.b64encode(iv + ct + tag).decode()
    return jsonify({"blob": blob})

@app.route("/ghash_probe", methods=["POST"])
def ghash_probe():
    data = request.get_json(force=True)
    who = data.get("who", "")
    if who != "attacker":
        return jsonify({"error": "ghash_probe allowed only for attacker in this lab"}), 403
    aad_hex = data.get("aad_hex", "")
    ct_b64 = data.get("ct_b64", "")
    try:
        aad = bytes.fromhex(aad_hex) if aad_hex != "" else b""
        ct = base64.b64decode(ct_b64) if ct_b64 != "" else b""
    except Exception:
        return jsonify({"error": "invalid inputs"}), 400
    g = ghash(GLOBAL_H, aad, ct)
    return jsonify({"ghash_b64": base64.b64encode(g).decode()})

@app.route("/verify", methods=["POST"])
def verify():
    data = request.get_json(force=True)
    who = data.get("who", "")
    if who not in ("victim", "attacker"):
        return jsonify({"error": "who must be 'victim' or 'attacker'"}), 400
    blob_b64 = data.get("blob", "")
    aad_hex = data.get("aad_hex", "")
    try:
        blob = base64.b64decode(blob_b64)
    except Exception:
        return jsonify({"error": "blob must be base64"}), 400
    if len(blob) < 12 + 16:
        return jsonify({"error": "blob too short"}), 400
    try:
        aad_from_client = bytes.fromhex(aad_hex) if aad_hex != "" else b""
    except Exception:
        return jsonify({"error": "aad_hex invalid"}), 400
    iv = blob[:12]; ct = blob[12:-16]; tag = blob[-16:]
    aad_actual = AAD_PREFIX + aad_from_client
    tenant_key = KEY_VICTIM if who == "victim" else KEY_ATTACKER
    pt, ok = buggy_decrypt_and_verify(tenant_key, iv, ct, tag, aad_actual, GLOBAL_H)
    if ok:
        return jsonify({"status": "valid", "plaintext_b64": base64.b64encode(pt).decode()})
    else:
        return jsonify({"status": "invalid"}), 403



if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000, debug=True)