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

TARGET=""
IFACE=""
REQUESTS=50
INTERVAL_MS=50
STOP_ON_KEYPRESS=0
PROTOCOLS_CSV="all"
DNS_PORT=53
HTTP_PORT=8080
HTTPS_PORT=8443
TCP_ECHO_PORT=9000
UDP_ECHO_PORT=9001
IPERF_TCP_PORT=5201
IPERF_UDP_PORT=5202

STOP=0
MAIN_PID=$$
KEYPRESS_PID=""

usage() {
  cat <<USAGE
Usage:
  $0 --target <ip-or-host> --iface <iface> [options]

Required:
  --target <ip-or-host>      Target host running services or receiving probes
  --iface <iface>            Source interface (used by arping/nping where applicable)

Optional:
  --protocols <list>         Comma-separated protocol set or 'all'
                             Available: arp,icmp,http,https,dns,tcp_echo,udp_echo,
                                        tcp_syn,tcp_flags,udp_raw,traceroute,
                                        iperf_tcp,iperf_udp,ssh_probe,ftp_probe,
                                        smtp_probe,ntp_probe,mdns_query
                             Default: all
  --requests <n>             Requests per selected protocol (default: 50)
  --interval-ms <n>          Delay between requests in ms (default: 50)
  --stop-on-keypress         Stop loop when any key is pressed

  --dns-port <port>          DNS port (default: 53)
  --http-port <port>         HTTP port (default: 8080)
  --https-port <port>        HTTPS port (default: 8443)
  --tcp-echo-port <port>     TCP echo/custom port (default: 9000)
  --udp-echo-port <port>     UDP echo/custom port (default: 9001)
  --iperf-tcp-port <port>    iPerf TCP server port (default: 5201)
  --iperf-udp-port <port>    iPerf UDP server port (default: 5202)

Examples:
  $0 --target 10.0.0.2 --iface enp3s0 --protocols icmp,http,https,dns --requests 100
  $0 --target 10.0.0.2 --iface enp3s0 --protocols all --requests 500 --stop-on-keypress
USAGE
}

log() {
  printf '[%s] %s\n' "$(date +'%H:%M:%S')" "$*"
}

has_cmd() {
  command -v "$1" >/dev/null 2>&1
}

sleep_ms() {
  local ms="$1"
  local sec=$((ms / 1000))
  local rem=$((ms % 1000))
  sleep "$(printf '%s.%03d' "$sec" "$rem")"
}

maybe_wait() {
  if [[ "$INTERVAL_MS" -gt 0 ]]; then
    sleep_ms "$INTERVAL_MS"
  fi
}

start_keypress_listener() {
  if [[ "$STOP_ON_KEYPRESS" -ne 1 ]]; then
    return
  fi

  if [[ ! -t 0 ]]; then
    log "--stop-on-keypress requested, but no interactive TTY detected; ignoring"
    STOP_ON_KEYPRESS=0
    return
  fi

  trap 'STOP=1' SIGUSR1
  log "Press any key to stop traffic generation"
  (
    while [[ "$STOP" -eq 0 ]]; do
      IFS= read -r -s -n 1 _key </dev/tty && kill -USR1 "$MAIN_PID"
    done
  ) &
  KEYPRESS_PID=$!
}

cleanup() {
  if [[ -n "$KEYPRESS_PID" ]]; then
    kill "$KEYPRESS_PID" >/dev/null 2>&1 || true
  fi
}

parse_args() {
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --target)
        TARGET="${2:-}"
        shift 2
        ;;
      --iface)
        IFACE="${2:-}"
        shift 2
        ;;
      --protocols)
        PROTOCOLS_CSV="${2:-all}"
        shift 2
        ;;
      --requests)
        REQUESTS="${2:-50}"
        shift 2
        ;;
      --interval-ms)
        INTERVAL_MS="${2:-50}"
        shift 2
        ;;
      --stop-on-keypress)
        STOP_ON_KEYPRESS=1
        shift
        ;;
      --dns-port)
        DNS_PORT="${2:-53}"
        shift 2
        ;;
      --http-port)
        HTTP_PORT="${2:-8080}"
        shift 2
        ;;
      --https-port)
        HTTPS_PORT="${2:-8443}"
        shift 2
        ;;
      --tcp-echo-port)
        TCP_ECHO_PORT="${2:-9000}"
        shift 2
        ;;
      --udp-echo-port)
        UDP_ECHO_PORT="${2:-9001}"
        shift 2
        ;;
      --iperf-tcp-port)
        IPERF_TCP_PORT="${2:-5201}"
        shift 2
        ;;
      --iperf-udp-port)
        IPERF_UDP_PORT="${2:-5202}"
        shift 2
        ;;
      -h|--help)
        usage
        exit 0
        ;;
      *)
        echo "Unknown argument: $1" >&2
        usage
        exit 1
        ;;
    esac
  done
}

assert_inputs() {
  if [[ -z "$TARGET" || -z "$IFACE" ]]; then
    usage
    exit 1
  fi

  if ! [[ "$REQUESTS" =~ ^[0-9]+$ ]] || [[ "$REQUESTS" -lt 1 ]]; then
    echo "--requests must be an integer >= 1" >&2
    exit 1
  fi

  if ! [[ "$INTERVAL_MS" =~ ^[0-9]+$ ]]; then
    echo "--interval-ms must be an integer >= 0" >&2
    exit 1
  fi
}

run_arp() {
  has_cmd arping || return 0
  for _ in $(seq 1 "$REQUESTS"); do
    [[ "$STOP" -eq 1 ]] && return
    arping -I "$IFACE" -c 1 "$TARGET" >/dev/null 2>&1 || true
    maybe_wait
  done
}

run_icmp() {
  has_cmd ping || return 0
  for _ in $(seq 1 "$REQUESTS"); do
    [[ "$STOP" -eq 1 ]] && return
    ping -c 1 -W 1 "$TARGET" >/dev/null 2>&1 || true
    maybe_wait
  done
}

run_http() {
  has_cmd curl || return 0
  for _ in $(seq 1 "$REQUESTS"); do
    [[ "$STOP" -eq 1 ]] && return
    curl -s "http://${TARGET}:${HTTP_PORT}/" >/dev/null 2>&1 || true
    maybe_wait
  done
}

run_https() {
  has_cmd curl || return 0
  for _ in $(seq 1 "$REQUESTS"); do
    [[ "$STOP" -eq 1 ]] && return
    curl -sk "https://${TARGET}:${HTTPS_PORT}/" >/dev/null 2>&1 || true
    maybe_wait
  done
}

run_dns() {
  if has_cmd dig; then
    for _ in $(seq 1 "$REQUESTS"); do
      [[ "$STOP" -eq 1 ]] && return
      dig @"$TARGET" -p "$DNS_PORT" example.com A +tries=1 +time=1 +short >/dev/null 2>&1 || true
      maybe_wait
    done
  fi
}

run_tcp_echo() {
  has_cmd nc || return 0
  for i in $(seq 1 "$REQUESTS"); do
    [[ "$STOP" -eq 1 ]] && return
    printf 'tcp-echo-%s\n' "$i" | nc -w1 "$TARGET" "$TCP_ECHO_PORT" >/dev/null 2>&1 || true
    maybe_wait
  done
}

run_udp_echo() {
  has_cmd nc || return 0
  for i in $(seq 1 "$REQUESTS"); do
    [[ "$STOP" -eq 1 ]] && return
    printf 'udp-echo-%s\n' "$i" | nc -u -w1 "$TARGET" "$UDP_ECHO_PORT" >/dev/null 2>&1 || true
    maybe_wait
  done
}

run_tcp_syn() {
  if has_cmd nping; then
    nping --interface "$IFACE" --tcp -p "$TCP_ECHO_PORT" --flags syn -c "$REQUESTS" "$TARGET" >/dev/null 2>&1 || true
  fi
}

run_tcp_flags() {
  if has_cmd nping; then
    nping --interface "$IFACE" --tcp -p "$TCP_ECHO_PORT" --flags syn,ack,fin,rst,psh,urg -c "$REQUESTS" "$TARGET" >/dev/null 2>&1 || true
  fi
}

run_udp_raw() {
  if has_cmd nping; then
    nping --interface "$IFACE" --udp -p "$UDP_ECHO_PORT" --data-length 256 -c "$REQUESTS" "$TARGET" >/dev/null 2>&1 || true
  fi
}

run_traceroute() {
  has_cmd traceroute || return 0
  for _ in $(seq 1 "$REQUESTS"); do
    [[ "$STOP" -eq 1 ]] && return
    traceroute -n -m 5 "$TARGET" >/dev/null 2>&1 || true
    maybe_wait
  done
}

run_iperf_tcp() {
  has_cmd iperf3 || return 0
  iperf3 -c "$TARGET" -p "$IPERF_TCP_PORT" -t 10 >/dev/null 2>&1 || true
}

run_iperf_udp() {
  has_cmd iperf3 || return 0
  iperf3 -c "$TARGET" -p "$IPERF_UDP_PORT" -u -b 100M -t 10 >/dev/null 2>&1 || true
}

run_port_probe() {
  local port="$1"
  has_cmd nc || return 0
  for _ in $(seq 1 "$REQUESTS"); do
    [[ "$STOP" -eq 1 ]] && return
    nc -z -w1 "$TARGET" "$port" >/dev/null 2>&1 || true
    maybe_wait
  done
}

run_ssh_probe() { run_port_probe 22; }
run_ftp_probe() { run_port_probe 21; }
run_smtp_probe() { run_port_probe 25; }
run_ntp_probe() {
  has_cmd nc || return 0
  for _ in $(seq 1 "$REQUESTS"); do
    [[ "$STOP" -eq 1 ]] && return
    printf 'ntp?\n' | nc -u -w1 "$TARGET" 123 >/dev/null 2>&1 || true
    maybe_wait
  done
}
run_mdns_query() {
  if has_cmd dig; then
    for _ in $(seq 1 "$REQUESTS"); do
      [[ "$STOP" -eq 1 ]] && return
      dig @224.0.0.251 -p 5353 _services._dns-sd._udp.local PTR +tries=1 +time=1 +short >/dev/null 2>&1 || true
      maybe_wait
    done
  fi
}

dispatch_protocol() {
  local proto="$1"
  log "Running protocol: ${proto}"
  case "$proto" in
    arp) run_arp ;;
    icmp) run_icmp ;;
    http) run_http ;;
    https) run_https ;;
    dns) run_dns ;;
    tcp_echo) run_tcp_echo ;;
    udp_echo) run_udp_echo ;;
    tcp_syn) run_tcp_syn ;;
    tcp_flags) run_tcp_flags ;;
    udp_raw) run_udp_raw ;;
    traceroute) run_traceroute ;;
    iperf_tcp) run_iperf_tcp ;;
    iperf_udp) run_iperf_udp ;;
    ssh_probe) run_ssh_probe ;;
    ftp_probe) run_ftp_probe ;;
    smtp_probe) run_smtp_probe ;;
    ntp_probe) run_ntp_probe ;;
    mdns_query) run_mdns_query ;;
    *)
      log "Unknown protocol '${proto}', skipping"
      ;;
  esac
}

main() {
  parse_args "$@"
  assert_inputs

  local protocols=(arp icmp http https dns tcp_echo udp_echo tcp_syn tcp_flags udp_raw traceroute iperf_tcp iperf_udp ssh_probe ftp_probe smtp_probe ntp_probe mdns_query)

  if [[ "$PROTOCOLS_CSV" != "all" ]]; then
    IFS=',' read -r -a protocols <<<"$PROTOCOLS_CSV"
  fi

  start_keypress_listener
  trap cleanup EXIT

  log "Target=${TARGET}, iface=${IFACE}, requests=${REQUESTS}, protocols=${PROTOCOLS_CSV}, stop_on_keypress=${STOP_ON_KEYPRESS}"
  for proto in "${protocols[@]}"; do
    [[ "$STOP" -eq 1 ]] && break
    dispatch_protocol "$proto"
  done

  log "Traffic generation complete"
}

main "$@"