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

log() {
  printf '[nat] %s\n' "$*"
}

usage() {
  cat <<'EOF'
Usage: ./nat.sh <up|down|status> <id-or-name-prefix>

Manage per-VM NAT rules for internet access.
EOF
}

get_prop() {
  local info="$1"
  local key="$2"
  awk -F= -v k="$key" '$1==k {print $2}' "$info"
}

find_vm_info() {
  local query="$1"
  local info match_count=0 match=""

  for info in state/vms/*/info; do
    [[ -f "$info" ]] || continue
    local id name
    id="$(get_prop "$info" "id")"
    name="$(get_prop "$info" "name")"
    if [[ "$id" == "$query"* || "$name" == "$query"* ]]; then
      match="$info"
      match_count=$((match_count + 1))
    fi
  done

  if (( match_count == 0 )); then
    log "no VM found for prefix: $query"
    exit 1
  fi
  if (( match_count > 1 )); then
    log "multiple VMs found for prefix: $query"
    exit 1
  fi

  printf '%s' "$match"
}

default_uplink() {
  ip route show default 2>/dev/null | awk '/default/ {print $5; exit}'
}

ensure_iptables() {
  if ! command -v iptables >/dev/null 2>&1; then
    log "iptables not found"
    exit 1
  fi
}

ACTION="${1:-}"
QUERY="${2:-}"
if [[ -z "$ACTION" || -z "$QUERY" ]]; then
  usage
  exit 1
fi

INFO_FILE="$(find_vm_info "$QUERY")"
GUEST_IP="$(get_prop "$INFO_FILE" "guest_ip")"
TAP="$(get_prop "$INFO_FILE" "tap")"

if [[ -z "$GUEST_IP" || -z "$TAP" ]]; then
  log "missing guest_ip or tap in $INFO_FILE"
  exit 1
fi

UPLINK="$(default_uplink)"
if [[ -z "$UPLINK" ]]; then
  log "failed to detect uplink interface"
  exit 1
fi

ensure_iptables

nat_rule=(-t nat -s "${GUEST_IP}/32" -o "$UPLINK" -j MASQUERADE)
fwd_out_rule=(-i "$TAP" -o "$UPLINK" -j ACCEPT)
fwd_in_rule=(-i "$UPLINK" -o "$TAP" -m state --state "RELATED,ESTABLISHED" -j ACCEPT)

case "$ACTION" in
  up)
    sudo sysctl -w net.ipv4.ip_forward=1 >/dev/null
    sudo iptables -t nat -C POSTROUTING "${nat_rule[@]}" 2>/dev/null || \
      sudo iptables -t nat -A POSTROUTING "${nat_rule[@]}"
    sudo iptables -C FORWARD "${fwd_out_rule[@]}" 2>/dev/null || \
      sudo iptables -A FORWARD "${fwd_out_rule[@]}"
    sudo iptables -C FORWARD "${fwd_in_rule[@]}" 2>/dev/null || \
      sudo iptables -A FORWARD "${fwd_in_rule[@]}"
    log "NAT enabled for $GUEST_IP via $UPLINK"
    ;;
  down)
    sudo iptables -t nat -C POSTROUTING "${nat_rule[@]}" 2>/dev/null && \
      sudo iptables -t nat -D POSTROUTING "${nat_rule[@]}" || true
    sudo iptables -C FORWARD "${fwd_out_rule[@]}" 2>/dev/null && \
      sudo iptables -D FORWARD "${fwd_out_rule[@]}" || true
    sudo iptables -C FORWARD "${fwd_in_rule[@]}" 2>/dev/null && \
      sudo iptables -D FORWARD "${fwd_in_rule[@]}" || true
    log "NAT disabled for $GUEST_IP"
    ;;
  status)
    sysctl net.ipv4.ip_forward | sed 's/^/[nat] /'
    if sudo iptables -t nat -C POSTROUTING "${nat_rule[@]}" 2>/dev/null; then
      log "nat: installed"
    else
      log "nat: missing"
    fi
    if sudo iptables -C FORWARD "${fwd_out_rule[@]}" 2>/dev/null; then
      log "forward out: installed"
    else
      log "forward out: missing"
    fi
    if sudo iptables -C FORWARD "${fwd_in_rule[@]}" 2>/dev/null; then
      log "forward in: installed"
    else
      log "forward in: missing"
    fi
    ;;
  *)
    usage
    exit 1
    ;;
esac