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

# Combined Wazuh simulator script (single entrypoint)
# Replays production-style sample logs from samples/*.log to Wazuh syslog UDP.
#
# Usage:
#   scripts/send-wazuh-sim-logs.sh [selector] [count] [delay_seconds] [--forever] [--dry-run] [--no-mock] [--no-guarantee-hits] [--random-types] [--include-nonalerts] [--docker-send]
#
# Selectors:
#   all
#   a|b|c|appendix-a|appendix-b|appendix-c
#   a1|a2|a3|a4|b1|b2|b3|c1|c2|c3
#   A1-01, A2-10, B3-06, C1-01, ...

SELECTOR="${1:-all}"
COUNT="${2:-1}"
DELAY="${3:-1}"
shift $(( $# >= 3 ? 3 : $# )) || true

FOREVER=0
DRY_RUN="${DRY_RUN:-0}"
MOCK_VALUES="${MOCK_VALUES:-1}"
GUARANTEE_HITS="${GUARANTEE_HITS:-1}"
RANDOM_TYPES="${RANDOM_TYPES:-0}"
INCLUDE_NONALERTS="${INCLUDE_NONALERTS:-0}"
DOCKER_SEND="${DOCKER_SEND:-0}"
WAZUH_MANAGER_CONTAINER="${WAZUH_MANAGER_CONTAINER:-wazuh-single-wazuh.manager-1}"

for arg in "$@"; do
  case "$arg" in
    --forever)
      FOREVER=1
      ;;
    --dry-run)
      DRY_RUN=1
      ;;
    --no-mock)
      MOCK_VALUES=0
      ;;
    --mock)
      MOCK_VALUES=1
      ;;
    --no-guarantee-hits)
      GUARANTEE_HITS=0
      ;;
    --guarantee-hits)
      GUARANTEE_HITS=1
      ;;
    --random-types)
      RANDOM_TYPES=1
      ;;
    --include-nonalerts)
      INCLUDE_NONALERTS=1
      ;;
    --docker-send)
      DOCKER_SEND=1
      ;;
    -h|--help)
      echo "usage: scripts/send-wazuh-sim-logs.sh [selector] [count] [delay_seconds] [--forever] [--dry-run] [--no-mock] [--no-guarantee-hits] [--random-types] [--include-nonalerts] [--docker-send]"
      exit 0
      ;;
    *)
      echo "error: unknown option '$arg'"
      exit 1
      ;;
  esac
done

if ! [[ "$COUNT" =~ ^[0-9]+$ ]] || [ "$COUNT" -lt 1 ]; then
  echo "error: count must be a positive integer"
  exit 1
fi

if ! [[ "$DELAY" =~ ^[0-9]+([.][0-9]+)?$ ]]; then
  echo "error: delay_seconds must be numeric"
  exit 1
fi

WAZUH_SYSLOG_HOST="${WAZUH_SYSLOG_HOST:-127.0.0.1}"
WAZUH_SYSLOG_PORT="${WAZUH_SYSLOG_PORT:-514}"
NC_WAIT_SECONDS="${NC_WAIT_SECONDS:-0}"
STRICT_SEND="${STRICT_SEND:-1}"
BASE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
SAMPLES_DIR="${BASE_DIR}/samples"

normalize() {
  printf '%s' "$1" | tr '[:upper:]' '[:lower:]'
}

is_valid_selector() {
  local sel
  sel="$(normalize "$1")"
  case "$sel" in
    all|a|b|c|appendix-a|appendix-b|appendix-c|a1|a2|a3|a4|b1|b2|b3|c1|c2|c3)
      return 0
      ;;
    [abc][0-9]-[0-9][0-9])
      return 0
      ;;
    *)
      return 1
      ;;
  esac
}

selector_matches_tag() {
  local selector tag sel tagl
  selector="$1"
  tag="$2"
  sel="$(normalize "$selector")"
  tagl="$(normalize "$tag")"

  if [ -z "$tagl" ]; then
    case "$sel" in
      all|a|b|c|appendix-a|appendix-b|appendix-c)
        return 0
        ;;
      *)
        return 1
        ;;
    esac
  fi

  case "$sel" in
    all)
      return 0
      ;;
    a|appendix-a)
      [[ "$tagl" == a* ]]
      return
      ;;
    b|appendix-b)
      [[ "$tagl" == b* ]]
      return
      ;;
    c|appendix-c)
      [[ "$tagl" == c* ]]
      return
      ;;
    a1|a2|a3|a4|b1|b2|b3|c1|c2|c3)
      [[ "$tagl" == "$sel"-* ]]
      return
      ;;
    [abc][0-9]-[0-9][0-9])
      [[ "$tagl" == "$sel" ]]
      return
      ;;
    *)
      return 1
      ;;
  esac
}

sample_files_for_selector() {
  local sel
  sel="$(normalize "$1")"
  case "$sel" in
    all)
      echo "${SAMPLES_DIR}/appendix-a-production-samples.log"
      echo "${SAMPLES_DIR}/appendix-b-production-samples.log"
      echo "${SAMPLES_DIR}/appendix-c-production-samples.log"
      ;;
    a|appendix-a|a1|a2|a3|a4|a[0-9]-[0-9][0-9])
      echo "${SAMPLES_DIR}/appendix-a-production-samples.log"
      ;;
    b|appendix-b|b1|b2|b3|b[0-9]-[0-9][0-9])
      echo "${SAMPLES_DIR}/appendix-b-production-samples.log"
      ;;
    c|appendix-c|c1|c2|c3|c[0-9]-[0-9][0-9])
      echo "${SAMPLES_DIR}/appendix-c-production-samples.log"
      ;;
    *)
      echo "error: unsupported selector '$1'" >&2
      return 1
      ;;
  esac
}

emit_syslog() {
  local line="$1"
  if [ "${DRY_RUN}" = "1" ]; then
    echo "DRY_RUN -> ${line}"
    return 0
  fi

  if [ "${DOCKER_SEND}" = "1" ]; then
    if ! printf '%s\n' "${line}" | docker exec -i "${WAZUH_MANAGER_CONTAINER}" bash -lc 'cat > /dev/udp/127.0.0.1/514'; then
      echo "send_failed target=${WAZUH_MANAGER_CONTAINER}:127.0.0.1:514/udp transport=docker-exec" >&2
      if [ "${STRICT_SEND}" = "1" ]; then
        return 1
      fi
    fi
    return 0
  fi

  if command -v nc >/dev/null 2>&1; then
    if ! printf '%s\n' "${line}" | nc -w "${NC_WAIT_SECONDS}" -u "${WAZUH_SYSLOG_HOST}" "${WAZUH_SYSLOG_PORT}"; then
      echo "send_failed target=${WAZUH_SYSLOG_HOST}:${WAZUH_SYSLOG_PORT}/udp transport=nc" >&2
      if [ "${STRICT_SEND}" = "1" ]; then
        return 1
      fi
    fi
  else
    if ! printf '%s\n' "${line}" >"/dev/udp/${WAZUH_SYSLOG_HOST}/${WAZUH_SYSLOG_PORT}"; then
      echo "send_failed target=${WAZUH_SYSLOG_HOST}:${WAZUH_SYSLOG_PORT}/udp transport=devudp" >&2
      if [ "${STRICT_SEND}" = "1" ]; then
        return 1
      fi
    fi
  fi
}

rand_between() {
  local min max
  min="$1"
  max="$2"
  echo $(( min + RANDOM % (max - min + 1) ))
}

random_public_ip() {
  local range last
  range="$(rand_between 0 2)"
  last="$(rand_between 2 254)"
  case "$range" in
    0) echo "198.51.100.${last}" ;;
    1) echo "203.0.113.${last}" ;;
    *) echo "192.0.2.${last}" ;;
  esac
}

random_private_ip() {
  echo "10.$(rand_between 10 30).$(rand_between 1 254).$(rand_between 1 254)"
}

random_user() {
  local users=(
    "admin01" "analyst01" "helpdesk01" "ops.admin" "jane.doe"
    "svc_backup$" "svc_dbbackup$" "finance.user" "it-admin" "guest"
  )
  echo "${users[$((RANDOM % ${#users[@]}))]}"
}

random_fgt_model() {
  local models=("FGT40F-Branch01" "FGT60F-Branch01" "FGT80F-Branch01" "FGT501E-DC01")
  echo "${models[$((RANDOM % ${#models[@]}))]}"
}

random_devid() {
  local n
  n="$(rand_between 10000000 99999999)"
  echo "FGT80FTK${n}"
}

random_domain() {
  echo "ioc-$(rand_between 1000 9999).malicious.example"
}

replace_kv() {
  local input key new
  input="$1"
  key="$2"
  new="$3"
  printf '%s' "$input" | sed -E "s#${key}=\"[^\"]*\"#${key}=\"${new}\"#g; s#${key}=([^\" ][^ ]*)#${key}=${new}#g"
}

mock_windows_json_line() {
  local line src_ip
  line="$1"
  src_ip="$(random_public_ip)"

  if command -v jq >/dev/null 2>&1; then
    printf '%s' "$line" | jq -c \
      --arg srcip "$src_ip" \
      '
      if (.win.eventdata | type) == "object" then
        .win.eventdata |= (
          if has("subjectUserName") then .subjectUserName = "SYSTEM" else . end |
          if has("workstationName") then .workstationName = ("WS-" + (($srcip|split("."))[3])) else . end |
          if has("ipAddress") then .ipAddress = $srcip else . end
        )
      else
        .
      end
      ' 2>/dev/null || printf '%s' "$line"
  else
    printf '%s' "$line"
  fi
}

mock_non_json_line() {
  local line now_date now_time now_iso epoch src_pub dst_priv prev_pub devname devid query
  line="$1"
  now_date="$(date '+%Y-%m-%d')"
  now_time="$(date '+%H:%M:%S')"
  now_iso="$(date -u '+%Y-%m-%dT%H:%M:%S.000Z')"
  epoch="$(date '+%s')"
  src_pub="$(random_public_ip)"
  dst_priv="$(random_private_ip)"
  prev_pub="$(random_public_ip)"
  devname="$(random_fgt_model)"
  devid="$(random_devid)"
  query="$(random_domain)"

  line="$(replace_kv "$line" "date" "$now_date")"
  line="$(replace_kv "$line" "time" "$now_time")"
  line="$(replace_kv "$line" "eventtime" "$epoch")"
  line="$(replace_kv "$line" "devname" "$devname")"
  line="$(replace_kv "$line" "devid" "$devid")"

  line="$(replace_kv "$line" "srcip" "$src_pub")"
  line="$(replace_kv "$line" "dstip" "$dst_priv")"
  line="$(replace_kv "$line" "src_ip" "$src_pub")"
  line="$(replace_kv "$line" "prev_ip" "$prev_pub")"

  line="$(replace_kv "$line" "query" "$query")"

  line="$(replace_kv "$line" "srcport" "$(rand_between 1025 65535)")"
  line="$(replace_kv "$line" "distance_km" "$(rand_between 500 16000)")"
  line="$(replace_kv "$line" "travel_minutes" "$(rand_between 5 180)")"

  if [[ "$line" =~ ^[0-9]{4}-[0-9]{2}-[0-9]{2}T ]]; then
    line="$(printf '%s' "$line" | sed -E "s#^[0-9]{4}-[0-9]{2}-[0-9]{2}T[^ ]+#${now_iso}#")"
  fi
  line="$(printf '%s' "$line" | sed -E "s#from [0-9]{1,3}(\.[0-9]{1,3}){3}#from ${src_pub}#g")"
  line="$(printf '%s' "$line" | sed -E "s# port [0-9]{2,5}# port $(rand_between 1025 65535)#g")"

  printf '%s' "$line"
}

mock_line() {
  local line="$1"
  if [ "$MOCK_VALUES" != "1" ]; then
    printf '%s' "$line"
    return 0
  fi

  if [[ "$line" =~ ^\{ ]]; then
    mock_windows_json_line "$line"
  else
    mock_non_json_line "$line"
  fi
}

send_file_once() {
  local file selector line sent current_tag extracted
  file="$1"
  selector="$2"
  sent=0
  current_tag=""

  while IFS= read -r line || [ -n "$line" ]; do
    if [[ "$line" =~ ^[[:space:]]*#[[:space:]]*([A-Za-z][0-9]-[0-9]{2})([[:space:]]|$) ]]; then
      extracted="${BASH_REMATCH[1]}"
      current_tag="$(normalize "$extracted")"
      continue
    fi

    if [[ -z "${line// }" ]] || [[ "$line" =~ ^[[:space:]]*# ]]; then
      continue
    fi

    if selector_matches_tag "$selector" "$current_tag"; then
      line="$(mock_line "$line")"
      emit_syslog "$line"
      sent=$((sent + 1))
      sleep "$DELAY"
    fi
  done < "$file"

  echo "sent=${sent} file=$(basename "$file") selector=$(normalize "$selector")"
}

send_guaranteed_hits_once() {
  local selector sent idx tag line
  selector="$1"
  sent=0

  local tags=(
    "a1-01"
    "a1-02"
    "a2-02"
    "a2-03"
    "a2-05"
    "a2-10"
    "c1-01"
    "c1-01"
  )

  local lines=(
    "soc_event=dns_ioc event_type=ioc_dns_traffic src_ip=10.26.45.214 query=ioc-2294.malicious.example action=blocked severity=medium"
    "soc_event=dns_ioc event_type=ioc_domain_match src_ip=10.26.45.214 query=bad-c2.example feed=internal_main confidence=high action=alert"
    "date=2026-03-09 time=10:02:04 devname=\"FGT80F-Branch01\" devid=\"FGT80FTK20000001\" eventtime=1773079324 vd=\"root\" logid=\"0100044547\" type=\"event\" subtype=\"system\" level=\"warning\" user=\"admin\" action=\"password-change\" ui=\"https(10.20.55.1)\""
    "date=2026-03-09 time=10:02:17 devname=\"FGT80F-Branch01\" devid=\"FGT80FTK20000001\" eventtime=1773079337 vd=\"root\" logid=\"0100044548\" type=\"event\" subtype=\"system\" level=\"warning\" user=\"admin\" action=\"create-admin\" target_user=\"soc-backup-admin\""
    "date=2026-03-09 time=10:04:03 devname=\"FGT80F-Branch01\" devid=\"FGT80FTK20000001\" eventtime=1773079443 vd=\"root\" logid=\"0100044552\" type=\"event\" subtype=\"system\" level=\"notice\" user=\"admin\" action=\"download-config\" dstip=10.20.50.33"
    "date=2026-03-09 time=10:07:59 devname=\"FGT80F-Branch01\" devid=\"FGT80FTK20000001\" eventtime=1773079679 vd=\"root\" logid=\"0000000014\" type=\"traffic\" subtype=\"forward\" level=\"warning\" srcip=10.20.55.50 dstip=203.0.113.60 dstport=443 threat_label=\"known-c2\" action=\"accept\""
    "date=2026-03-09 time=10:31:00 devname=\"FGT80F-Branch01\" devid=\"FGT80FTK20000001\" eventtime=1773081060 vd=\"root\" logid=\"0101037135\" type=\"event\" subtype=\"vpn\" tunneltype=\"ssl\" action=\"ssl-login-success\" user=\"analyst01\" srcip=203.0.113.71 previous_country=TH current_country=US"
    "soc_event=correlation event_type=c1_impossible_travel user=\"analyst01\" src_ip=203.0.113.71 prev_ip=203.0.113.11 prev_country=TH current_country=US distance_km=13890 travel_minutes=18"
  )

  for idx in "${!lines[@]}"; do
    tag="${tags[$idx]}"
    if selector_matches_tag "$selector" "$tag"; then
      line="$(mock_line "${lines[$idx]}")"
      emit_syslog "$line"
      sent=$((sent + 1))
      sleep "$DELAY"
    fi
  done

  echo "guaranteed_sent=${sent} selector=$(normalize "$selector")"
}

EVENT_POOL_READY=0
EVENT_POOL_SELECTOR=""
declare -a EVENT_POOL_TAGS
declare -a EVENT_POOL_LINES

build_event_pool() {
  local selector file line current_tag extracted idx tag
  selector="$1"
  EVENT_POOL_READY=0
  EVENT_POOL_SELECTOR="$selector"
  EVENT_POOL_TAGS=()
  EVENT_POOL_LINES=()

  if [ "$INCLUDE_NONALERTS" = "1" ]; then
    for file in "${FILES[@]}"; do
      current_tag=""
      while IFS= read -r line || [ -n "$line" ]; do
        if [[ "$line" =~ ^[[:space:]]*#[[:space:]]*([A-Za-z][0-9]-[0-9]{2})([[:space:]]|$) ]]; then
          extracted="${BASH_REMATCH[1]}"
          current_tag="$(normalize "$extracted")"
          continue
        fi
        if [[ -z "${line// }" ]] || [[ "$line" =~ ^[[:space:]]*# ]]; then
          continue
        fi
        if selector_matches_tag "$selector" "$current_tag"; then
          EVENT_POOL_TAGS+=("$current_tag")
          EVENT_POOL_LINES+=("$line")
        fi
      done < "$file"
    done
  fi

  if [ "$GUARANTEE_HITS" = "1" ]; then
    local guaranteed_tags=(
      "a1-01" "a1-02" "a2-02" "a2-03" "a2-05" "a2-10" "c1-01" "c1-01"
    )
    local guaranteed_lines=(
      "soc_event=dns_ioc event_type=ioc_dns_traffic src_ip=10.26.45.214 query=ioc-2294.malicious.example action=blocked severity=medium"
      "soc_event=dns_ioc event_type=ioc_domain_match src_ip=10.26.45.214 query=bad-c2.example feed=internal_main confidence=high action=alert"
      "date=2026-03-09 time=10:02:04 devname=\"FGT80F-Branch01\" devid=\"FGT80FTK20000001\" eventtime=1773079324 vd=\"root\" logid=\"0100044547\" type=\"event\" subtype=\"system\" level=\"warning\" user=\"admin\" action=\"password-change\" ui=\"https(10.20.55.1)\""
      "date=2026-03-09 time=10:02:17 devname=\"FGT80F-Branch01\" devid=\"FGT80FTK20000001\" eventtime=1773079337 vd=\"root\" logid=\"0100044548\" type=\"event\" subtype=\"system\" level=\"warning\" user=\"admin\" action=\"create-admin\" target_user=\"soc-backup-admin\""
      "date=2026-03-09 time=10:04:03 devname=\"FGT80F-Branch01\" devid=\"FGT80FTK20000001\" eventtime=1773079443 vd=\"root\" logid=\"0100044552\" type=\"event\" subtype=\"system\" level=\"notice\" user=\"admin\" action=\"download-config\" dstip=10.20.50.33"
      "date=2026-03-09 time=10:07:59 devname=\"FGT80F-Branch01\" devid=\"FGT80FTK20000001\" eventtime=1773079679 vd=\"root\" logid=\"0000000014\" type=\"traffic\" subtype=\"forward\" level=\"warning\" srcip=10.20.55.50 dstip=203.0.113.60 dstport=443 threat_label=\"known-c2\" action=\"accept\""
      "date=2026-03-09 time=10:31:00 devname=\"FGT80F-Branch01\" devid=\"FGT80FTK20000001\" eventtime=1773081060 vd=\"root\" logid=\"0101037135\" type=\"event\" subtype=\"vpn\" tunneltype=\"ssl\" action=\"ssl-login-success\" user=\"analyst01\" srcip=203.0.113.71 previous_country=TH current_country=US"
      "soc_event=correlation event_type=c1_impossible_travel user=\"analyst01\" src_ip=203.0.113.71 prev_ip=203.0.113.11 prev_country=TH current_country=US distance_km=13890 travel_minutes=18"
    )
    for idx in "${!guaranteed_lines[@]}"; do
      tag="${guaranteed_tags[$idx]}"
      if selector_matches_tag "$selector" "$tag"; then
        EVENT_POOL_TAGS+=("$tag")
        EVENT_POOL_LINES+=("${guaranteed_lines[$idx]}")
      fi
    done
  fi

  EVENT_POOL_READY=1
}

send_random_event_once() {
  local selector size idx line tag
  selector="$1"
  if [ "$EVENT_POOL_READY" -ne 1 ] || [ "$EVENT_POOL_SELECTOR" != "$selector" ]; then
    build_event_pool "$selector"
  fi
  size="${#EVENT_POOL_LINES[@]}"
  if [ "$size" -eq 0 ]; then
    echo "random_sent=0 selector=$(normalize "$selector")"
    return 0
  fi
  idx=$((RANDOM % size))
  tag="${EVENT_POOL_TAGS[$idx]}"
  line="${EVENT_POOL_LINES[$idx]}"
  line="$(mock_line "$line")"
  emit_syslog "$line"
  echo "random_sent=1 tag=${tag} selector=$(normalize "$selector")"
}

if ! is_valid_selector "$SELECTOR"; then
  echo "error: selector must be one of all|a|b|c|appendix-a|appendix-b|appendix-c|a1..a4|b1..b3|c1..c3|A1-01..C3-04"
  exit 1
fi

FILES=()
while IFS= read -r f; do
  [ -n "$f" ] && FILES+=("$f")
done < <(sample_files_for_selector "$SELECTOR")

for f in "${FILES[@]}"; do
  if [ ! -f "$f" ]; then
    echo "error: missing sample file '$f'"
    exit 1
  fi
done

echo "selector=${SELECTOR} count=${COUNT} delay=${DELAY}s forever=${FOREVER} dry_run=${DRY_RUN} mock_values=${MOCK_VALUES} guarantee_hits=${GUARANTEE_HITS} random_types=${RANDOM_TYPES} include_nonalerts=${INCLUDE_NONALERTS} docker_send=${DOCKER_SEND} nc_wait=${NC_WAIT_SECONDS}s strict_send=${STRICT_SEND}"
echo "target=${WAZUH_SYSLOG_HOST}:${WAZUH_SYSLOG_PORT}/udp"

if [ "$RANDOM_TYPES" = "1" ]; then
  if [ "$FOREVER" -eq 1 ]; then
    loop=1
    while true; do
      send_random_event_once "$SELECTOR"
      echo "loop=${loop} complete"
      loop=$((loop + 1))
      sleep "$DELAY"
    done
  else
    for ((i=1; i<=COUNT; i++)); do
      send_random_event_once "$SELECTOR"
      echo "iteration=${i}/${COUNT} complete"
      sleep "$DELAY"
    done
  fi
elif [ "$FOREVER" -eq 1 ]; then
  loop=1
  while true; do
    for f in "${FILES[@]}"; do
      send_file_once "$f" "$SELECTOR"
    done
    if [ "$GUARANTEE_HITS" = "1" ]; then
      send_guaranteed_hits_once "$SELECTOR"
    fi
    echo "loop=${loop} complete"
    loop=$((loop + 1))
  done
else
  for ((i=1; i<=COUNT; i++)); do
    for f in "${FILES[@]}"; do
      send_file_once "$f" "$SELECTOR"
    done
    if [ "$GUARANTEE_HITS" = "1" ]; then
      send_guaranteed_hits_once "$SELECTOR"
    fi
    echo "iteration=${i}/${COUNT} complete"
  done
fi