setup.sh

#!/bin/bash
# setup.sh - Red Team Implant Bootstrap Script
# Target: Raspberry Pi 5 | 16GB RAM | NVMe (PCIe Gen 3) | Debian Trixie

set -euo pipefail

# ==============================================================================
# CONFIGURABLE MODEL NAMES
# Update these to match the models set in src/red_audit/registry.py before
# running this script. The defaults below are validated Ollama tags.
# ==============================================================================
OLLAMA_PASSIVE_MODEL="phi:latest"
#OLLAMA_TACTICAL_MODEL="xploiter/pentester:v2"               # Fast per-module tactical summaries
#OLLAMA_STRATEGIC_MODEL="xploiter/pentester:latest"          # Full engagement Executive Summary
OLLAMA_AUDIT_MODEL="hf.co/BugTraceAI/BugTraceAI-CORE-Pro"   # Used by scripts/ai_code_auditor.py

NUCLEI_VERSION="3.3.8"

# ==============================================================================

echo "[*] Bootstrapping Red Audit Environment..."

ARCH=$(uname -m)
echo "[*] Detected Architecture: $ARCH"

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

BASE_PACKAGES="build-essential git python3-venv python3-dev libpcap-dev libkrb5-dev libssl-dev nmap pipx tmux unzip wget hostapd dnsmasq iw fonts-dejavu-core zstd network-manager libjpeg-dev zlib1g-dev"

# ------------------------------------------------------------------------------
# 1. OS-Level Dependencies
# ------------------------------------------------------------------------------
if command -v apt-get &> /dev/null; then
    echo "[*] Checking for OS dependencies..."
    sudo apt-get update -y

    if [[ "$ARCH" == "aarch64" ]]; then
        echo "[*] Adding Pi-specific hardware drivers..."
        sudo apt-get install -y $BASE_PACKAGES python3-lgpio gpiod liblgpio-dev swig
    else
        echo "[*] Standard x86 hardware detected. Skipping GPIO drivers."
        sudo apt-get install -y $BASE_PACKAGES
    fi
else
    echo "[!] apt-get not found. Ensure all required system packages are installed manually."
fi

# ------------------------------------------------------------------------------
# 1b. NetExec (nxc) — SMB recon layer, invoked via subprocess in smb.py
# Not available in Debian repos; must be installed from GitHub via pipx.
# ------------------------------------------------------------------------------
if ! command -v nxc &> /dev/null; then
    echo "[*] Installing NetExec (nxc) via pipx..."
    pipx install git+https://github.com/Pennyw0rth/NetExec
    pipx ensurepath
    echo "[+] NetExec installed. 'nxc' is now available in PATH."
else
    echo "[*] NetExec (nxc) already installed. Skipping."
fi

# ------------------------------------------------------------------------------
# 2. Hardware Tuning (Pi 5 only)
# ------------------------------------------------------------------------------
if [[ "$ARCH" == "aarch64" ]]; then
    echo "[*] Optimizing Pi 5 Hardware (NVMe PCIe Gen 3 & SPI Interface)..."
    CONFIG_FILE="/boot/firmware/config.txt"

    if ! grep -q "dtparam=pciex1_gen=3" "$CONFIG_FILE"; then
        echo "dtparam=pciex1_gen=3" | sudo tee -a "$CONFIG_FILE" > /dev/null
        echo "[+] PCIe Gen 3 enabled."
    fi

    if ! grep -q "dtparam=spi=on" "$CONFIG_FILE"; then
        echo "dtparam=spi=on" | sudo tee -a "$CONFIG_FILE" > /dev/null
        echo "[+] SPI Interface enabled for LCD screen."
    fi
else
    echo "[*] Skipping /boot/firmware tuning for x86 architecture."
fi

# ------------------------------------------------------------------------------
# 3. Ollama (Offline LLM Engine)
# ------------------------------------------------------------------------------
if ! command -v ollama &> /dev/null; then
    echo "[*] Installing Ollama..."
    curl -fsSL https://ollama.com/install.sh | sh
else
    echo "[*] Ollama already installed. Skipping."
fi

sudo systemctl enable --now ollama

echo "[*] Waiting for Ollama daemon to initialize..."
TRIES=0
while ! curl -s http://127.0.0.1:11434 > /dev/null; do
    sleep 1
    TRIES=$((TRIES + 1))
    if [[ $TRIES -ge 30 ]]; then
        echo "[!] Ollama did not respond within 30 seconds."
        echo "[!] Check daemon status: sudo systemctl status ollama"
        exit 1
    fi
done
echo "[+] Ollama daemon is ready."

# ------------------------------------------------------------------------------
# 4. Ollama Resource Limits (Pi 5 only)
# Prevents Ollama from saturating all 4 cores during inference, which causes
# the kernel scheduler to starve sshd and drop SSH connections.
# ------------------------------------------------------------------------------
if [[ "$ARCH" == "aarch64" ]]; then
    echo "[*] Applying Ollama CPU resource limits..."
    OLLAMA_DROP_IN="/etc/systemd/system/ollama.service.d/resource-limits.conf"
    sudo mkdir -p "$(dirname "$OLLAMA_DROP_IN")"

    sudo tee "$OLLAMA_DROP_IN" > /dev/null << 'DROPIN'
[Service]
# Caps Ollama to 2 of the Pi 5's 4 Cortex-A76 cores, guaranteeing 2 cores
# remain available for sshd, systemd, and kernel I/O at all times.
CPUQuota=200%
# Drops Ollama below normal interactive scheduling priority.
Nice=10
# Prevents model weight reads from monopolizing the NVMe queue.
IOSchedulingClass=best-effort
IOSchedulingPriority=5
# Aligns llama.cpp thread count with the CPUQuota ceiling.
Environment="OLLAMA_NUM_THREAD=2"
# One parallel request on a single-user implant is sufficient.
Environment="OLLAMA_NUM_PARALLEL=1"
# Prevents two loaded models from competing for the same quota-limited cores.
Environment="OLLAMA_MAX_LOADED_MODELS=1"
# Binds Ollama to all interfaces so LAN clients can reach the API.
# Default is 127.0.0.1 (localhost only).
Environment="OLLAMA_HOST=0.0.0.0:11434"
DROPIN

    sudo systemctl daemon-reload
    sudo systemctl restart ollama
    echo "[+] Ollama resource limits applied (CPUQuota=200%, Nice=10, 2 threads)."

    echo "[*] Hardening SSH keepalive settings..."
    set_sshd_option() {
        local key="$1" val="$2"
        # Replace existing (commented or active) entry, or append if absent
        if grep -qE "^\s*#?\s*${key}\s" /etc/ssh/sshd_config; then
            sudo sed -i "s|^\s*#\?\s*${key}\s.*|${key} ${val}|" /etc/ssh/sshd_config
        else
            echo "${key} ${val}" | sudo tee -a /etc/ssh/sshd_config > /dev/null
        fi
    }
    set_sshd_option "ClientAliveInterval" "15"
    set_sshd_option "ClientAliveCountMax" "4"
    set_sshd_option "TCPKeepAlive"        "yes"
    sudo systemctl reload sshd
    echo "[+] SSH keepalive hardened (ClientAliveInterval=15, CountMax=4)."
else
    echo "[*] Skipping Ollama resource limits (not aarch64)."
fi

# ------------------------------------------------------------------------------
# 5. Pull AI Models
# ------------------------------------------------------------------------------
echo "[*] Pulling passive model: $OLLAMA_PASSIVE_MODEL..."
ollama pull "$OLLAMA_PASSIVE_MODEL"

#echo "[*] Pulling tactical model: $OLLAMA_TACTICAL_MODEL..."
#ollama pull "$OLLAMA_TACTICAL_MODEL"

#echo "[*] Pulling strategic  model: $OLLAMA_STRATEGIC_MODEL..."
#ollama pull "$OLLAMA_STRATEGIC_MODEL"

echo "[*] Pulling code audit model: $OLLAMA_AUDIT_MODEL..."
ollama pull "$OLLAMA_AUDIT_MODEL"

# ------------------------------------------------------------------------------
# 6. Nuclei (Web Vulnerability Scanner)
# ------------------------------------------------------------------------------
if ! command -v nuclei &> /dev/null; then
    TMP_DIR=$(mktemp -d)

    if [[ "$ARCH" == "aarch64" ]]; then
        echo "[*] Downloading Nuclei v${NUCLEI_VERSION} (ARM64)..."
        NUCLEI_URL="https://github.com/projectdiscovery/nuclei/releases/download/v${NUCLEI_VERSION}/nuclei_${NUCLEI_VERSION}_linux_arm64.zip"
    else
        echo "[*] Downloading Nuclei v${NUCLEI_VERSION} (x86_64)..."
        NUCLEI_URL="https://github.com/projectdiscovery/nuclei/releases/download/v${NUCLEI_VERSION}/nuclei_${NUCLEI_VERSION}_linux_amd64.zip"
    fi

    wget "$NUCLEI_URL" -q -O "$TMP_DIR/nuclei.zip"
    unzip -q "$TMP_DIR/nuclei.zip" -d "$TMP_DIR/"
    sudo mv "$TMP_DIR/nuclei" /usr/local/bin/
    sudo chmod +x /usr/local/bin/nuclei
    rm -rf "$TMP_DIR"
    echo "[+] Nuclei v${NUCLEI_VERSION} installed."
else
    echo "[*] Nuclei already installed. Skipping."
fi

echo "[*] Downloading Nuclei templates..."
nuclei -update-templates -silent || echo "[!] Template download failed. Run manually: nuclei -update-templates"

# ------------------------------------------------------------------------------
# 7. CISA Vulnrichment (Offline CVE Intelligence)
# Provides KEV flags, SSVC scores, CVSS, and CWE for local CVE lookups.
# Stored in /opt/vulnrichment — no internet required during engagements.
# ------------------------------------------------------------------------------
VULNRICH_DIR="/opt/vulnrichment"
if [ ! -d "$VULNRICH_DIR/.git" ]; then
    echo "[*] Cloning CISA Vulnrichment database to $VULNRICH_DIR..."
    sudo git clone --depth=1 https://github.com/cisagov/vulnrichment.git "$VULNRICH_DIR"
    echo "[+] Vulnrichment database cloned."
else
    echo "[*] Vulnrichment already present at $VULNRICH_DIR. Skipping."
fi

# ------------------------------------------------------------------------------
# 8. uv (Python package manager)
# ------------------------------------------------------------------------------
if ! command -v uv &>/dev/null; then
    echo "[*] Installing uv..."
    curl -LsSf https://astral.sh/uv/install.sh | sh
    export PATH="$HOME/.local/bin:$PATH"
else
    echo "[*] uv already installed. Skipping."
fi

# ------------------------------------------------------------------------------
# 9. Install Python Dependencies
# ------------------------------------------------------------------------------
cd "$PROJECT_ROOT"
echo "[*] Installing Red Audit package via uv..."

# WHY: cryptography is a transitive dep of impacket/bloodhound. uv prefers
# binary wheels by default, but if no ARM64 wheel exists for the running Python
# version it falls back to building from source (requires Rust compiler).
if ! uv sync --no-build 2>/dev/null; then
    echo "[!] Some packages require building from source. Installing Rust compiler..."
    sudo apt-get install -y cargo
    uv sync
fi

# ------------------------------------------------------------------------------
# 10. Global Execution Wrappers
# ------------------------------------------------------------------------------
echo "[*] Creating global execution wrappers..."
cat << EOF | sudo tee /usr/local/bin/red-audit > /dev/null
#!/bin/bash
SESSION="red-audit"
BINARY="$PROJECT_ROOT/.venv/bin/red-audit"

if [ -n "\$TMUX" ]; then
    exec "\$BINARY" "\$@"
fi

tmux kill-session -t "\$SESSION" 2>/dev/null || true

QUOTED_ARGS=\$(printf '%q ' "\$@")

tmux new-session -d -s "\$SESSION" \
    -e "PYTHONUNBUFFERED=1" \
    "\$BINARY \$QUOTED_ARGS"

tmux set-option -t "\$SESSION" remain-on-exit on
exec tmux attach-session -t "\$SESSION"
EOF
sudo chmod +x /usr/local/bin/red-audit

if [[ "$ARCH" == "aarch64" ]]; then
    cat << EOF | sudo tee /usr/local/bin/red-menu > /dev/null
#!/bin/bash
exec $PROJECT_ROOT/.venv/bin/python3 $PROJECT_ROOT/src/red_audit/menu.py "\$@"
EOF
    sudo chmod +x /usr/local/bin/red-menu
fi
# ------------------------------------------------------------------------------
echo ""
echo "============================================================"
echo "[+] Red Audit environment ready."
echo "============================================================"
echo ""
if [[ "$ARCH" == "aarch64" ]]; then
    echo "[!] A system reboot is REQUIRED to apply PCIe Gen 3 and SPI changes."
    echo "    Run: sudo reboot"
    echo ""
fi
echo "Usage (CLI):  sudo red-audit <module> -t <target>"
echo "Usage (LCD):  sudo red-menu"
echo "Help:         sudo red-audit --help"
echo ""
echo "Models pulled:  $OLLAMA_TACTICAL_MODEL  |  $OLLAMA_STRATEGIC_MODEL  |  $OLLAMA_AUDIT_MODEL"
echo "Verify they match registry.ai_config in src/red_audit/registry.py"
echo ""
echo "Vulnrichment:   $VULNRICH_DIR  (run update_env.sh to refresh)"
echo ""
echo "Dev install:  To include test/lint tools:"
echo "              uv sync --group dev"
echo "============================================================"