#!/bin/bash
set -euo pipefail

# ===================================
# PENTEST AUTOMATION SCRIPT v2.0
# Tools: Wapiti, Nuclei, Trivy
# ===================================

# ---------- 1. CONFIG ----------
declare -A CONFIG=(
  [dev]="http://localhost:3388|gosot"
  [prod]="https://gosot.example.com|gosot"
)

# ---------- 2. CLI ARGS ----------
ENV="${1:-dev}"
OVERRIDE_URL="${2:-}"
OVERRIDE_IMAGE="${3:-}"

# Validate environment
if [[ -z "${CONFIG[$ENV]:-}" ]]; then
  echo "Error: Unknown environment '$ENV'. Use 'dev' or 'prod'."
  exit 1
fi

# Split config
IFS='|' read -r TARGET_URL IMAGE_NAME <<< "${CONFIG[$ENV]}"

# Apply overrides
[[ -n "$OVERRIDE_URL" ]] && TARGET_URL="$OVERRIDE_URL"
[[ -n "$OVERRIDE_IMAGE" ]] && IMAGE_NAME="$OVERRIDE_IMAGE"

# Validate URL
if [[ ! "$TARGET_URL" =~ ^https?:// ]]; then
  echo "Error: Invalid URL: $TARGET_URL"
  exit 1
fi

# ---------- 3. SETUP ----------
TIMESTAMP=$(date +%F_%H%M%S)
REPORT_DIR="pentest_reports/${ENV}_${TIMESTAMP}"
mkdir -p "$REPORT_DIR"

echo "=== PENTEST STARTED ==="
echo "Environment: $ENV"
echo "Target URL: $TARGET_URL"
[[ -n "$IMAGE_NAME" ]] && echo "Docker Image: $IMAGE_NAME" || echo "Docker Image: (skipped)"
echo "Reports → $REPORT_DIR"
echo

# ---------- 4. TOOL CHECKS ----------
check_tool() {
  if ! command -v "$1" &> /dev/null; then
    echo "Error: $1 is not installed or not in PATH."
    exit 1
  fi
}

check_tool nuclei
check_tool trivy

# Activate venv once
if [[ -f "./scripts/.wapiti-venv/bin/activate" ]]; then
  source "./scripts/.wapiti-venv/bin/activate"
else
  echo "Warning: Wapiti venv not found. Assuming global install."
fi

# ---------- 5. WAPITI SCAN ----------
echo "[1/3] Running Wapiti (full crawl + all modules)..."
wapiti \
  -u "$TARGET_URL" \
  --scope domain \
  -m all \
  --level 2 \
  --timeout 30 \
  --max-files-per-dir 1000 \
  --max-parameters 500 \
  --color \
  -f html -o "$REPORT_DIR/wapiti_$TIMESTAMP/" || echo "Wapiti completed with warnings."

# ---------- 6. NUCLEI SCAN ----------
echo "[2/3] Running Nuclei (critical/high/medium, tagged templates)..."
nuclei \
  -u "$TARGET_URL" \
  -t cves/ -t vulnerabilities/ -t exposures/ -t misconfiguration/ -t ssl/ \
  -tags web,api,exposed,backup,token \
  -severity critical,high,medium \
  -c 30 -rl 30 -retries 2 \
  -timeout 10 \
  -es info,low \
  -jsonl "$REPORT_DIR/nuclei.jsonl" \
  -me "$REPORT_DIR/nuclei.md" \
  -silent \
  -stats -stats-interval 10 || echo "Nuclei completed with warnings."

# ---------- 7. TRIVY IMAGE SCAN ----------
if [[ -n "$IMAGE_NAME" ]]; then
  echo "[3/3] Scanning Docker image: $IMAGE_NAME"

  # Pull image if not exists
  if ! docker image inspect "$IMAGE_NAME" &>/dev/null; then
    echo "Pulling image: $IMAGE_NAME"
    docker pull "$IMAGE_NAME" || echo "Failed to pull image. Continuing with local cache."
  fi

  # Single Trivy call with multiple outputs
  trivy image \
    --severity CRITICAL,HIGH,MEDIUM \
    --scanners vuln,misconfig \
    --format table \
    --output "$REPORT_DIR/trivy.txt" \
    "$IMAGE_NAME" || echo "Trivy table scan failed."

  trivy image \
    --format json \
    --output "$REPORT_DIR/trivy.json" \
    "$IMAGE_NAME" || echo "Trivy JSON scan failed."
else
  echo "[3/3] No Docker image specified. Skipping Trivy."
fi

# ---------- 8. FINALIZE ----------
deactivate 2>/dev/null || true

echo
echo "=== PENTEST COMPLETE ==="
echo "HTML Report: $REPORT_DIR/wapiti_$TIMESTAMP/"
echo "Nuclei JSON: $REPORT_DIR/nuclei.json"
echo "Nuclei MD:   $REPORT_DIR/nuclei.md"
[[ -f "$REPORT_DIR/trivy.txt" ]] && echo "Trivy TXT:   $REPORT_DIR/trivy.txt"
[[ -f "$REPORT_DIR/trivy.json" ]] && echo "Trivy JSON:  $REPORT_DIR/trivy.json"
echo
