#!/usr/bin/env bash

set -e

ROOT="$(cd "$(dirname "$(readlink -f "$0")")/.." && pwd)"

# Formatting
BOLD=$(tput bold 2>/dev/null || true)
NORMAL=$(tput sgr0 2>/dev/null || true)
MAROON='\033[38;5;1m'
RED='\033[31m'
GREEN='\033[32m'
YELLOW='\033[33m'
BLUE='\033[34m'
CYAN='\033[36m'

# Banner
echo -e "${BOLD}${MAROON}"
cat <<BANNER
██╗      █████╗ ███╗   ██╗██╗███╗   ███╗ █████╗ ██╗     ███████╗
██║     ██╔══██╗████╗  ██║██║████╗ ████║██╔══██╗██║     ██╔════╝
██║     ███████║██╔██╗ ██║██║██╔████╔██║███████║██║     ███████╗
██║     ██╔══██║██║╚██╗██║██║██║╚██╔╝██║██╔══██║██║     ╚════██║
███████╗██║  ██║██║ ╚████║██║██║ ╚═╝ ██║██║  ██║███████╗███████║
╚══════╝╚═╝  ╚═╝╚═╝  ╚═══╝╚═╝╚═╝     ╚═╝╚═╝  ╚═╝╚══════╝╚══════╝
               VULNERABILITY SCANNER
BANNER
echo -e "${NORMAL}"

echo -e "${BOLD}${MAROON}[ INIT ]${NORMAL} Vulnerability Scanner Initializing..."
sleep 0.3

# Create log directory
mkdir -p ~/.lanimals/logs/vulnscan

# Setup variables
SCAN_TARGET=""
SCAN_TYPE="basic"
OUTPUT_FILE=~/.lanimals/logs/vulnscan/scan_$(date +%Y%m%d_%H%M%S).log
SCAN_PORTS="21,22,23,25,53,80,443,445,3306,3389,8080,8443"

# Check for nmap
if ! command -v nmap &> /dev/null; then
    echo -e "${BOLD}${RED}[ ERROR ]${NORMAL} nmap not found. Install with 'sudo apt install nmap'"
    exit 1
fi

# Get primary interface
PRIMARY_INTERFACE=$(ip route | grep default | awk '{print $5}')
LOCAL_IP=$(ip -o -4 addr show $PRIMARY_INTERFACE | awk '{print $4}' | cut -d/ -f1)
SUBNET=$(ip -o -4 addr show $PRIMARY_INTERFACE | awk '{print $4}' | cut -d/ -f1 | sed 's/\.[0-9]*$/.0\/24/')

# Function to display help
show_help() {
    echo "Usage: lanimals vulnscan [OPTIONS]"
    echo ""
    echo "Options:"
    echo "  --target=HOST/NETWORK    Approved private target (default: local subnet)"
    echo "  --ports=PORTS            Specify ports (default: common ports)"
    echo "  --intense                Perform intense scan (more thorough but slower)"
    echo "  --stealth                Perform stealth scan (less detectable)"
    echo "  --vuln                   Scan for known vulnerabilities"
    echo "  --ssl                    Check for SSL/TLS vulnerabilities"
    echo "  --output=FILE            Custom output file"
    echo "  --help                   Show this help message"
    echo ""
    echo "Examples:"
    echo "  lanimals vulnscan --target=192.168.1.1"
    echo "  lanimals vulnscan --target=192.168.1.0/24 --intense"
    echo "  LANIMALS_ALLOWED_CIDRS=192.168.1.0/24 lanimals vulnscan --target=192.168.1.10 --vuln"
}

# Function to perform basic scan
basic_scan() {
    local target=$1
    local ports=$2
    local output=$3
    
    echo -e "${BOLD}${MAROON}[ SCAN ]${NORMAL} Performing basic vulnerability scan on $target..."
    echo -e "${BOLD}${MAROON}[ INFO ]${NORMAL} Scanning ports: $ports"
    echo -e "${BOLD}${MAROON}[ INFO ]${NORMAL} Results will be saved to $output"
    
    # Run the scan, redirect to log file and display progress
    # Removed -F flag as it conflicts with -p
    sudo nmap -sV -T4 -O --version-light -p "$ports" "$target" | tee "$output"
    
    echo -e "${BOLD}${GREEN}[ OK ]${NORMAL} Basic scan completed."
}

# Function to perform intense scan
intense_scan() {
    local target=$1
    local ports=$2
    local output=$3
    
    echo -e "${BOLD}${MAROON}[ SCAN ]${NORMAL} Performing intense vulnerability scan on $target..."
    echo -e "${BOLD}${MAROON}[ INFO ]${NORMAL} Scanning ports: $ports"
    echo -e "${BOLD}${MAROON}[ INFO ]${NORMAL} Results will be saved to $output"
    echo -e "${BOLD}${YELLOW}[ WARN ]${NORMAL} This scan is more thorough but slower and more detectable."
    
    # Run the scan, redirect to log file and display progress
    sudo nmap -sS -sV -T4 -A -O -p "$ports" "$target" | tee "$output"
    
    echo -e "${BOLD}${GREEN}[ OK ]${NORMAL} Intense scan completed."
}

# Function to perform stealth scan
stealth_scan() {
    local target=$1
    local ports=$2
    local output=$3
    
    echo -e "${BOLD}${MAROON}[ SCAN ]${NORMAL} Performing stealth vulnerability scan on $target..."
    echo -e "${BOLD}${MAROON}[ INFO ]${NORMAL} Scanning ports: $ports"
    echo -e "${BOLD}${MAROON}[ INFO ]${NORMAL} Results will be saved to $output"
    echo -e "${BOLD}${YELLOW}[ WARN ]${NORMAL} This scan is less likely to be detected but may miss some information."
    
    # Run the scan, redirect to log file and display progress
    sudo nmap -sS -T2 -f -D RND:5 -p "$ports" "$target" | tee "$output"
    
    echo -e "${BOLD}${GREEN}[ OK ]${NORMAL} Stealth scan completed."
}

# Function to perform vulnerability scan
vuln_scan() {
    local target=$1
    local output=$2
    
    echo -e "${BOLD}${MAROON}[ SCAN ]${NORMAL} Scanning for known vulnerabilities on $target..."
    echo -e "${BOLD}${MAROON}[ INFO ]${NORMAL} Results will be saved to $output"
    echo -e "${BOLD}${YELLOW}[ WARN ]${NORMAL} This scan can take a long time to complete."
    
    # Run the scan, redirect to log file and display progress
    sudo nmap --script vuln "$target" | tee "$output"
    
    echo -e "${BOLD}${GREEN}[ OK ]${NORMAL} Vulnerability scan completed."
}

# Function to scan for SSL/TLS vulnerabilities
ssl_scan() {
    local target=$1
    local output=$2
    
    echo -e "${BOLD}${MAROON}[ SCAN ]${NORMAL} Scanning for SSL/TLS vulnerabilities on $target..."
    echo -e "${BOLD}${MAROON}[ INFO ]${NORMAL} Results will be saved to $output"
    
    # Run the scan, redirect to log file and display progress
    sudo nmap --script ssl-enum-ciphers,ssl-heartbleed,ssl-poodle,ssl-ccs-injection "$target" | tee "$output"
    
    echo -e "${BOLD}${GREEN}[ OK ]${NORMAL} SSL/TLS vulnerability scan completed."
}

# Parse command line arguments
if [ $# -eq 0 ]; then
    # Default scan if no arguments provided
    SCAN_TARGET="$SUBNET"
    echo -e "${BOLD}${MAROON}[ INFO ]${NORMAL} No target specified, using default: $SCAN_TARGET"
else
    for arg in "$@"; do
        case $arg in
            --target=*)
                SCAN_TARGET="${arg#*=}"
                shift
                ;;
            --ports=*)
                SCAN_PORTS="${arg#*=}"
                shift
                ;;
            --intense)
                SCAN_TYPE="intense"
                shift
                ;;
            --stealth)
                SCAN_TYPE="stealth"
                shift
                ;;
            --vuln)
                SCAN_TYPE="vuln"
                shift
                ;;
            --ssl)
                SCAN_TYPE="ssl"
                shift
                ;;
            --output=*)
                OUTPUT_FILE="${arg#*=}"
                shift
                ;;
            --help)
                show_help
                exit 0
                ;;
            *)
                echo -e "${BOLD}${RED}[ ERROR ]${NORMAL} Unknown option: $arg"
                show_help
                exit 1
                ;;
        esac
    done
fi

# If no target specified after parsing, use default
if [ -z "$SCAN_TARGET" ]; then
    SCAN_TARGET="$SUBNET"
    echo -e "${BOLD}${MAROON}[ INFO ]${NORMAL} No target specified, using default: $SCAN_TARGET"
fi

if [[ ! "$SCAN_PORTS" =~ ^[0-9,-]+$ ]]; then
    echo -e "${BOLD}${RED}[ ERROR ]${NORMAL} Invalid port list. Use digits, commas, and ranges only." >&2
    exit 2
fi

if ! SCAN_TARGET=$(PYTHONPATH="$ROOT" python3 -m core.validate_scope_cli "$SCAN_TARGET"); then
    echo -e "${BOLD}${RED}[ ERROR ]${NORMAL} Refusing target outside the approved LAN scope." >&2
    exit 2
fi

# Run the selected scan type
case $SCAN_TYPE in
    basic)
        basic_scan "$SCAN_TARGET" "$SCAN_PORTS" "$OUTPUT_FILE"
        ;;
    intense)
        intense_scan "$SCAN_TARGET" "$SCAN_PORTS" "$OUTPUT_FILE"
        ;;
    stealth)
        stealth_scan "$SCAN_TARGET" "$SCAN_PORTS" "$OUTPUT_FILE"
        ;;
    vuln)
        vuln_scan "$SCAN_TARGET" "$OUTPUT_FILE"
        ;;
    ssl)
        ssl_scan "$SCAN_TARGET" "$OUTPUT_FILE"
        ;;
    *)
        echo -e "${BOLD}${RED}[ ERROR ]${NORMAL} Unknown scan type: $SCAN_TYPE"
        show_help
        exit 1
        ;;
esac

# Display summary
echo
echo -e "${BOLD}${MAROON}[ SUMMARY ]${NORMAL} Scan completed"
echo -e "${BOLD}Target:${NORMAL} $SCAN_TARGET"
echo -e "${BOLD}Scan Type:${NORMAL} $SCAN_TYPE"
echo -e "${BOLD}Log File:${NORMAL} $OUTPUT_FILE"
echo
echo -e "${BOLD}${MAROON}[ INFO ]${NORMAL} View scan report anytime with: cat $OUTPUT_FILE"
