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

HTTP_PORT=8080
HTTPS_PORT=8443
TCP_ECHO_PORT=9000
UDP_ECHO_PORT=9001
IPERF_TCP_PORT=5201
IPERF_UDP_PORT=5202
SOCKPERF_TCP_PORT=11111
SOCKPERF_UDP_PORT=11112
DNS_PORT=5353
WORKDIR="/tmp/mitm-traffic-target"
CERT_DAYS=2

PIDS=()

usage() {
  cat <<USAGE
Usage:
  $0 [options]

Options:
  --http-port <port>       HTTP port (default: 8080)
  --https-port <port>      HTTPS port (default: 8443)
  --tcp-echo-port <port>   TCP echo port (default: 9000)
  --udp-echo-port <port>   UDP echo port (default: 9001)
  --iperf-tcp-port <port>  iPerf TCP server port (default: 5201)
  --iperf-udp-port <port>  iPerf UDP server port (default: 5202)
  --sockperf-tcp-port <port> Sockperf TCP server port (default: 11111)
  --sockperf-udp-port <port> Sockperf UDP server port (default: 11112)
  --dns-port <port>        Lightweight DNS UDP port (default: 5353)
  --workdir <dir>          Working directory (default: /tmp/mitm-traffic-target)
  --help                   Show this help

Notes:
  - DNS service is a lightweight UDP responder using socat (not full DNS).
  - For DNS on port 53, run as root/cap_net_bind_service or choose high port and set generator --dns-port.
USAGE
}

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

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

need_cmds() {
  local missing=0
  for cmd in "$@"; do
    if ! has_cmd "$cmd"; then
      echo "Missing command: $cmd" >&2
      missing=1
    fi
  done
  if [[ "$missing" -ne 0 ]]; then
    echo "Install missing tools and retry." >&2
    exit 1
  fi
}

parse_args() {
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --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 ;;
      --sockperf-tcp-port)
        SOCKPERF_TCP_PORT="${2:-11111}"; shift 2 ;;
      --sockperf-udp-port)
        SOCKPERF_UDP_PORT="${2:-11112}"; shift 2 ;;
      --dns-port)
        DNS_PORT="${2:-5353}"; shift 2 ;;
      --workdir)
        WORKDIR="${2:-/tmp/mitm-traffic-target}"; shift 2 ;;
      -h|--help)
        usage; exit 0 ;;
      *)
        echo "Unknown argument: $1" >&2
        usage
        exit 1
        ;;
    esac
  done
}

start_bg() {
  local name="$1"
  shift
  "$@" >"${WORKDIR}/${name}.log" 2>&1 &
  local pid=$!
  PIDS+=("$pid")
  log "Started ${name} (pid=${pid})"
}

cleanup() {
  log 'Stopping services...'
  for pid in "${PIDS[@]:-}"; do
    kill "$pid" >/dev/null 2>&1 || true
  done
  wait >/dev/null 2>&1 || true
  log "Stopped. Logs in ${WORKDIR}"
}

create_cert() {
  if [[ -f "${WORKDIR}/cert.pem" && -f "${WORKDIR}/key.pem" ]]; then
    return
  fi

  openssl req -x509 -newkey rsa:2048 -nodes \
    -keyout "${WORKDIR}/key.pem" \
    -out "${WORKDIR}/cert.pem" \
    -days "$CERT_DAYS" \
    -subj '/CN=mitm-traffic-target.local' >/dev/null 2>&1
}

main() {
  parse_args "$@"

  need_cmds python3 openssl socat
  if ! has_cmd iperf3; then
    log 'iperf3 not found, iperf services will be skipped'
  fi
  if ! has_cmd sockperf; then
    log 'sockperf not found, sockperf services will be skipped'
  fi
  if ! has_cmd netserver; then
    log 'netserver not found, Flent netperf-based tests may not work'
  fi

  mkdir -p "$WORKDIR"
  trap cleanup EXIT INT TERM

  create_cert

  log "Workdir: ${WORKDIR}"
  log 'Starting target services...'

  start_bg http python3 -m http.server "$HTTP_PORT" --directory "$WORKDIR"
  start_bg https openssl s_server -quiet -accept "$HTTPS_PORT" -cert "${WORKDIR}/cert.pem" -key "${WORKDIR}/key.pem"

  start_bg tcp_echo socat "TCP-LISTEN:${TCP_ECHO_PORT},reuseaddr,fork" SYSTEM:'cat'
  start_bg udp_echo socat "UDP-LISTEN:${UDP_ECHO_PORT},reuseaddr,fork" SYSTEM:'cat'

  if has_cmd iperf3; then
    start_bg iperf_tcp iperf3 -s -p "$IPERF_TCP_PORT"
    start_bg iperf_udp iperf3 -s -p "$IPERF_UDP_PORT"
  fi

  if has_cmd sockperf; then
    start_bg sockperf_tcp sockperf server --tcp --port "$SOCKPERF_TCP_PORT"
    start_bg sockperf_udp sockperf server --udp --port "$SOCKPERF_UDP_PORT"
  fi

  if has_cmd netserver; then
    start_bg netserver netserver -D
  fi

  # Lightweight DNS-like UDP responder to generate DNS-shaped traffic.
  # It echoes fixed bytes and is intended only for packet-generation tests.
  start_bg dns_dummy socat "UDP-LISTEN:${DNS_PORT},reuseaddr,fork" SYSTEM:'printf "\\x81\\x80"'

  cat <<INFO

Target services are running.

Ports:
  HTTP      : ${HTTP_PORT}
  HTTPS     : ${HTTPS_PORT}
  TCP echo  : ${TCP_ECHO_PORT}
  UDP echo  : ${UDP_ECHO_PORT}
  iPerf TCP : ${IPERF_TCP_PORT}
  iPerf UDP : ${IPERF_UDP_PORT}
  Sockperf TCP : ${SOCKPERF_TCP_PORT}
  Sockperf UDP : ${SOCKPERF_UDP_PORT}
  DNS dummy : ${DNS_PORT}

Matching benchmark options:
  --iperf-tcp-port ${IPERF_TCP_PORT} --iperf-udp-port ${IPERF_UDP_PORT}
  --sockperf-tcp-port ${SOCKPERF_TCP_PORT} --sockperf-udp-port ${SOCKPERF_UDP_PORT}

For Flent RRUL/netperf-based tests, install netperf/netserver on this host.

Keep this process running. Press Ctrl+C to stop all services.

INFO

  while true; do
    sleep 1
  done
}

main "$@"