#!/usr/bin/env bash
set -euo pipefail

# Script Name: pscan
# Description: Unified port scanner wrapper (nmap/masscan/rustscan)
# Usage: pscan <target>             # Quick scan with best available tool
#        pscan <target> -f          # Full port scan (all 65535)
#        pscan <target> -u          # UDP scan
#        pscan <target> -v          # Version detection
#        pscan <target> -s          # Stealth scan

VERSION="1.0.0"

# Colors
readonly RED='\033[0;31m'
readonly GREEN='\033[0;32m'
readonly YELLOW='\033[1;33m'
readonly BLUE='\033[0;34m'
readonly CYAN='\033[0;36m'
readonly MAGENTA='\033[0;35m'
readonly BOLD='\033[1m'
readonly NC='\033[0m'

show_help() {
  echo -e "${BOLD}pscan${NC} - Unified Port Scanner Wrapper v${VERSION}"
  echo
  echo -e "${BOLD}USAGE:${NC}"
  echo "  pscan <target> [OPTIONS]"
  echo
  echo -e "${BOLD}OPTIONS:${NC}"
  echo -e "  ${CYAN}-f, --full${NC}       Scan all 65535 ports"
  echo -e "  ${CYAN}-t, --top${NC}        Scan top 1000 ports (default)"
  echo -e "  ${CYAN}-q, --quick${NC}      Quick scan (top 100 ports)"
  echo -e "  ${CYAN}-u, --udp${NC}        UDP scan"
  echo -e "  ${CYAN}-v, --version${NC}    Version detection"
  echo -e "  ${CYAN}-s, --stealth${NC}    Stealth SYN scan"
  echo -e "  ${CYAN}-a, --aggressive${NC} Aggressive scan (OS, version, scripts, traceroute)"
  echo -e "  ${CYAN}-o, --output FILE${NC} Save output to file"
  echo -e "  ${CYAN}-h, --help${NC}       Show this help"
  echo
  echo -e "${BOLD}TOOL PREFERENCE:${NC}"
  echo "  1. rustscan  (fastest, if available)"
  echo "  2. masscan   (fast, if available)"
  echo "  3. nmap      (fallback, always available)"
  echo
  echo -e "${BOLD}EXAMPLES:${NC}"
  echo "  pscan 192.168.1.1              # Quick scan"
  echo "  pscan 192.168.1.0/24           # Scan subnet"
  echo "  pscan 10.10.10.5 -f            # Full port scan"
  echo "  pscan target.com -v            # Version detection"
  echo "  pscan 10.0.0.1 -s              # Stealth scan"
  echo "  pscan 192.168.1.1 -o scan.txt  # Save output"
  echo
  echo -e "${BOLD}INSTALLED TOOLS:${NC}"
  command -v rustscan &>/dev/null && echo -e "  ${GREEN}✓${NC} rustscan" || echo -e "  ${RED}✗${NC} rustscan (install: cargo install rustscan)"
  command -v masscan &>/dev/null && echo -e "  ${GREEN}✓${NC} masscan" || echo -e "  ${RED}✗${NC} masscan (install: sudo apt install masscan)"
  command -v nmap &>/dev/null && echo -e "  ${GREEN}✓${NC} nmap" || echo -e "  ${RED}✗${NC} nmap (install: sudo apt install nmap)"
}

# Detect best scanner
get_scanner() {
  if command -v rustscan &>/dev/null; then
    echo "rustscan"
  elif command -v masscan &>/dev/null; then
    echo "masscan"
  elif command -v nmap &>/dev/null; then
    echo "nmap"
  else
    echo -e "${RED}Error:${NC} No port scanner found" >&2
    echo "Install one: sudo apt install nmap masscan" >&2
    exit 1
  fi
}

# Rustscan wrapper
scan_rustscan() {
  local target="$1"
  local ports="${2:-1-65535}"
  local extra_args="${3:-}"

  echo -e "${CYAN}[*]${NC} Using rustscan for ${BOLD}$target${NC}"
  echo -e "${CYAN}[*]${NC} Ports: $ports"
  echo

  rustscan -a "$target" -r "$ports" --ulimit 5000 $extra_args
}

# Masscan wrapper
scan_masscan() {
  local target="$1"
  local ports="${2:-0-65535}"
  local rate="${3:-1000}"

  echo -e "${CYAN}[*]${NC} Using masscan for ${BOLD}$target${NC}"
  echo -e "${CYAN}[*]${NC} Ports: $ports | Rate: $rate pps"
  echo

  sudo masscan "$target" -p"$ports" --rate="$rate" --open
}

# Nmap wrapper
scan_nmap() {
  local target="$1"
  local scan_type="${2:--sS}"
  local ports="${3:--p-}"
  local extra_args="${4:-}"

  echo -e "${CYAN}[*]${NC} Using nmap for ${BOLD}$target${NC}"
  echo -e "${CYAN}[*]${NC} Scan type: $scan_type | Ports: $ports"
  echo

  sudo nmap "$scan_type" "$ports" $extra_args "$target"
}

# Parse arguments
if [[ $# -eq 0 ]] || [[ "$1" =~ ^(-h|--help|help)$ ]]; then
  show_help
  exit 0
fi

target="$1"
shift

# Default settings
mode="top"
scan_type="tcp"
output_file=""
scanner=$(get_scanner)
aggressive=false
version_detect=false
stealth=false

# Parse options
while [[ $# -gt 0 ]]; do
  case $1 in
    -f|--full)
      mode="full"
      shift
      ;;
    -t|--top)
      mode="top"
      shift
      ;;
    -q|--quick)
      mode="quick"
      shift
      ;;
    -u|--udp)
      scan_type="udp"
      shift
      ;;
    -v|--version)
      version_detect=true
      shift
      ;;
    -s|--stealth)
      stealth=true
      shift
      ;;
    -a|--aggressive)
      aggressive=true
      shift
      ;;
    -o|--output)
      output_file="$2"
      shift 2
      ;;
    *)
      echo -e "${RED}Error:${NC} Unknown option: $1"
      exit 1
      ;;
  esac
done

# Build scan command based on scanner and options
case "$scanner" in
  rustscan)
    case "$mode" in
      full)
        ports="1-65535"
        ;;
      top)
        ports="1-10000"
        ;;
      quick)
        ports="1-1000"
        ;;
    esac

    extra_args=""
    [[ -n "$output_file" ]] && extra_args="$extra_args -o $output_file"

    scan_rustscan "$target" "$ports" "$extra_args"
    ;;

  masscan)
    case "$mode" in
      full)
        ports="0-65535"
        rate="10000"
        ;;
      top)
        ports="1-10000"
        rate="5000"
        ;;
      quick)
        ports="1-1000"
        rate="1000"
        ;;
    esac

    [[ -n "$output_file" ]] && extra_args="$extra_args -oL $output_file"

    scan_masscan "$target" "$ports" "$rate"
    ;;

  nmap)
    # Build nmap arguments
    if [[ "$stealth" == "true" ]]; then
      nmap_scan="-sS"
    elif [[ "$scan_type" == "udp" ]]; then
      nmap_scan="-sU"
    else
      nmap_scan="-sS"
    fi

    case "$mode" in
      full)
        nmap_ports="-p-"
        ;;
      top)
        nmap_ports="--top-ports 1000"
        ;;
      quick)
        nmap_ports="--top-ports 100"
        ;;
    esac

    nmap_extra=""
    [[ "$version_detect" == "true" ]] && nmap_extra="$nmap_extra -sV"
    [[ "$aggressive" == "true" ]] && nmap_extra="$nmap_extra -A"
    [[ -n "$output_file" ]] && nmap_extra="$nmap_extra -oN $output_file"

    scan_nmap "$target" "$nmap_scan" "$nmap_ports" "$nmap_extra"
    ;;
esac

echo
echo -e "${GREEN}[✓]${NC} Scan complete"
[[ -n "$output_file" ]] && echo -e "${GREEN}[✓]${NC} Output saved to: $output_file"
