#!/usr/bin/env bash
# Traffic control for AmneziaWG (production-ready)

set -euo pipefail

LOG="/var/log/amnezia-tc.log"
LOCK_FILE="/tmp/amnezia-tc.lock"
LIMIT="${LIMIT:-70mbit}"
CONTAINER=""
WG_IF="${WG_IF:-wg0}"

# Prefer new container name, fall back to old
for c in amnezia-awg2 amnezia-awg; do
  if docker ps --format '{{.Names}}' 2>/dev/null | grep -xq "$c"; then
    CONTAINER="$c"
    break
  fi
done

log() {
  echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" | tee -a "$LOG"
}

# Lock to avoid concurrent runs
exec 200>"$LOCK_FILE"
flock -n 200 || { log "Script already running, exit"; exit 0; }

if [ -z "$CONTAINER" ]; then
  log "Container amnezia-awg2/amnezia-awg not found, exit"
  exit 0
fi

# detect interface inside container if default not present there
if ! docker exec "$CONTAINER" ip link show "$WG_IF" >/dev/null 2>&1; then
  DET_IF=$(docker exec "$CONTAINER" sh -c "ip -o link show 2>/dev/null | awk -F': ' '{print $2}'" 2>/dev/null \
           | egrep -o '^(wg[0-9]+|wg|awg[0-9]*|amnezia[0-9]*|tun[0-9]+)' | head -n1 || true)

  if [ -n "$DET_IF" ]; then
    WG_IF="$DET_IF"
    log "Detected interface inside container: $WG_IF"
  else
    log "Interface $WG_IF not found inside container $CONTAINER, exit"
    exit 0
  fi
fi

log "=== Start updating rules === (container=$CONTAINER, iface=$WG_IF, limit=$LIMIT)"

# Get current clients
CLIENT_IPS=$(docker exec "$CONTAINER" sh -c "wg show all allowed-ips 2>/dev/null" 2>/dev/null \
           | awk '{print $3}' | grep -oE '10\.8\.1\.[0-9]{1,3}' | sort -u || true)

# fallback
if [ -z "$CLIENT_IPS" ]; then
  CLIENT_IPS=$(docker exec "$CONTAINER" sh -c "wg show 2>/dev/null" 2>/dev/null \
             | grep -Eo '10\.8\.1\.[0-9]{1,3}' | sort -u || true)
fi

if [ -z "$CLIENT_IPS" ]; then
  log "No clients found, exit"
  exit 0
fi

log "Clients: $(echo $CLIENT_IPS | tr '\n' ' ')"

EXISTING_RULES=$(docker exec "$CONTAINER" sh -c "tc filter show dev $WG_IF parent 1: 2>/dev/null || true" 2>/dev/null \
               | grep -oE 'match [0-9a-f]+' | awk '{print $2}' | sort -u || true)

NEEDS_UPDATE=false

num_ips=$(echo "$CLIENT_IPS" | wc -w)
num_rules=$(echo "$EXISTING_RULES" | wc -w)
if [ "$num_rules" -ne "$num_ips" ] || [ -z "$EXISTING_RULES" ]; then
  NEEDS_UPDATE=true
else
  for ip in $CLIENT_IPS; do
    hex_ip=$(printf '%02x' ${ip//./ } | tr -d '\n' || true)
    if ! echo "$EXISTING_RULES" | grep -q "$hex_ip"; then
      NEEDS_UPDATE=true
      break
    fi
  done
fi

if [ "$NEEDS_UPDATE" = false ]; then
  log "Rules are already up-to-date, skip"
  exit 0
fi

log "Updating rules (old will be removed)"

docker exec "$CONTAINER" sh -c "
  tc qdisc del dev $WG_IF root 2>/dev/null || true
  tc qdisc del dev $WG_IF ingress 2>/dev/null || true
  ip link del ifb0 2>/dev/null || true
" 2>/dev/null || true

docker exec "$CONTAINER" sh -c "
  tc qdisc add dev $WG_IF root handle 1: htb 2>/dev/null || true
  tc class add dev $WG_IF parent 1: classid 1:1 htb rate 1000mbit ceil 1000mbit 2>/dev/null || true
" 2>/dev/null || true

docker exec "$CONTAINER" sh -c "
  modprobe ifb 2>/dev/null || true
  ip link add name ifb0 type ifb 2>/dev/null || true
  ip link set dev ifb0 up 2>/dev/null || true

  tc qdisc add dev $WG_IF handle ffff: ingress 2>/dev/null || true
  tc filter add dev $WG_IF parent ffff: protocol all u32 match u32 0 0 action mirred egress redirect dev ifb0 2>/dev/null || true

  tc qdisc add dev ifb0 root handle 2: htb 2>/dev/null || true
  tc class add dev ifb0 parent 2: classid 2:1 htb rate 1000mbit ceil 1000mbit 2>/dev/null || true
" 2>/dev/null || true

class_id=10
count=0
for ip in $CLIENT_IPS; do
  log "Adding limit for $ip (class $class_id)"

  docker exec "$CONTAINER" sh -c "
    tc class add dev $WG_IF parent 1:1 classid 1:$class_id htb rate $LIMIT ceil $LIMIT 2>/dev/null || true
    tc qdisc add dev $WG_IF parent 1:$class_id fq_codel 2>/dev/null || true
    tc filter add dev $WG_IF parent 1: protocol ip u32 match ip dst $ip flowid 1:$class_id 2>/dev/null || true
  " 2>/dev/null || true

  docker exec "$CONTAINER" sh -c "
    tc class add dev ifb0 parent 2:1 classid 2:$class_id htb rate $LIMIT ceil $LIMIT 2>/dev/null || true
    tc qdisc add dev ifb0 parent 2:$class_id fq_codel 2>/dev/null || true
    tc filter add dev ifb0 parent 2: protocol ip u32 match ip src $ip flowid 2:$class_id 2>/dev/null || true
  " 2>/dev/null || true

  class_id=$((class_id + 1))
  count=$((count + 1))
done

log "Limits applied for $count clients"
log "=== Completed successfully ==="
exit 0