#!/usr/bin/env bash
# DOUAWÊ Remote port manager — prepare/test before production use.
set -Eeuo pipefail

WG_IF="wg0"
WG_NET="10.50.0."
PUBLIC_IF="eth0"
PORT_MIN=48292
PORT_MAX=49291
RESERVED_IP="10.50.0.2"
RESERVED_PORT=48291
BASE="/etc/douawe-vpn"
STATE_DIR="$BASE/remote-ports"
LOCK_FILE="$BASE/remote-ports.lock"
RULES_FILE="/etc/iptables/rules.v4"

usage() {
  cat <<'USAGE'
Usage:
  douawe-remote-port add <wireguard-ip> <client-id>
  douawe-remote-port remove <wireguard-ip>
  douawe-remote-port list

This tool manages TCP public-port -> MikroTik WinBox TCP/8291 forwarding.
It does not create WireGuard peers or modify port 48291 (reserved for 10.50.0.2).
USAGE
}

need_root() { [[ ${EUID:-$(id -u)} -eq 0 ]] || { echo "Run as root (sudo)." >&2; exit 1; }; }
valid_ip() {
  [[ "$1" =~ ^10\.50\.0\.([0-9]{1,3})$ ]] || return 1
  local n="${BASH_REMATCH[1]}"
  (( n >= 2 && n <= 254 )) || return 1
  [[ "$1" != "$RESERVED_IP" ]]
}
valid_id() { [[ "$1" =~ ^[A-Za-z0-9._-]{1,64}$ ]]; }

persist_rules() {
  local tmp backup
  [[ -f "$RULES_FILE" ]] || { echo "Missing persistent rules file: $RULES_FILE" >&2; return 1; }
  tmp="$(mktemp /etc/iptables/rules.v4.tmp.XXXXXX)"
  backup="${RULES_FILE}.backup-remote-$(date +%Y%m%d-%H%M%S)"
  cp -a "$RULES_FILE" "$backup"
  if ! iptables-save > "$tmp"; then rm -f "$tmp"; echo "iptables-save failed; persistent file unchanged." >&2; return 1; fi
  chmod --reference="$RULES_FILE" "$tmp"
  chown --reference="$RULES_FILE" "$tmp"
  if ! iptables-restore --test < "$tmp"; then rm -f "$tmp"; echo "Saved rules failed validation; persistent file unchanged." >&2; return 1; fi
  mv "$tmp" "$RULES_FILE"
  echo "Persistent rules updated. Backup: $backup"
}

has_peer() {
  wg show "$WG_IF" allowed-ips | awk -v ip="$1/32" '$2 == ip {found=1} END {exit !found}'
}
get_state_file() { printf '%s/%s.state' "$STATE_DIR" "$1"; }

add_port() {
  local ip="$1" client_id="$2" state port target existing
  valid_ip "$ip" || { echo "Invalid or reserved WireGuard IP (10.50.0.3-10.50.0.254)." >&2; exit 2; }
  valid_id "$client_id" || { echo "client-id must contain 1-64 letters, digits, dot, underscore or hyphen." >&2; exit 2; }
  has_peer "$ip" || { echo "No active WireGuard peer found for $ip/32. Nothing changed." >&2; exit 3; }
  state="$(get_state_file "$ip")"
  if [[ -f "$state" ]]; then
    # Idempotent: return existing mapping if the corresponding rule is still present.
    port="$(sed -n 's/^PORT=//p' "$state" | head -n1)"
    if [[ "$port" =~ ^[0-9]+$ ]] && iptables -t nat -C PREROUTING -i "$PUBLIC_IF" -p tcp --dport "$port" -j DNAT --to-destination "$ip:8291" 2>/dev/null; then
      printf 'SUCCESS=1\nIP=%s\nPORT=%s\nDESTINATION=%s:8291\n' "$ip" "$port" "$ip"
      return 0
    fi
    echo "State exists but firewall mapping is missing for $ip. Review manually; refusing to guess." >&2
    exit 4
  fi

  # Refuse duplicate mapping of same IP or any existing public port use.
  if iptables -t nat -S PREROUTING | grep -F -- "--to-destination $ip:8291" >/dev/null; then
    echo "A PREROUTING mapping to $ip:8291 already exists but has no state file. Review manually." >&2
    exit 4
  fi

  port=""
  for ((p=PORT_MIN; p<=PORT_MAX; p++)); do
    if ! iptables -t nat -S PREROUTING | grep -Eq -- "--dport ${p}([[:space:]]|$)" \
       && ! ss -H -ltn 2>/dev/null | awk '{print $4}' | grep -Eq ":${p}$"; then
      port="$p"; break
    fi
  done
  [[ -n "$port" ]] || { echo "No free port in ${PORT_MIN}-${PORT_MAX}." >&2; exit 5; }

  # Back up live IPv4 rules before changing them; remove partial rules if any command fails.
  local live_backup
  live_backup="/root/iptables-before-remote-${ip//./-}-$(date +%Y%m%d-%H%M%S).v4"
  iptables-save > "$live_backup"
  if ! {
    iptables -t nat -A PREROUTING -i "$PUBLIC_IF" -p tcp --dport "$port" -j DNAT --to-destination "$ip:8291"
    iptables -A FORWARD -d "$ip/32" -i "$PUBLIC_IF" -o "$WG_IF" -p tcp --dport 8291 -m conntrack --ctstate NEW -j ACCEPT
    iptables -t nat -A POSTROUTING -d "$ip/32" -o "$WG_IF" -p tcp --dport 8291 -j SNAT --to-source 10.50.0.1
    iptables-save > /dev/null
    persist_rules
  }; then
    iptables-restore < "$live_backup" || true
    echo "Failed to install/persist rules; live rules restored from $live_backup." >&2
    exit 6
  fi

  umask 077
  printf 'IP=%s\nPORT=%s\nCLIENT_ID=%s\nCREATED_AT=%s\n' "$ip" "$port" "$client_id" "$(date -Is)" > "$state"
  chmod 600 "$state"
  printf 'SUCCESS=1\nIP=%s\nPORT=%s\nDESTINATION=%s:8291\nCLIENT_ID=%s\n' "$ip" "$port" "$ip" "$client_id"
  echo "Live rules backup: $live_backup" >&2
}

remove_port() {
  local ip="$1" state port
  valid_ip "$ip" || { echo "Invalid or reserved WireGuard IP." >&2; exit 2; }
  state="$(get_state_file "$ip")"
  [[ -f "$state" ]] || { echo "No managed port mapping found for $ip. Nothing changed." >&2; exit 3; }
  port="$(sed -n 's/^PORT=//p' "$state" | head -n1)"
  [[ "$port" =~ ^[0-9]+$ ]] && (( port >= PORT_MIN && port <= PORT_MAX )) || { echo "Invalid port in state file; refusing to remove rules." >&2; exit 4; }
  local live_backup
  live_backup="/root/iptables-before-remote-remove-${ip//./-}-$(date +%Y%m%d-%H%M%S).v4"
  iptables-save > "$live_backup"
  # Remove only exact rules associated with this mapping.
  while iptables -t nat -C PREROUTING -i "$PUBLIC_IF" -p tcp --dport "$port" -j DNAT --to-destination "$ip:8291" 2>/dev/null; do
    iptables -t nat -D PREROUTING -i "$PUBLIC_IF" -p tcp --dport "$port" -j DNAT --to-destination "$ip:8291"
  done
  while iptables -C FORWARD -d "$ip/32" -i "$PUBLIC_IF" -o "$WG_IF" -p tcp --dport 8291 -m conntrack --ctstate NEW -j ACCEPT 2>/dev/null; do
    iptables -D FORWARD -d "$ip/32" -i "$PUBLIC_IF" -o "$WG_IF" -p tcp --dport 8291 -m conntrack --ctstate NEW -j ACCEPT
  done
  while iptables -t nat -C POSTROUTING -d "$ip/32" -o "$WG_IF" -p tcp --dport 8291 -j SNAT --to-source 10.50.0.1 2>/dev/null; do
    iptables -t nat -D POSTROUTING -d "$ip/32" -o "$WG_IF" -p tcp --dport 8291 -j SNAT --to-source 10.50.0.1
  done
  if ! persist_rules; then
    iptables-restore < "$live_backup"
    echo "Persistence failed; live rules restored. State file retained." >&2
    exit 6
  fi
  rm -f "$state"
  printf 'SUCCESS=1\nREMOVED_IP=%s\nREMOVED_PORT=%s\n' "$ip" "$port"
  echo "Live rules backup: $live_backup" >&2
}

list_ports() {
  shopt -s nullglob
  local f
  printf 'IP\tPORT\tCLIENT_ID\tCREATED_AT\n'
  for f in "$STATE_DIR"/*.state; do
    awk -F= 'BEGIN{ip="";port="";id="";created=""} $1=="IP"{ip=$2} $1=="PORT"{port=$2} $1=="CLIENT_ID"{id=$2} $1=="CREATED_AT"{created=$2} END{printf "%s\t%s\t%s\t%s\n",ip,port,id,created}' "$f"
  done
}

need_root
mkdir -p "$STATE_DIR"
chmod 700 "$STATE_DIR"
exec 9>"$LOCK_FILE"
flock -x 9
case "${1:-}" in
  add) [[ $# -eq 3 ]] || { usage; exit 2; }; add_port "$2" "$3" ;;
  remove) [[ $# -eq 2 ]] || { usage; exit 2; }; remove_port "$2" ;;
  list) [[ $# -eq 1 ]] || { usage; exit 2; }; list_ports ;;
  *) usage; exit 2 ;;
esac