setup.sh

#!/usr/bin/env bash
# =============================================================================
# Risk Platform — Full Setup & Start Script
# =============================================================================
# Usage:
#   ./setup.sh            — first-time setup + start everything
#   ./setup.sh --restart  — tear down and restart all services cleanly
#   ./setup.sh --runner   — (re)start only the host runner (port 9999)
# =============================================================================

set -euo pipefail

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ENV_FILE="$ROOT/.env"
VENV="$ROOT/riskenv"
PYTHON="$VENV/bin/python"
PROWLER="$VENV/bin/prowler"
CARTOGRAPHY="$VENV/bin/cartography"
RUNNER_LOG="/tmp/runner.log"
RUNNER_PID_FILE="/tmp/risk-runner.pid"

# ── Colours ───────────────────────────────────────────────────────────────────
RED='\033[0;31m'; YELLOW='\033[1;33m'; GREEN='\033[0;32m'
CYAN='\033[0;36m'; BOLD='\033[1m'; RESET='\033[0m'

info()    { echo -e "${CYAN}[INFO]${RESET}  $*"; }
ok()      { echo -e "${GREEN}[OK]${RESET}    $*"; }
warn()    { echo -e "${YELLOW}[WARN]${RESET}  $*"; }
error()   { echo -e "${RED}[ERROR]${RESET} $*"; }
section() { echo -e "\n${BOLD}━━━  $*  ━━━${RESET}"; }
die()     { error "$*"; exit 1; }

# ── Argument handling ─────────────────────────────────────────────────────────
MODE="setup"
[[ "${1:-}" == "--restart" ]] && MODE="restart"
[[ "${1:-}" == "--runner"  ]] && MODE="runner"

# =============================================================================
# SECTION 1 — Prerequisites
# =============================================================================
section "Checking prerequisites"

check_cmd() {
    if ! command -v "$1" &>/dev/null; then
        die "'$1' not found. $2"
    fi
    ok "$1 found ($(command -v "$1"))"
}

check_cmd docker   "Install Docker: https://docs.docker.com/get-docker/"
check_cmd python3  "Install Python 3.11+: https://python.org"
check_cmd curl     "Install curl via your package manager"

# Docker daemon running?
if ! docker info &>/dev/null; then
    die "Docker daemon is not running. Start it and retry."
fi
ok "Docker daemon is running"

# Docker Compose (v2 plugin or standalone)
if docker compose version &>/dev/null 2>&1; then
    COMPOSE="docker compose"
elif docker-compose version &>/dev/null 2>&1; then
    COMPOSE="docker-compose"
else
    die "Docker Compose not found. Install it: https://docs.docker.com/compose/install/"
fi
ok "Docker Compose found ($($COMPOSE version --short 2>/dev/null || echo 'ok'))"

# Python version >= 3.11
PY_VER=$(python3 -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')")
PY_MAJOR=$(echo "$PY_VER" | cut -d. -f1)
PY_MINOR=$(echo "$PY_VER" | cut -d. -f2)
if [[ "$PY_MAJOR" -lt 3 ]] || [[ "$PY_MAJOR" -eq 3 && "$PY_MINOR" -lt 11 ]]; then
    die "Python 3.11+ required, found $PY_VER"
fi
ok "Python $PY_VER"

# =============================================================================
# SECTION 2 — .env file
# =============================================================================
section "Environment configuration"

if [[ ! -f "$ENV_FILE" ]]; then
    if [[ ! -f "$ROOT/.env.example" ]]; then
        die ".env.example not found. Is this the right directory?"
    fi
    cp "$ROOT/.env.example" "$ENV_FILE"
    warn ".env created from .env.example"
    ok ".env file created — configure AWS credentials via the dashboard ⚙ Account Setup panel after startup"
fi
ok ".env file exists"

# Load env vars (skip comments and blank lines)
set -o allexport
# shellcheck disable=SC2046
eval $(grep -v '^#' "$ENV_FILE" | grep -v '^[[:space:]]*$' | sed 's/[[:space:]]*$//')
set +o allexport

# Validate GitHub token if set
if [[ -n "${GITHUB_TOKEN:-}" ]]; then
    info "Validating GitHub token..."
    GH_STATUS=$(curl -s -o /dev/null -w "%{http_code}" \
        -H "Authorization: token ${GITHUB_TOKEN}" \
        https://api.github.com/user)
    if [[ "$GH_STATUS" == "200" ]]; then
        ok "GitHub token valid"
    else
        warn "GitHub token returned HTTP $GH_STATUS — GitHub pipeline may fail"
    fi
else
    warn "GITHUB_TOKEN not set — GitHub pipeline will be unavailable (set it via ⚙ Account Setup)"
fi

# =============================================================================
# SECTION 3 — Python virtualenv + dependencies
# =============================================================================
section "Python virtualenv"

if [[ ! -d "$VENV" ]]; then
    info "Creating virtualenv at $VENV ..."
    python3 -m venv "$VENV"
    ok "Virtualenv created"
else
    ok "Virtualenv already exists"
fi

info "Installing/updating Python dependencies..."
"$VENV/bin/pip" install --quiet --upgrade pip
"$VENV/bin/pip" install --quiet -r "$ROOT/engine/requirements.txt"
ok "Dependencies installed"

# Verify key tools are present in venv
for tool in prowler cartography; do
    if [[ ! -f "$VENV/bin/$tool" ]]; then
        die "$tool not found in venv after pip install. Check engine/requirements.txt"
    fi
    ok "$tool installed ($("$VENV/bin/$tool" --version 2>/dev/null | head -1 || echo 'ok'))"
done

# =============================================================================
# SECTION 4 — Docker services
# =============================================================================
section "Docker services"

cd "$ROOT"

if [[ "$MODE" == "restart" ]]; then
    info "Tearing down existing containers..."
    $COMPOSE down --remove-orphans
    ok "Containers stopped"
fi

info "Starting Docker services (this may take a minute on first run)..."
$COMPOSE up -d --build

# Wait for Neo4j
info "Waiting for Neo4j to be ready..."
NEO4J_READY=false
for i in $(seq 1 30); do
    if curl -s http://localhost:7474 &>/dev/null; then
        NEO4J_READY=true
        break
    fi
    sleep 3
done
if [[ "$NEO4J_READY" != "true" ]]; then
    die "Neo4j did not become ready in 90s. Check: docker logs risk-platform-neo4j-1"
fi
ok "Neo4j is ready (http://localhost:7474)"

# Wait for DefectDojo — it runs DB migrations on first boot, takes ~60s
info "Waiting for DefectDojo to be ready (DB migrations may take ~60s)..."
DOJO_READY=false
for i in $(seq 1 40); do
    HTTP=$(curl -s -o /dev/null -w "%{http_code}" http://localhost:8080/login 2>/dev/null || true)
    if [[ "$HTTP" == "200" ]]; then
        DOJO_READY=true
        break
    fi
    sleep 5
done
if [[ "$DOJO_READY" != "true" ]]; then
    die "DefectDojo did not become ready in 200s. Check: docker logs risk-platform-defectdojo-1"
fi
ok "DefectDojo is ready (http://localhost:8080)"

# Wait for dashboard
info "Waiting for dashboard to be ready..."
DASH_READY=false
for i in $(seq 1 15); do
    if curl -s http://localhost:8888 &>/dev/null; then
        DASH_READY=true
        break
    fi
    sleep 3
done
if [[ "$DASH_READY" != "true" ]]; then
    warn "Dashboard not responding yet — it may still be starting"
else
    ok "Dashboard is ready (http://localhost:8888)"
fi

# =============================================================================
# SECTION 5 — DefectDojo API token
# =============================================================================
section "DefectDojo API token"

DOJO_ADMIN_PASSWORD="${DOJO_ADMIN_PASSWORD:-admin}"

# Try to fetch a fresh token regardless — always use the live one
info "Fetching DefectDojo API token..."
DOJO_TOKEN_RESPONSE=$(curl -s -X POST http://localhost:8080/api/v2/api-token-auth/ \
    -H "Content-Type: application/json" \
    -d "{\"username\":\"admin\",\"password\":\"${DOJO_ADMIN_PASSWORD}\"}" 2>/dev/null || true)

FRESH_TOKEN=$(echo "$DOJO_TOKEN_RESPONSE" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('token',''))" 2>/dev/null || true)

if [[ -z "$FRESH_TOKEN" ]]; then
    if [[ -n "${DOJO_API_TOKEN:-}" ]]; then
        warn "Could not fetch fresh token — using existing DOJO_API_TOKEN from .env"
        FRESH_TOKEN="$DOJO_API_TOKEN"
    else
        die "Could not fetch DefectDojo API token. DefectDojo may still be initialising — wait 30s and retry."
    fi
else
    ok "DefectDojo API token obtained"
    # Update .env if token changed
    if [[ "${DOJO_API_TOKEN:-}" != "$FRESH_TOKEN" ]]; then
        if grep -q "^DOJO_API_TOKEN=" "$ENV_FILE"; then
            sed -i "s|^DOJO_API_TOKEN=.*|DOJO_API_TOKEN=${FRESH_TOKEN}|" "$ENV_FILE"
        else
            echo "DOJO_API_TOKEN=${FRESH_TOKEN}" >> "$ENV_FILE"
        fi
        export DOJO_API_TOKEN="$FRESH_TOKEN"
        ok "DOJO_API_TOKEN updated in .env"
    fi
fi

# Verify token works
TOKEN_CHECK=$(curl -s -o /dev/null -w "%{http_code}" \
    -H "Authorization: Token ${FRESH_TOKEN}" \
    http://localhost:8080/api/v2/users/ 2>/dev/null || true)
if [[ "$TOKEN_CHECK" != "200" ]]; then
    die "DefectDojo token verification failed (HTTP $TOKEN_CHECK). Check admin password in .env."
fi
ok "DefectDojo token verified"

# =============================================================================
# SECTION 6 — Host runner (port 9999)
# =============================================================================
section "Pipeline runner (port 9999)"

start_runner() {
    # Kill any existing runner
    if [[ -f "$RUNNER_PID_FILE" ]]; then
        OLD_PID=$(cat "$RUNNER_PID_FILE")
        if kill -0 "$OLD_PID" 2>/dev/null; then
            info "Stopping existing runner (PID $OLD_PID)..."
            kill "$OLD_PID" 2>/dev/null || true
            sleep 1
        fi
        rm -f "$RUNNER_PID_FILE"
    fi
    # Also kill anything on port 9999
    fuser -k 9999/tcp 2>/dev/null || true
    sleep 1

    info "Starting runner..."
    cd "$ROOT"
    # Re-export env so runner subprocess has all vars
    set -o allexport
    eval $(grep -v '^#' "$ENV_FILE" | grep -v '^[[:space:]]*$' | sed 's/[[:space:]]*$//')
    set +o allexport

    nohup "$PYTHON" engine/runner.py >> "$RUNNER_LOG" 2>&1 &
    RUNNER_PID=$!
    echo "$RUNNER_PID" > "$RUNNER_PID_FILE"
    disown "$RUNNER_PID"

    # Wait for it to bind
    RUNNER_READY=false
    for i in $(seq 1 15); do
        if curl -s http://localhost:9999/preflight/aws &>/dev/null; then
            RUNNER_READY=true
            break
        fi
        sleep 1
    done

    if [[ "$RUNNER_READY" != "true" ]]; then
        die "Runner did not start. Check $RUNNER_LOG"
    fi
    ok "Runner started (PID $RUNNER_PID, log: $RUNNER_LOG)"
}

if [[ "$MODE" == "runner" ]]; then
    start_runner
    echo ""
    ok "Runner restarted. Done."
    exit 0
fi

start_runner

# =============================================================================
# SECTION 7 — Final verification
# =============================================================================
section "System verification"

ALL_OK=true

check_endpoint() {
    local label="$1" url="$2" expected="$3"
    HTTP=$(curl -s -o /dev/null -w "%{http_code}" "$url" 2>/dev/null || true)
    if [[ "$HTTP" == "$expected" ]]; then
        ok "$label ($url)"
    else
        error "$label — expected HTTP $expected, got $HTTP ($url)"
        ALL_OK=false
    fi
}

check_endpoint "Dashboard"   "http://localhost:8888"       "200"
check_endpoint "Neo4j"       "http://localhost:7474"       "200"
check_endpoint "DefectDojo"  "http://localhost:8080/login" "200"
check_endpoint "Runner"      "http://localhost:9999/preflight/aws" "200"

# Verify runner preflight details
info "Checking runner preflight status..."
AWS_PREFLIGHT=$(curl -s http://localhost:9999/preflight/aws 2>/dev/null || true)
CART_PREFLIGHT=$(curl -s http://localhost:9999/preflight/cartography 2>/dev/null || true)

echo ""
echo "  AWS Pipeline preflight:"
echo "$AWS_PREFLIGHT" | python3 -c "
import sys, json
checks = json.load(sys.stdin)
for c in checks:
    icon = '  ✓' if c['ok'] else '  ✗'
    print(f\"    {icon}  {c['label']}\")
" 2>/dev/null || echo "    (could not parse)"

echo ""
echo "  Cartography preflight:"
echo "$CART_PREFLIGHT" | python3 -c "
import sys, json
checks = json.load(sys.stdin)
for c in checks:
    icon = '  ✓' if c['ok'] else '  ✗'
    print(f\"    {icon}  {c['label']}\")
" 2>/dev/null || echo "    (could not parse)"

# =============================================================================
# SECTION 8 — Summary
# =============================================================================
section "Setup complete"

echo ""
echo -e "  ${GREEN}${BOLD}All services are running.${RESET}"
echo ""
echo "  ┌─────────────────────────────────────────────────────┐"
echo "  │  Dashboard      →  http://localhost:8888            │"
echo "  │  DefectDojo     →  http://localhost:8080            │"
echo "  │  Neo4j Browser  →  http://localhost:7474            │"
echo "  │  Runner API     →  http://localhost:9999            │"
echo "  └─────────────────────────────────────────────────────┘"
echo ""
echo "  Next steps:"
echo "    1. Open http://localhost:8888 → ⚙ Account Setup — enter AWS credentials"
echo "    2. Open http://localhost:8888 → ▶ Run Pipeline"
echo "    3. Run Cartography for AWS graph discovery (after credentials are set):"
echo ""
echo "       export \$(grep -v '^#' $ENV_FILE | xargs)"
echo "       $CARTOGRAPHY \\"
echo "         --neo4j-uri bolt://localhost:7687 \\"
echo "         --neo4j-user neo4j \\"
echo "         --neo4j-password-env-var NEO4J_PASSWORD \\"
echo "         --neo4j-database neo4j \\"
echo "         --aws-best-effort-mode \\"
echo "         --aws-regions \"\$AWS_DEFAULT_REGION\" \\"
echo "         --aws-requested-syncs \"s3,ec2:instance,ec2:security_group,rds,lambda_function\""
echo ""
echo "  Runner log:  $RUNNER_LOG"
echo "  Env file:    $ENV_FILE"
echo ""

if [[ "$ALL_OK" != "true" ]]; then
    warn "Some checks failed — review errors above before running pipelines."
    exit 1
fi