#!/usr/bin/env bash
# setup_pentest_env.sh - Set up the penetration testing environment
# Installs Python dependencies and verifies tool availability.
#
# Usage:
#   ./setup_pentest_env.sh [--venv]
#
# Options:
#   --venv    Create and use a virtual environment (recommended)

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REQUIREMENTS="$SCRIPT_DIR/requirements.txt"
MIN_PYTHON_MAJOR=3
MIN_PYTHON_MINOR=9

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

info()  { printf "${BLUE}[*]${NC} %s\n" "$1"; }
ok()    { printf "${GREEN}[+]${NC} %s\n" "$1"; }
warn()  { printf "${YELLOW}[!]${NC} %s\n" "$1"; }
fail()  { printf "${RED}[-]${NC} %s\n" "$1"; }

header() {
    echo ""
    echo "======================================"
    echo "  Penetration Tester - Environment Setup"
    echo "======================================"
    echo ""
}

check_python() {
    info "Checking Python version..."

    if ! command -v python3 &>/dev/null; then
        fail "Python 3 is not installed."
        echo "  Install Python ${MIN_PYTHON_MAJOR}.${MIN_PYTHON_MINOR}+ from https://python.org"
        return 1
    fi

    local version
    version=$(python3 -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}')")
    local major minor
    major=$(python3 -c "import sys; print(sys.version_info.major)")
    minor=$(python3 -c "import sys; print(sys.version_info.minor)")

    if (( major < MIN_PYTHON_MAJOR )) || { (( major == MIN_PYTHON_MAJOR )) && (( minor < MIN_PYTHON_MINOR )); }; then
        fail "Python $version found, but ${MIN_PYTHON_MAJOR}.${MIN_PYTHON_MINOR}+ is required."
        return 1
    fi

    ok "Python $version"
}

setup_venv() {
    local venv_dir="$SCRIPT_DIR/.venv"

    if [[ -d "$venv_dir" ]]; then
        info "Activating existing virtual environment..."
    else
        info "Creating virtual environment at $venv_dir..."
        python3 -m venv "$venv_dir"
    fi

    # shellcheck disable=SC1091
    source "$venv_dir/bin/activate"
    ok "Virtual environment active: $venv_dir"
}

install_deps() {
    info "Installing Python dependencies from requirements.txt..."

    if [[ ! -f "$REQUIREMENTS" ]]; then
        fail "requirements.txt not found at $REQUIREMENTS"
        return 1
    fi

    if pip install -r "$REQUIREMENTS" --quiet 2>/dev/null; then
        ok "All dependencies installed."
    else
        fail "Failed to install some dependencies. Try: pip install -r $REQUIREMENTS"
        return 1
    fi
}

verify_tool() {
    local name="$1"
    local cmd="$2"

    if eval "$cmd" &>/dev/null; then
        local ver
        ver=$(eval "$cmd" 2>&1 | head -1)
        ok "$name: $ver"
    else
        warn "$name: not available (optional — some scans may be limited)"
    fi
}

verify_tools() {
    info "Verifying tool availability..."
    echo ""

    verify_tool "bandit"    "bandit --version"
    verify_tool "pip-audit" "pip-audit --version"
    verify_tool "npm"       "npm --version"
    verify_tool "nmap"      "nmap --version"
    verify_tool "curl"      "curl --version"
    verify_tool "openssl"   "openssl version"
}

check_system_tools() {
    info "Checking system tools..."

    local missing=0
    for tool in curl openssl; do
        if command -v "$tool" &>/dev/null; then
            ok "$tool is available"
        else
            warn "$tool is not installed (some checks may be limited)"
            missing=$((missing + 1))
        fi
    done

    return 0
}

print_summary() {
    echo ""
    echo "======================================"
    echo "  Setup Complete"
    echo "======================================"
    echo ""
    echo "Available scanners:"
    echo "  - security_scanner.py    HTTP headers, SSL/TLS, endpoint probing"
    echo "  - dependency_auditor.py  npm audit, pip-audit wrapper"
    echo "  - code_security_scanner.py  bandit + regex pattern analysis"
    echo ""
    echo "Quick start:"
    echo "  python3 $SCRIPT_DIR/security_scanner.py https://example.com"
    echo "  python3 $SCRIPT_DIR/dependency_auditor.py /path/to/project"
    echo "  python3 $SCRIPT_DIR/code_security_scanner.py /path/to/code"
    echo ""
}

main() {
    header

    local use_venv=false
    for arg in "$@"; do
        case "$arg" in
            --venv) use_venv=true ;;
            --help|-h)
                echo "Usage: $0 [--venv]"
                echo ""
                echo "Options:"
                echo "  --venv    Create and use a virtual environment (recommended)"
                echo "  --help    Show this help message"
                exit 0
                ;;
            *)
                fail "Unknown argument: $arg"
                exit 1
                ;;
        esac
    done

    local errors=0

    check_python || errors=$((errors + 1))

    if (( errors > 0 )); then
        fail "Cannot continue without Python ${MIN_PYTHON_MAJOR}.${MIN_PYTHON_MINOR}+."
        exit 1
    fi

    if $use_venv; then
        setup_venv
    fi

    install_deps || errors=$((errors + 1))
    check_system_tools
    verify_tools

    if (( errors > 0 )); then
        echo ""
        warn "Setup completed with warnings. Some features may be limited."
        exit 1
    fi

    print_summary
}

main "$@"
