setup.sh

#!/usr/bin/env bash
#
# setup.sh — Claude-IR Quick Start Setup
#
# Installs core MCPs and databases with zero prompts.
# Run again with flags to add optional components.
#
# Usage:
#   ./setup.sh                              # Install core MCPs
#   ./setup.sh --opencti                    # Add OpenCTI (guided)
#   ./setup.sh --remnux                     # Add REMnux (guided)
#   ./setup.sh --help                       # Full usage
#

set -o pipefail

# =============================================================================
# Constants
# =============================================================================

GITHUB_ORG="https://github.com/scriptedstatement"
PROJECT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Triage data source repos (for --build-from-source)
TRIAGE_DATA_REPOS=(
    "https://github.com/AndrewRathbun/VanillaWindowsReference.git"
    "https://github.com/LOLBAS-Project/LOLBAS.git"
    "https://github.com/magicsword-io/LOLDrivers.git"
    "https://github.com/wietze/HijackLibs.git"
)
TRIAGE_REGISTRY_REPO="https://github.com/AndrewRathbun/VanillaWindowsRegistryHives.git"

# =============================================================================
# Color Setup
# =============================================================================

setup_colors() {
    if [[ -t 1 ]] && [[ "${TERM:-dumb}" != "dumb" ]]; then
        RED='\033[0;31m'
        GREEN='\033[0;32m'
        YELLOW='\033[1;33m'
        BLUE='\033[0;34m'
        BOLD='\033[1m'
        DIM='\033[2m'
        NC='\033[0m'
    else
        RED='' GREEN='' YELLOW='' BLUE='' BOLD='' DIM='' NC=''
    fi
}

# =============================================================================
# Helpers
# =============================================================================

info()    { echo -e "  ${GREEN}[OK]${NC} $1"; }
skip()    { echo -e "  ${DIM}[SKIP]${NC} $1"; }
warn()    { echo -e "  ${YELLOW}[WARN]${NC} $1"; }
fail()    { echo -e "  ${RED}[FAIL]${NC} $1"; }
step()    { echo -e "\n${BOLD}$1${NC}"; }
die()     { echo -e "\n${RED}Error: $1${NC}" >&2; exit 1; }

handle_error() {
    local msg="$1"
    fail "$msg"
    # In non-interactive contexts (no terminal), abort on errors
    if [[ ! -t 0 ]]; then
        echo "  Non-interactive mode: aborting on error."
        exit 1
    fi
    echo ""
    read -p "  Continue anyway? [Y/n] " -n 1 -r
    echo
    if [[ "$REPLY" =~ ^[Nn]$ ]]; then
        echo "Aborted."
        exit 1
    fi
    return 0  # User chose to continue
}

get_mcp_value() {
    # Read a credential value from existing .mcp.json
    # Returns empty string if not found or if value is a placeholder
    # Uses env vars to pass values to Python (avoids shell interpolation in code)
    local key="$1"
    [[ -f "$PROJECT_DIR/.mcp.json" ]] || return 0
    _MCP_CONFIG="$PROJECT_DIR/.mcp.json" _MCP_KEY="$key" python3 -c "
import json, re, os, sys
PLACEHOLDERS = {'YOUR_OPENCTI_TOKEN','YOUR_TOKEN_HERE','your-api-token-here',
                'your-token-here','YOUR_OPENCTI_URL','YOUR_REMNUX_HOST','YOUR_REMNUX_TOKEN'}
try:
    cfg = json.load(open(os.environ['_MCP_CONFIG']))
    mcps = cfg.get('mcpServers', {})
    k = os.environ['_MCP_KEY']
    v = ''
    if k in ('OPENCTI_URL','OPENCTI_TOKEN','OPENCTI_SSL_VERIFY'):
        v = mcps.get('opencti-mcp', {}).get('env', {}).get(k, '')
    elif k == 'REMNUX_HOST':
        url = mcps.get('remnux-mcp', {}).get('url', '')
        m = re.match(r'https?://([^:/]+)', url)
        v = m.group(1) if m else ''
    elif k == 'REMNUX_PORT':
        url = mcps.get('remnux-mcp', {}).get('url', '')
        if '://' in url:
            m = re.search(r':(\d+)', url.split('://',1)[-1])
            v = m.group(1) if m else ''
    elif k == 'REMNUX_TOKEN':
        auth = mcps.get('remnux-mcp',{}).get('headers',{}).get('Authorization','')
        v = auth[7:] if auth.startswith('Bearer ') else ''
    if v in PLACEHOLDERS or 'YOUR_' in v:
        v = ''
    print(v)
except: pass
" 2>/dev/null
}

# =============================================================================
# Argument Parsing
# =============================================================================

show_help() {
    cat <<'HELP'
Usage: ./setup.sh [OPTIONS]

Quick start (no options needed):
  ./setup.sh                          Install core MCPs with sensible defaults

Add optional MCPs (interactive — guides you through configuration):
  --opencti                           Add OpenCTI threat intelligence
  --remnux                            Add REMnux malware analysis

Database options:
  --build-from-source                 Build triage DBs from source instead of downloading
  --with-registry                     Include full registry baseline (~12GB, built from source)
  --rag-no-bundled                    Build RAG without bundled SANS/AppliedIR content

Other:
  --force-reprompt                    Re-prompt for tokens even if already configured
  --skip-tools-inventory              Skip forensic tools detection
  --skip-hooks                        Skip hooks and /case command installation
  --help                              Show this help

Notes:
  To remove an optional MCP, re-run setup.sh without that flag.
  Running ./setup.sh (no flags) keeps only the 4 core MCPs in .mcp.json.
HELP
}

parse_args() {
    # Defaults
    OPENCTI_URL=""
    OPENCTI_TOKEN=""
    OPENCTI_SSL_NO_VERIFY=false
    REMNUX_HOST=""
    REMNUX_PORT="3000"
    REMNUX_TOKEN=""
    WANT_OPENCTI=false
    WANT_REMNUX=false
    BUILD_FROM_SOURCE=false
    WITH_REGISTRY=false
    RAG_NO_BUNDLED=false
    FORCE_REPROMPT=false
    SKIP_TOOLS_INVENTORY=false
    SKIP_HOOKS=false

    while [[ $# -gt 0 ]]; do
        case "$1" in
            --opencti)
                WANT_OPENCTI=true; shift ;;
            --remnux)
                WANT_REMNUX=true; shift ;;
            --build-from-source)
                BUILD_FROM_SOURCE=true; shift ;;
            --with-registry)
                WITH_REGISTRY=true; shift ;;
            --rag-no-bundled)
                RAG_NO_BUNDLED=true; shift ;;
            --force-reprompt)
                FORCE_REPROMPT=true; shift ;;
            --skip-tools-inventory)
                SKIP_TOOLS_INVENTORY=true; shift ;;
            --skip-hooks)
                SKIP_HOOKS=true; shift ;;
            --help|-h)
                show_help; exit 0 ;;
            *)
                die "Unknown option: $1\nRun './setup.sh --help' for usage." ;;
        esac
    done
}

# =============================================================================
# Prerequisites
# =============================================================================

check_project_dir() {
    if [[ ! -f "$PROJECT_DIR/CLAUDE.md" ]] || [[ ! -f "$PROJECT_DIR/.mcp.json.example" ]]; then
        die "Must run from claude-ir project root.\nExpected CLAUDE.md and .mcp.json.example in: $PROJECT_DIR"
    fi
}

check_prerequisites() {
    step "Checking prerequisites..."

    local missing=()

    # git
    if command -v git &>/dev/null; then
        info "git $(git --version | cut -d' ' -f3)"
    else
        missing+=("git")
    fi

    # python3
    if command -v python3 &>/dev/null; then
        local pyver
        pyver=$(python3 -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')
        local pymajor pyminor
        pymajor=$(echo "$pyver" | cut -d. -f1)
        pyminor=$(echo "$pyver" | cut -d. -f2)
        if [[ "$pymajor" -gt 3 ]] || { [[ "$pymajor" -eq 3 ]] && [[ "$pyminor" -ge 10 ]]; }; then
            info "python3 $pyver"
        else
            die "Python 3.10+ required (found $pyver)"
        fi
    else
        missing+=("python3 (3.10+)")
    fi

    # pip (via python3)
    if python3 -m pip --version &>/dev/null 2>&1; then
        info "pip available"
    else
        missing+=("pip (python3 -m pip)")
    fi

    # Network (quick check)
    if curl -sf --max-time 5 "https://github.com" >/dev/null 2>&1; then
        info "Network connectivity"
    elif wget -q --timeout=5 --spider "https://github.com" 2>/dev/null; then
        info "Network connectivity"
    else
        warn "Network check failed (may still work on local network)"
    fi

    if [[ ${#missing[@]} -gt 0 ]]; then
        die "Missing prerequisites: ${missing[*]}"
    fi
}

# =============================================================================
# Interactive Configuration Wizards
# =============================================================================

wizard_opencti() {
    step "OpenCTI Configuration"
    echo ""
    echo "  Connect to your OpenCTI threat intelligence platform."
    echo "  You'll need the URL and an API token."
    echo ""

    # URL
    local default_url="http://localhost:8080"
    local existing_url
    existing_url=$(get_mcp_value "OPENCTI_URL")
    if [[ -n "$existing_url" ]]; then
        default_url="$existing_url"
    fi

    read -p "  OpenCTI URL [$default_url]: " input_url
    OPENCTI_URL="${input_url:-$default_url}"

    # Validate URL format
    if [[ ! "$OPENCTI_URL" =~ ^https?:// ]]; then
        OPENCTI_URL="http://${OPENCTI_URL}"
        warn "Added http:// prefix → $OPENCTI_URL"
    fi

    # Basic URL sanity check
    if [[ ! "$OPENCTI_URL" =~ ^https?://[a-zA-Z0-9._:-]+(/.*)?$ ]]; then
        warn "URL looks malformed: $OPENCTI_URL"
    fi

    # Self-signed cert check (only for HTTPS)
    if [[ "$OPENCTI_URL" =~ ^https:// ]]; then
        echo ""
        read -p "  Does it use a self-signed certificate? [y/N]: " -n 1 -r
        echo
        if [[ "$REPLY" =~ ^[Yy]$ ]]; then
            OPENCTI_SSL_NO_VERIFY=true
            info "SSL verification disabled (self-signed cert)"
        fi
    fi

    # HTTP warning for non-localhost
    if [[ "$OPENCTI_URL" =~ ^http:// ]] && [[ ! "$OPENCTI_URL" =~ localhost|127\.0\.0\.1 ]]; then
        echo ""
        warn "This URL uses unencrypted HTTP."
        echo "       API tokens will be sent in cleartext."
        echo "       This is acceptable on isolated lab networks."
        echo "       For anything else, use HTTPS."
        echo ""
    fi

    # Token
    local existing_token
    existing_token=$(get_mcp_value "OPENCTI_TOKEN")
    if [[ -n "$existing_token" ]] && [[ "$FORCE_REPROMPT" != true ]]; then
        OPENCTI_TOKEN="$existing_token"
        skip "Token already configured"
    else
        echo ""
        echo "  Find your token at:"
        echo "    OpenCTI web UI > Settings (gear) > Security > API Access"
        echo ""
        read -s -p "  API token: " opencti_token
        echo
        if [[ -n "$opencti_token" ]]; then
            OPENCTI_TOKEN="$opencti_token"
            info "Token accepted"
        else
            warn "No token entered — re-run with --opencti to add it later"
        fi
    fi

    info "OpenCTI URL: $OPENCTI_URL"
    echo ""
}

wizard_remnux() {
    step "REMnux Configuration"
    echo ""
    echo "  Connect to your REMnux instance for malware analysis."
    echo "  The REMnux MCP server must be running on the remote machine."
    echo ""

    # Host
    local default_host=""
    local existing_host
    existing_host=$(get_mcp_value "REMNUX_HOST")
    if [[ -n "$existing_host" ]]; then
        default_host="$existing_host"
    fi

    if [[ -n "$default_host" ]]; then
        read -p "  REMnux host/IP [$default_host]: " input_host
        REMNUX_HOST="${input_host:-$default_host}"
    else
        while [[ -z "$REMNUX_HOST" ]]; do
            read -p "  REMnux host/IP: " REMNUX_HOST
            if [[ -z "$REMNUX_HOST" ]]; then
                warn "Host is required"
            fi
        done
    fi

    # Port
    local default_port="3000"
    local existing_port
    existing_port=$(get_mcp_value "REMNUX_PORT")
    if [[ -n "$existing_port" ]]; then
        default_port="$existing_port"
    fi
    read -p "  MCP server port [$default_port]: " input_port
    REMNUX_PORT="${input_port:-$default_port}"
    if ! [[ "$REMNUX_PORT" =~ ^[0-9]+$ ]] || [[ "$REMNUX_PORT" -lt 1 ]] || [[ "$REMNUX_PORT" -gt 65535 ]]; then
        warn "Invalid port '$REMNUX_PORT', using default $default_port"
        REMNUX_PORT="$default_port"
    fi

    # Token
    local existing_token
    existing_token=$(get_mcp_value "REMNUX_TOKEN")
    if [[ -n "$existing_token" ]] && [[ "$FORCE_REPROMPT" != true ]]; then
        REMNUX_TOKEN="$existing_token"
        skip "Token already configured"
    else
        echo ""
        echo "  The bearer token was set when you started the REMnux MCP server."
        echo "  Check the server's startup command or environment."
        echo ""
        read -s -p "  Bearer token: " remnux_token
        echo
        if [[ -n "$remnux_token" ]]; then
            REMNUX_TOKEN="$remnux_token"
            info "Token accepted"
        else
            warn "No token entered — re-run with --remnux to add it later"
        fi
    fi

    info "REMnux: ${REMNUX_HOST}:${REMNUX_PORT}"
    echo ""
}

# =============================================================================
# Component Determination
# =============================================================================

determine_components() {
    # Always included
    WANT_RAG=true
    WANT_TRIAGE=true
    WANT_MSLEARN=true
    WANT_ZELTSER=true

    # WANT_OPENCTI / WANT_REMNUX already set by parse_args

    # Run interactive wizards to collect connection details
    if [[ "$WANT_OPENCTI" == true ]]; then
        wizard_opencti
    fi
    if [[ "$WANT_REMNUX" == true ]]; then
        wizard_remnux
    fi

    # Build --only list for generate-mcp-config.py
    ONLY_LIST="forensic-rag,windows-triage,microsoft-learn,zeltser-ir-writing"
    if [[ "$WANT_OPENCTI" == true ]]; then
        ONLY_LIST="${ONLY_LIST},opencti-mcp"
    fi
    if [[ "$WANT_REMNUX" == true ]]; then
        ONLY_LIST="${ONLY_LIST},remnux-mcp"
    fi
}

# =============================================================================
# Phase: Install Local MCPs
# =============================================================================

install_local_mcp() {
    local name="$1" repo_dir="$2" note="${3:-}"
    local repo_url="${GITHUB_ORG}/${repo_dir}.git"

    echo ""
    echo -e "  ${BOLD}${name}${NC}${note:+ ${DIM}(${note})${NC}}"

    # Clone
    if [[ -d "$PROJECT_DIR/$repo_dir" ]]; then
        skip "Already cloned: ${repo_dir}/"
    else
        echo -n "    Cloning ${repo_dir}..."
        if git clone "$repo_url" "$PROJECT_DIR/$repo_dir" >/dev/null 2>&1; then
            echo -e " ${GREEN}done${NC}"
        else
            handle_error "Failed to clone $repo_url"
            # User chose continue — skip venv/pip since dir doesn't exist
            return 0
        fi
    fi

    # Create venv
    if [[ -d "$PROJECT_DIR/$repo_dir/.venv" ]]; then
        skip "Venv exists: ${repo_dir}/.venv/"
    else
        echo -n "    Creating venv..."
        if python3 -m venv "$PROJECT_DIR/$repo_dir/.venv" 2>/dev/null; then
            echo -e " ${GREEN}done${NC}"
        else
            handle_error "Failed to create venv for $repo_dir (is python3-venv installed?)"
            # User chose continue — skip pip since venv doesn't exist
            return 0
        fi
    fi

    # Install dependencies
    # Check if already installed by trying to import the main module
    local venv_python="$PROJECT_DIR/$repo_dir/.venv/bin/python"
    local check_module
    case "$repo_dir" in
        forensic-rag-mcp)  check_module="rag_mcp" ;;
        windows-triage-mcp) check_module="windows_triage" ;;
        opencti-mcp)       check_module="opencti_mcp" ;;
        *)                 check_module="" ;;
    esac

    if [[ -n "$check_module" ]] && "$venv_python" -c "import $check_module" 2>/dev/null; then
        skip "Dependencies installed: ${repo_dir}"
    else
        echo "    Installing dependencies..."
        "$venv_python" -m pip install --upgrade pip >/dev/null 2>&1
        if "$venv_python" -m pip install -e "$PROJECT_DIR/$repo_dir/" >/dev/null 2>&1; then
            info "Dependencies installed: ${repo_dir}"
        else
            handle_error "pip install failed for $repo_dir"
        fi
    fi
}

phase_install_local_mcps() {
    step "Installing local MCPs..."

    local count=0 total=2
    [[ "$WANT_OPENCTI" == true ]] && ((total++))

    ((count++))
    echo -e "\n  ${DIM}[${count}/${total}]${NC}"
    install_local_mcp "forensic-rag" "forensic-rag-mcp" "includes ML libraries, ~500MB"

    ((count++))
    echo -e "\n  ${DIM}[${count}/${total}]${NC}"
    install_local_mcp "windows-triage" "windows-triage-mcp"

    if [[ "$WANT_OPENCTI" == true ]]; then
        ((count++))
        echo -e "\n  ${DIM}[${count}/${total}]${NC}"
        install_local_mcp "opencti" "opencti-mcp"
    fi
}

# =============================================================================
# Phase: Generate Config
# =============================================================================

phase_generate_config() {
    step "Generating MCP configuration..."

    local args=("--only" "$ONLY_LIST" "--force")

    # URLs/hosts pass as CLI args (not sensitive)
    if [[ -n "$OPENCTI_URL" ]]; then
        args+=("--opencti-url" "$OPENCTI_URL")
    fi
    if [[ "$OPENCTI_SSL_NO_VERIFY" == true ]]; then
        args+=("--opencti-ssl-no-verify")
    fi
    if [[ -n "$REMNUX_HOST" ]]; then
        args+=("--remnux-host" "$REMNUX_HOST")
        [[ "$REMNUX_PORT" != "3000" ]] && args+=("--remnux-port" "$REMNUX_PORT")
    fi

    # Tokens pass as env vars (not visible in ps)
    if [[ -n "${OPENCTI_TOKEN:-}" ]]; then
        export OPENCTI_TOKEN
    fi
    if [[ -n "${REMNUX_TOKEN:-}" ]]; then
        export REMNUX_TOKEN
    fi

    if python3 "$PROJECT_DIR/scripts/generate-mcp-config.py" "${args[@]}"; then
        info "Generated .mcp.json"
    else
        handle_error "Failed to generate .mcp.json"
    fi
}

# =============================================================================
# Phase: Database Setup
# =============================================================================

phase_database_setup() {
    step "Setting up databases..."

    # --- RAG Index ---
    if [[ "$WANT_RAG" == true ]]; then
        if [[ ! -d "$PROJECT_DIR/forensic-rag-mcp/.venv" ]]; then
            skip "RAG index: forensic-rag-mcp not installed"
        elif [[ -d "$PROJECT_DIR/forensic-rag-mcp/data/chroma" ]] && \
             [[ -n "$(ls -A "$PROJECT_DIR/forensic-rag-mcp/data/chroma" 2>/dev/null)" ]]; then
            skip "RAG index already built: forensic-rag-mcp/data/chroma/"
        else
            echo ""
            echo -e "  ${BOLD}RAG index${NC} ${DIM}(building in background)${NC}"
            local build_cmd="$PROJECT_DIR/forensic-rag-mcp/.venv/bin/python -m rag_mcp.build"
            if [[ "$RAG_NO_BUNDLED" == true ]]; then
                build_cmd="$build_cmd --no-bundled"
            fi
            # Store PID for summary
            nohup bash -c 'cd "$1" && shift && "$@"' _ "$PROJECT_DIR/forensic-rag-mcp" $build_cmd > /tmp/rag-build.log 2>&1 &
            RAG_BUILD_PID=$!
            ( umask 077; echo "$RAG_BUILD_PID" > /tmp/rag-build.pid )
            info "RAG build started (PID $RAG_BUILD_PID, log: /tmp/rag-build.log)"
        fi
    fi

    # --- Triage Databases ---
    if [[ "$WANT_TRIAGE" == true ]]; then
        local triage_data="$PROJECT_DIR/windows-triage-mcp/data"
        if [[ ! -d "$PROJECT_DIR/windows-triage-mcp" ]]; then
            skip "Triage databases: windows-triage-mcp not installed"
        elif [[ -f "$triage_data/known_good.db" ]] && [[ -f "$triage_data/context.db" ]]; then
            skip "Triage databases already exist"
        elif [[ "$BUILD_FROM_SOURCE" == true ]]; then
            # Build ALL from source (includes registry extraction)
            echo ""
            echo -e "  ${BOLD}Triage databases${NC} ${DIM}(building from source in background)${NC}"

            # Clone data source repos
            local sources_dir="$triage_data/sources"
            mkdir -p "$sources_dir"
            for repo_url in "${TRIAGE_DATA_REPOS[@]}"; do
                local repo_name
                repo_name=$(basename "$repo_url" .git)
                if [[ -d "$sources_dir/$repo_name" ]]; then
                    skip "Data source: $repo_name"
                else
                    echo -n "    Cloning $repo_name..."
                    if git clone --depth 1 "$repo_url" "$sources_dir/$repo_name" >/dev/null 2>&1; then
                        echo -e " ${GREEN}done${NC}"
                    else
                        handle_error "Failed to clone $repo_name"
                    fi
                fi
            done

            # Registry repo (optional, ~12GB)
            if [[ "$WITH_REGISTRY" == true ]]; then
                local reg_name="VanillaWindowsRegistryHives"
                if [[ -d "$sources_dir/$reg_name" ]]; then
                    skip "Data source: $reg_name"
                else
                    echo -n "    Cloning $reg_name (large, ~12GB)..."
                    if git clone --depth 1 "$TRIAGE_REGISTRY_REPO" "$sources_dir/$reg_name" >/dev/null 2>&1; then
                        echo -e " ${GREEN}done${NC}"
                    else
                        handle_error "Failed to clone $reg_name"
                    fi
                fi
            fi

            # Run import
            local import_cmd="$PROJECT_DIR/windows-triage-mcp/.venv/bin/python scripts/import_all.py"
            if [[ "$WITH_REGISTRY" != true ]]; then
                import_cmd="$import_cmd --skip-registry"
            fi
            nohup bash -c 'cd "$1" && shift && "$@"' _ "$PROJECT_DIR/windows-triage-mcp" $import_cmd > /tmp/triage-import.log 2>&1 &
            TRIAGE_BUILD_PID=$!
            ( umask 077; echo "$TRIAGE_BUILD_PID" > /tmp/triage-import.pid )
            info "Triage import started (PID $TRIAGE_BUILD_PID, log: /tmp/triage-import.log)"
        else
            # Download pre-built databases (default)
            echo ""
            echo -e "  ${BOLD}Triage databases${NC} ${DIM}(downloading pre-built)${NC}"
            if bash "$PROJECT_DIR/windows-triage-mcp/scripts/download-databases.sh"; then
                info "Triage databases downloaded"
            else
                handle_error "Database download failed — try --build-from-source"
            fi
        fi

        # --- Optional: Full Registry Baseline ---
        # Separate from core DBs — downloads core + builds registry from source
        if [[ "$WITH_REGISTRY" == true ]] && [[ "$BUILD_FROM_SOURCE" != true ]]; then
            local reg_db="$triage_data/known_good_registry.db"
            if [[ -f "$reg_db" ]]; then
                skip "Registry baseline already exists"
            elif [[ ! -d "$PROJECT_DIR/windows-triage-mcp/.venv" ]]; then
                skip "Registry baseline: windows-triage-mcp venv not ready"
            else
                echo ""
                echo -e "  ${BOLD}Registry baseline${NC} ${DIM}(building from source in background)${NC}"

                # Clone registry repo
                local sources_dir="$triage_data/sources"
                local reg_name="VanillaWindowsRegistryHives"
                mkdir -p "$sources_dir"
                if [[ -d "$sources_dir/$reg_name" ]]; then
                    skip "Data source: $reg_name"
                else
                    echo -n "    Cloning $reg_name (large, ~12GB)..."
                    if git clone --depth 1 "$TRIAGE_REGISTRY_REPO" "$sources_dir/$reg_name" >/dev/null 2>&1; then
                        echo -e " ${GREEN}done${NC}"
                    else
                        handle_error "Failed to clone $reg_name"
                    fi
                fi

                # Init + import registry DB in background
                local py="$PROJECT_DIR/windows-triage-mcp/.venv/bin/python"
                nohup bash -c '
                    cd "$1" && shift
                    "$1" scripts/init_registry_db.py --force &&
                    "$1" scripts/import_registry_full.py
                ' _ "$PROJECT_DIR/windows-triage-mcp" "$py" > /tmp/registry-build.log 2>&1 &
                REG_BUILD_PID=$!
                ( umask 077; echo "$REG_BUILD_PID" > /tmp/registry-build.pid )
                info "Registry import started (PID $REG_BUILD_PID, log: /tmp/registry-build.log)"
            fi
        fi
    fi
}

# =============================================================================
# Phase: Hooks & Skill
# =============================================================================

phase_hooks() {
    if [[ "$SKIP_HOOKS" == true ]]; then
        skip "Hooks installation (--skip-hooks)"
        return 0
    fi

    step "Installing hooks and /case command..."

    if [[ -f "$PROJECT_DIR/scripts/install-hooks-and-skill.sh" ]]; then
        if ! bash "$PROJECT_DIR/scripts/install-hooks-and-skill.sh" --non-interactive; then
            warn "Hooks installation had errors (non-fatal)"
        fi
    else
        warn "install-hooks-and-skill.sh not found — skipping"
    fi
}

# =============================================================================
# Phase: Tools Inventory
# =============================================================================

phase_tools_inventory() {
    if [[ "$SKIP_TOOLS_INVENTORY" == true ]]; then
        skip "Tools inventory (--skip-tools-inventory)"
        return 0
    fi

    if [[ -f "$PROJECT_DIR/local/FORENSIC_TOOLS.md" ]]; then
        skip "Tools inventory already exists: local/FORENSIC_TOOLS.md"
        return 0
    fi

    step "Detecting forensic tools..."

    mkdir -p "$PROJECT_DIR/local"

    # Detect environment
    local env_type="Linux" os_info=""
    if [[ "$(uname)" == "Darwin" ]]; then
        env_type="macOS"
        os_info=$(sw_vers -productName 2>/dev/null && sw_vers -productVersion 2>/dev/null | tr '\n' ' ' || uname -a)
    elif [[ -f /proc/version ]] && grep -qi microsoft /proc/version 2>/dev/null; then
        env_type="WSL"
        os_info=$(lsb_release -d 2>/dev/null | cut -f2 || uname -a)
    else
        os_info=$(lsb_release -d 2>/dev/null | cut -f2 || uname -a)
    fi

    # Tool categories and binaries to check
    # Format: "Category|tool1,tool2,tool3"
    local categories=(
        "Memory Analysis|vol,volatility3,volatility2,rekall,strings,bulk_extractor"
        "Zimmerman Tools|MFTECmd,EvtxECmd,RECmd,AmcacheParser,AppCompatCacheParser,JLECmd,LECmd,SBECmd,RBCmd,WxTCmd,SQLECmd,bstrings"
        "Timeline|log2timeline.py,psort.py,pinfo.py,mactime,timesketch"
        "Sleuth Kit|mmls,fls,icat,fsstat,blkcat,blkls,blkstat,ifind,ffind,img_stat,img_cat,istat,jcat,jls,hfind,tsk_recover,tsk_gettimes,tsk_loaddb,sorter,sigfind,fiwalk"
        "Disk Imaging|ewfinfo,ewfverify,ewfmount,ewfacquire,affinfo,affcat,affconvert,affcompare,qemu-img,qemu-nbd,xmount,dc3dd,dcfldd,ddrescue"
        "Registry|rip.pl,regfinfo,regfexport,regipy"
        "Windows Artifacts|esedbinfo,esedbexport,evtinfo,evtexport,evtxinfo,evtxexport,fvdeinfo,fvdemount,pffinfo,pffexport,vshadowinfo,vshadowmount,readpst"
        "Malware Analysis|clamscan,freshclam,readpe,pehash,pescan,pestr,r2,radare2,upx,ssdeep,floss"
        "Hash Tools|md5deep,sha256deep,hashdeep,md5sum,sha256sum,shasum"
        "Crypto|aeskeyfind,rsakeyfind,cryptsetup,dislocker,ccrypt"
        "Recovery & Carving|photorec,testdisk,extundelete,foremost,scalpel,bulk_extractor"
        "Network Forensics|tcpdump,tcpflow,wireshark,tshark,ngrep,tcpreplay,tcpstat,tcptrace,tcptrack,tcpxtract,tcpick,netcat,nc,dsniff,ettercap,arp-scan,nbtscan,nfdump"
        "File Analysis|exiftool,exif,file,strings,xxd,hexdump,hexedit,ghex"
        "Utilities|sqlite3,jq,7z,p7zip,unrar,cabextract,autopsy"
        "Optional|hayabusa,chainsaw,yara,tshark,binwalk,hashcat,john"
    )

    local total_found=0
    local total_missing=0
    local cat_count=0
    local outfile="$PROJECT_DIR/local/FORENSIC_TOOLS.md"
    local tmpfile
    tmpfile=$(mktemp "$PROJECT_DIR/local/.FORENSIC_TOOLS.md.XXXXXX")

    # Write header
    {
        echo "# Forensic Tools Inventory"
        echo ""
        echo "**Auto-detected** by setup.sh on $(date -u '+%Y-%m-%d %H:%M UTC')"
        echo "**Environment:** ${env_type}"
        echo "**OS:** ${os_info}"
        echo ""
    } > "$tmpfile"

    for entry in "${categories[@]}"; do
        local category="${entry%%|*}"
        local tools_str="${entry#*|}"
        IFS=',' read -ra tools <<< "$tools_str"

        local found_lines=()
        local missing_names=()
        for tool in "${tools[@]}"; do
            if command -v "$tool" &>/dev/null; then
                local tpath
                tpath=$(command -v "$tool")
                found_lines+=("- [x] ${tool} → ${tpath}")
                ((total_found++))
            else
                missing_names+=("$tool")
                ((total_missing++))
            fi
        done

        # Write category if any tools were checked (show found AND missing)
        if [[ ${#found_lines[@]} -gt 0 ]] || [[ ${#missing_names[@]} -gt 0 ]]; then
            ((cat_count++))
            echo "## ${category}" >> "$tmpfile"
            echo "" >> "$tmpfile"
            for line in "${found_lines[@]}"; do
                echo "$line" >> "$tmpfile"
            done
            for m in "${missing_names[@]}"; do
                echo "- [ ] ${m} — not found" >> "$tmpfile"
            done
            echo "" >> "$tmpfile"
        fi
    done

    # Windows-only tools section (expected missing on Linux/macOS)
    {
        echo "## Windows-Only (Not Available on Linux/macOS)"
        echo ""
        echo "These tools only run on Windows. Use workarounds:"
        echo "- PECmd → use \`log2timeline.py --parsers prefetch\`"
        echo "- SrumECmd → use \`esedbexport\` on SRUDB.dat"
        echo ""
    } >> "$tmpfile"

    # Environment-specific notes
    {
        echo "## Environment Notes"
        echo ""
        case "$env_type" in
            WSL)
                echo "- **WSL detected** — Some .NET-based tools may have issues"
                echo "- PECmd and SrumECmd are Windows-only (see workarounds above)"
                echo "- GUI tools (wireshark, ghex, autopsy) require X11/WSLg"
                ;;
            macOS)
                echo "- **macOS detected** — Many Windows artifact tools are Linux-only"
                echo "- Zimmerman Tools require .NET runtime (may not work natively)"
                echo "- Focus: memory analysis, timeline, Sleuth Kit, general utilities"
                echo "- Use \`shasum -a 256\` instead of \`sha256sum\`"
                ;;
            *)
                echo "- Standard Linux environment"
                ;;
        esac
        echo ""
        echo "---"
        echo ""
        echo "**Total:** ${total_found} tools found, ${total_missing} not found across ${cat_count} categories"
        echo ""
        echo "*Claude will review this inventory on first launch and suggest improvements.*"
    } >> "$tmpfile"

    # Atomic move — partial file on interrupt won't corrupt outfile
    mv "$tmpfile" "$outfile"

    info "${total_found} tools found across ${cat_count} categories → local/FORENSIC_TOOLS.md"
}

# =============================================================================
# Phase: Summary
# =============================================================================

phase_summary() {
    # Track status for each component
    local rag_status triage_status opencti_status
    local mslearn_status zeltser_status

    # RAG
    if [[ -d "$PROJECT_DIR/forensic-rag-mcp/.venv" ]]; then
        if [[ -d "$PROJECT_DIR/forensic-rag-mcp/data/chroma" ]] && \
           [[ -n "$(ls -A "$PROJECT_DIR/forensic-rag-mcp/data/chroma" 2>/dev/null)" ]]; then
            rag_status="cloned, deps installed, index built"
        elif [[ -n "${RAG_BUILD_PID:-}" ]] && kill -0 "$RAG_BUILD_PID" 2>/dev/null; then
            rag_status="cloned, deps installed, building (PID $RAG_BUILD_PID)"
        else
            rag_status="cloned, deps installed"
        fi
    elif [[ -d "$PROJECT_DIR/forensic-rag-mcp" ]]; then
        rag_status="cloned"
    else
        rag_status="not installed"
    fi

    # Triage
    if [[ -d "$PROJECT_DIR/windows-triage-mcp/.venv" ]]; then
        if [[ -f "$PROJECT_DIR/windows-triage-mcp/data/known_good.db" ]]; then
            triage_status="cloned, deps installed, DBs ready"
        elif [[ -n "${TRIAGE_BUILD_PID:-}" ]] && kill -0 "$TRIAGE_BUILD_PID" 2>/dev/null; then
            triage_status="cloned, deps installed, importing (PID $TRIAGE_BUILD_PID)"
        else
            triage_status="cloned, deps installed"
        fi
    elif [[ -d "$PROJECT_DIR/windows-triage-mcp" ]]; then
        triage_status="cloned"
    else
        triage_status="not installed"
    fi

    # OpenCTI
    if [[ "$WANT_OPENCTI" == true ]]; then
        if [[ -d "$PROJECT_DIR/opencti-mcp/.venv" ]]; then
            opencti_status="cloned, deps installed"
        elif [[ -d "$PROJECT_DIR/opencti-mcp" ]]; then
            opencti_status="cloned"
        else
            opencti_status="not installed"
        fi
    fi

    # Remote MCPs (always configured if selected)
    mslearn_status="configured (remote)"
    zeltser_status="configured (remote)"

    # Print summary
    echo ""
    echo -e "${BOLD}══════════════════════════════════════════════════════${NC}"
    echo -e "${BOLD}  Claude-IR Setup Complete${NC}"
    echo -e "${BOLD}══════════════════════════════════════════════════════${NC}"
    echo ""
    echo "  Installed:"

    # Checkboxes
    if [[ "$rag_status" != "not installed" ]]; then
        echo -e "    ${GREEN}[x]${NC} forensic-rag          $rag_status"
    else
        echo -e "    ${RED}[ ]${NC} forensic-rag          $rag_status"
    fi

    if [[ "$triage_status" != "not installed" ]]; then
        echo -e "    ${GREEN}[x]${NC} windows-triage        $triage_status"
    else
        echo -e "    ${RED}[ ]${NC} windows-triage        $triage_status"
    fi

    echo -e "    ${GREEN}[x]${NC} microsoft-learn       $mslearn_status"
    echo -e "    ${GREEN}[x]${NC} zeltser-ir-writing    $zeltser_status"

    if [[ "$WANT_OPENCTI" == true ]]; then
        if [[ -d "$PROJECT_DIR/opencti-mcp/.venv" ]]; then
            echo -e "    ${GREEN}[x]${NC} opencti-mcp           $opencti_status"
        else
            echo -e "    ${RED}[ ]${NC} opencti-mcp           $opencti_status"
        fi
    else
        echo -e "    ${DIM}[ ] opencti-mcp           (use --opencti to add)${NC}"
    fi

    if [[ "$WANT_REMNUX" == true ]]; then
        echo -e "    ${GREEN}[x]${NC} remnux-mcp            configured (remote)"
    else
        echo -e "    ${DIM}[ ] remnux-mcp            (use --remnux to add)${NC}"
    fi

    # Database status
    echo ""
    echo "  Databases:"
    if [[ -n "${RAG_BUILD_PID:-}" ]] && kill -0 "$RAG_BUILD_PID" 2>/dev/null; then
        echo "    RAG index:     building in background (PID $RAG_BUILD_PID)"
    elif [[ -d "$PROJECT_DIR/forensic-rag-mcp/data/chroma" ]] && \
         [[ -n "$(ls -A "$PROJECT_DIR/forensic-rag-mcp/data/chroma" 2>/dev/null)" ]]; then
        echo "    RAG index:     ready"
    else
        echo "    RAG index:     not built yet"
    fi

    if [[ -f "$PROJECT_DIR/windows-triage-mcp/data/known_good.db" ]]; then
        echo "    Triage DBs:    ready"
    elif [[ -n "${TRIAGE_BUILD_PID:-}" ]] && kill -0 "$TRIAGE_BUILD_PID" 2>/dev/null; then
        echo "    Triage DBs:    importing in background (PID $TRIAGE_BUILD_PID)"
    else
        echo "    Triage DBs:    not built yet"
    fi

    if [[ -f "$PROJECT_DIR/windows-triage-mcp/data/known_good_registry.db" ]]; then
        echo "    Registry DB:   ready"
    elif [[ -n "${REG_BUILD_PID:-}" ]] && kill -0 "$REG_BUILD_PID" 2>/dev/null; then
        echo "    Registry DB:   building in background (PID $REG_BUILD_PID)"
    elif [[ "$WITH_REGISTRY" == true ]]; then
        echo "    Registry DB:   not built yet"
    fi

    # Tools & hooks
    echo ""
    if [[ -f "$PROJECT_DIR/local/FORENSIC_TOOLS.md" ]]; then
        local tool_count
        tool_count=$(grep -c "^\- \[x\]" "$PROJECT_DIR/local/FORENSIC_TOOLS.md" 2>/dev/null || echo "?")
        echo "    Tools:         ${tool_count} found (local/FORENSIC_TOOLS.md)"
    else
        echo "    Tools:         skipped"
    fi

    if [[ -f "$PROJECT_DIR/.claude/settings.json" ]]; then
        echo "    Hooks:         installed"
    else
        echo "    Hooks:         skipped"
    fi

    echo "    Config:        .mcp.json generated"

    # Next steps
    echo ""
    echo -e "${BOLD}══════════════════════════════════════════════════════${NC}"
    echo -e "${BOLD}  Next Steps${NC}"
    echo -e "${BOLD}══════════════════════════════════════════════════════${NC}"
    echo ""

    if ! command -v claude &>/dev/null; then
        echo "  Claude Code not found in PATH."
        echo "  Install: https://docs.anthropic.com/en/docs/claude-code/overview"
        echo ""
    fi

    echo "  1. Run claude"
    echo "  2. Accept the 4 MCP servers when prompted"
    echo "  3. Run /setup-review"
    echo ""

    # Show build monitoring hint if builds are running
    local builds_running=false
    if [[ -n "${RAG_BUILD_PID:-}" ]] && kill -0 "$RAG_BUILD_PID" 2>/dev/null; then
        builds_running=true
    fi
    if [[ -n "${TRIAGE_BUILD_PID:-}" ]] && kill -0 "$TRIAGE_BUILD_PID" 2>/dev/null; then
        builds_running=true
    fi
    if [[ -n "${REG_BUILD_PID:-}" ]] && kill -0 "$REG_BUILD_PID" 2>/dev/null; then
        builds_running=true
    fi

    if [[ "$builds_running" == true ]]; then
        echo "  Note: databases are building in background."
        echo "  /setup-review will monitor them automatically."
        echo ""
    fi

    echo -e "${BOLD}══════════════════════════════════════════════════════${NC}"
}

# =============================================================================
# Banner
# =============================================================================

print_banner() {
    echo ""
    echo -e "${BOLD}══════════════════════════════════════════════════════${NC}"
    echo -e "${BOLD}  Claude-IR Setup${NC}"
    echo -e "${BOLD}══════════════════════════════════════════════════════${NC}"
    echo ""
    echo "  Forensic investigation workstation powered by Claude Code"
    echo ""
}

# =============================================================================
# Main
# =============================================================================

main() {
    setup_colors
    parse_args "$@"
    print_banner
    check_project_dir
    check_prerequisites

    determine_components
    phase_install_local_mcps
    phase_generate_config
    phase_database_setup

    phase_hooks
    phase_tools_inventory
    phase_summary
}

# Initialize build PID tracking
RAG_BUILD_PID=""
TRIAGE_BUILD_PID=""
REG_BUILD_PID=""

main "$@"