#!/bin/sh
# wg-provision.sh — mt7628 网关侧 WireGuard 收敛器 (注册之外的全部设备侧逻辑)
#
# 注册 (POST /api/register) 由 ha_install_register 负责,它把校验过的响应写到
#   $WG_DIR/register_resp.json,并把持久私钥写到 $WG_DIR/gw_private.key。
# 本脚本只负责把设备实况收敛到该响应描述的目标状态,并保证幂等。
#
# 用法:
#   wg-provision preflight   只读探测 + 依赖/能力自检,绝不改动设备
#   wg-provision reconcile   收敛: 生成配置 -> 拉起隧道 -> 装防火墙 -> 校验实况
#
# 退出码 (与 ha_install_register 的 ClassifyProvisionExit 及设计 §16 契约一致):
#   0        成功 (调用方重置退避)
#   10       缺少依赖 / 配置非法              (永久错误,等待下一次事件)
#   11       注册响应缺失或非法                (永久错误,等待下一次事件)
#   12       内核 WireGuard / fw4 表 / 链 / 前缀-NAT 不支持 (永久错误)
#   13       WireGuard 配置失败 (wg)          (可重试)
#   14       接口/地址/路由/转发失败 (ip)      (可重试)
#   15       防火墙生成或语法校验失败          (可重试)
#   16       fw4 reload 或实况规则校验失败(已回滚) (可重试)
#   17       另一个实例持有锁                  (可重试)
#   2        用法错误
#
# 测试注入 (host 侧无 root 时):
#   WG_ROOT      所有可变路径的前缀 (默认空 = 真实 /)
#   PATH         测试把 fake 的 wg/ip/nft/fw4/uci/sysctl 目录前置到 PATH
#   FAKE_*       fake 命令自身识别的失败开关 (脚本不感知)

set -eu
umask 077

# ---- 退出码常量 ----
E_DEP=10
E_RESP=11
E_CAP=12
E_WG=13
E_IP=14
E_FWGEN=15
E_FW4=16
E_LOCK=17
E_USAGE=2

# ---- 可配置路径 (全部经 WG_ROOT 前缀,便于测试重定向) ----
WG_ROOT=${WG_ROOT:-}
WG_DIR=$WG_ROOT/etc/wireguard
PRIV=$WG_DIR/gw_private.key
CONF=$WG_DIR/wg0.conf
RESP=$WG_DIR/register_resp.json
FW_FILE=$WG_ROOT/etc/firewall.wg
CAMERA_FILE=$WG_DIR/cameras.json
SYSCTL_FILE=$WG_ROOT/etc/sysctl.conf
HOTPLUG=$WG_ROOT/etc/hotplug.d/iface/99-wg-provision
KERNEL_WG=$WG_ROOT/sys/module/wireguard
LOCKDIR=${WG_LOCKDIR:-${WG_ROOT:-}/var/run/wg-provision.lock}

WG_MTU=${WG_MTU:-1420}
WGIF=wg0

# 防火墙规则的唯一注释,用于幂等探测与实况校验
C_DNAT=wg-provision-dnat
C_FORWARD=wg-provision-forward
C_MASQ=wg-provision-masq

log() {
  echo "[wg] $*"
  logger -t wg-provision "$*" 2>/dev/null || true
}
err() {
  echo "[wg][ERR] $*" >&2
  logger -t wg-provision "ERR: $*" 2>/dev/null || true
}
die() {
  code=$1
  shift
  err "$*"
  exit "$code"
}

# ---- 从注册响应里取一个字符串字段 (扁平 string 字段,兼容紧凑/美化 JSON) ----
json_str() {
  sed -n "s/.*\"$2\"[[:space:]]*:[[:space:]]*\"\\([^\"]*\\)\".*/\\1/p" "$1" | head -1
}

# ---- 依赖检查 (只读) ----
need_deps() {
  miss=
  for c in wg nft ip fw4 uci sysctl jsonfilter; do
    command -v "$c" >/dev/null 2>&1 || miss="$miss $c"
  done
  [ -z "$miss" ] || die "$E_DEP" "missing dependencies:$miss (opkg install wireguard-tools nftables)"
}

# ---- 内核 WireGuard 支持 (只读) ----
need_kernel_wg() {
  [ -e "$KERNEL_WG" ] || die "$E_CAP" "kernel WireGuard module not present ($KERNEL_WG)"
}

# ---- fw4 表与三条基础链存在 (只读) ----
need_fw4() {
  nft list tables 2>/dev/null | grep -q 'inet fw4' ||
    die "$E_CAP" "nftables 'inet fw4' table not found (fw4 firmware required)"
  for ch in dstnat forward srcnat; do
    nft list chain inet fw4 "$ch" >/dev/null 2>&1 ||
      die "$E_CAP" "required fw4 chain missing: $ch"
  done
}

# ---- 前缀-NAT (netmap) 能力 (只读: nft --check 不落地) ----
need_netmap() {
  printf 'add rule inet fw4 dstnat ip daddr 198.18.0.0/24 dnat ip prefix to ip daddr map { 198.18.0.0/24 : 192.168.0.0/24 }\n' |
    nft --check -f - >/dev/null 2>&1 ||
    die "$E_CAP" "kernel nftables lacks prefix NAT (netmap) support"
}

# ---- 只读探测上联口与其 IPv4 (排除 APIPA);失败即 no-uplink ----
probe_uplink() {
  UPLINK=$(ip route show default 2>/dev/null | sort -k11 -n | head -1 |
    sed -n 's/.* dev \([^ ]*\).*/\1/p')
  [ -n "$UPLINK" ] || return 1
  cidr=$(ip -o -f inet addr show dev "$UPLINK" 2>/dev/null |
    sed -n 's/.* inet \([0-9.]*\/[0-9]*\).*/\1/p' | head -1)
  [ -n "$cidr" ] || return 1
  UPLINK_IP=${cidr%/*}
  case "$UPLINK_IP" in 169.254.*) return 1 ;; esac
  return 0
}

# ---- 解析并校验注册响应 ----
parse_resp() {
  [ -f "$RESP" ] || die "$E_RESP" "register response missing: $RESP"
  TENANT_IP=$(json_str "$RESP" gateway_tunnel_ip)
  TENANT_NET=$(json_str "$RESP" tenant_subnet)
  HOME_NET=$(json_str "$RESP" home_subnet) # 192.168.x.0/24
  SRV_PUB=$(json_str "$RESP" server_public_key)
  SRV_EP=$(json_str "$RESP" server_endpoint)
  SRV_AIP=$(json_str "$RESP" server_allowed_ips)
  for v in TENANT_IP TENANT_NET HOME_NET SRV_PUB SRV_EP SRV_AIP; do
    eval "val=\$$v"
    [ -n "$val" ] || die "$E_RESP" "register response missing field: $v"
  done
  SUPERNET=$SRV_AIP # AllowedIPs == tunnel supernet
}

# ---- 开启并持久化 ip_forward ----
ensure_forwarding() {
  cur=$(sysctl -n net.ipv4.ip_forward 2>/dev/null || echo 0)
  if [ "$cur" != 1 ]; then
    sysctl -w net.ipv4.ip_forward=1 >/dev/null 2>&1 || die "$E_IP" "sysctl ip_forward failed"
  fi
  mkdir -p "$(dirname "$SYSCTL_FILE")"
  grep -q '^net.ipv4.ip_forward=1' "$SYSCTL_FILE" 2>/dev/null ||
    echo 'net.ipv4.ip_forward=1' >>"$SYSCTL_FILE"
}

# ---- 生成 wg0.conf 到临时文件 (0600);先应用/校验,成功后才替换持久文件 ----
prepare_conf() {
  [ -f "$PRIV" ] || die "$E_RESP" "private key missing: $PRIV (written by registration)"
  mkdir -p "$WG_DIR"
  CONF_TMP="$CONF.tmp.$$"
  {
    echo "[Interface]"
    echo "PrivateKey = $(cat "$PRIV")"
    echo "Address = $TENANT_IP"
    echo ""
    echo "[Peer]"
    echo "PublicKey = $SRV_PUB"
    echo "Endpoint = $SRV_EP"
    echo "AllowedIPs = $SRV_AIP"
    echo "PersistentKeepalive = 25"
  } >"$CONF_TMP"
  chmod 600 "$CONF_TMP"
}

# ---- 拉起/修复隧道链路,逐项校验实况并按需修复;仅在应用成功后替换持久 wg0.conf ----
bringup() {
  ip link show "$WGIF" >/dev/null 2>&1 ||
    ip link add dev "$WGIF" type wireguard ||
    die "$E_IP" "ip link add $WGIF failed"

  # 从临时 conf(去掉 Address 行)应用 peer;wg setconf 会校验语法/密钥
  setconf="$CONF_TMP.setconf"
  grep -v '^Address' "$CONF_TMP" >"$setconf"
  wg setconf "$WGIF" "$setconf" || {
    rm -f "$setconf"
    die "$E_WG" "wg setconf failed"
  }
  rm -f "$setconf"
  wg show "$WGIF" dump 2>/dev/null | grep -q "$SRV_PUB" ||
    die "$E_WG" "peer public key not present after setconf"

  # peer 已成功应用 -> 现在才把临时 conf 提升为持久文件 (失败则保留旧文件)
  if [ ! -f "$CONF" ] || ! cmp -s "$CONF_TMP" "$CONF"; then
    mv -f "$CONF_TMP" "$CONF"
  else
    rm -f "$CONF_TMP"
  fi

  # 地址: 必须恰好等于 tenant tunnel IP
  want_addr=$TENANT_IP
  have_addr=$(ip -o -f inet addr show dev "$WGIF" 2>/dev/null |
    sed -n 's/.* inet \([0-9./]*\).*/\1/p' | head -1)
  if [ "$have_addr" != "$want_addr" ]; then
    [ -z "$have_addr" ] || ip addr del "$have_addr" dev "$WGIF" 2>/dev/null || true
    ip addr add "$want_addr" dev "$WGIF" || die "$E_IP" "ip addr add failed"
  fi

  ip link set dev "$WGIF" mtu "$WG_MTU" up || die "$E_IP" "ip link set mtu/up failed"

  # supernet 路由: replace 幂等地新增或覆盖
  ip route replace "$SUPERNET" dev "$WGIF" || die "$E_IP" "ip route replace $SUPERNET failed"
}

# ---- 生成 fw4 include 到临时文件并语法校验,变化时原子替换;返回 0=有变化 1=未变 ----
# 规则严格限定在探测到的家庭 /24,不暴露网关自身服务 (设计 §14/§18)。
gen_firewall() {
  mkdir -p "$(dirname "$FW_FILE")"
  tmp="$FW_FILE.tmp.$$"
  cat >"$tmp" <<EOF
#!/bin/sh
# wg-provision fw4 include — 幂等; 每次运行先按注释删旧规则再加新规则,
# 网络环境变化 (home/tenant 变更) 时旧规则必被清理,新规则必定生效,
# 不依赖 fw4 reload 是否 flush。
wg_del() {  # \$1=chain \$2=注释(唯一)  删除该链上所有带此注释的规则
  for h in \$(nft -a list chain inet fw4 "\$1" 2>/dev/null | grep -F "\$2" | sed -n 's/.*# handle \([0-9]*\)\$/\1/p'); do
    nft delete rule inet fw4 "\$1" handle "\$h" 2>/dev/null || true
  done
}
wg_del dstnat wg-camera-dnat
if [ -r "$CAMERA_FILE" ] && \
   [ "\$(jsonfilter -i "$CAMERA_FILE" -e '@.tenant_subnet' 2>/dev/null)" = "$TENANT_NET" ]; then
  index=0
  while :; do
    alias=\$(jsonfilter -i "$CAMERA_FILE" -e "@.entries[\$index].alias" 2>/dev/null) || break
    real_ip=\$(jsonfilter -i "$CAMERA_FILE" -e "@.entries[\$index].real_ip" 2>/dev/null) || break
    [ -n "\$alias" ] && [ -n "\$real_ip" ] && \
      nft insert rule inet fw4 dstnat iifname \"$WGIF\" ip daddr "\$alias" dnat to "\$real_ip" comment \"wg-camera-dnat\"
    index=\$((index + 1))
  done
fi
wg_del dstnat $C_DNAT
nft add rule inet fw4 dstnat iifname \\"$WGIF\\" ip daddr $TENANT_NET dnat ip prefix to ip daddr map { $TENANT_NET : $HOME_NET } comment \\"$C_DNAT\\"
wg_del forward $C_FORWARD
nft insert rule inet fw4 forward iifname \\"$WGIF\\" ip daddr $HOME_NET oifname != \\"$WGIF\\" accept comment \\"$C_FORWARD\\"
wg_del srcnat $C_MASQ
nft insert rule inet fw4 srcnat ip saddr $SUPERNET ip daddr $HOME_NET oifname != \\"$WGIF\\" masquerade comment \\"$C_MASQ\\"
EOF
  chmod 600 "$tmp"
  sh -n "$tmp" || {
    rm -f "$tmp"
    die "$E_FWGEN" "generated firewall include failed syntax check"
  }
  if [ -f "$FW_FILE" ] && cmp -s "$tmp" "$FW_FILE"; then
    rm -f "$tmp"
    return 1
  fi
  mv -f "$tmp" "$FW_FILE"
  return 0
}

# ---- UCI 里登记 fw4 include (去重,恰好一条) ----
ensure_uci_include() {
  if uci show firewall 2>/dev/null | grep -qF "path='$FW_FILE'"; then
    return 0
  fi
  uci add firewall include >/dev/null
  uci set firewall.@include[-1].type='script'
  uci set firewall.@include[-1].path="$FW_FILE"
  uci set firewall.@include[-1].fw4_compatible='1'
  uci commit firewall
  log "registered fw4 include $FW_FILE"
}

# ---- 实况防火墙三条唯一注释规则都在,且 DNAT 家庭段与目标一致 ----
fw_rules_live() {
  nft list table inet fw4 2>/dev/null | grep -q "$C_DNAT" || return 1
  nft list table inet fw4 2>/dev/null | grep -q "$C_FORWARD" || return 1
  nft list table inet fw4 2>/dev/null | grep -q "$C_MASQ" || return 1
  nft list chain inet fw4 dstnat 2>/dev/null | grep "$C_DNAT" | grep -q "$HOME_NET" || return 1
  return 0
}

# ---- 应用防火墙: 仅在文件变化或实况缺失时 reload;reload 或校验失败则回滚 (§15) ----
apply_firewall() {
  changed=$1 # gen_firewall 的返回: 0=变化 1=未变
  ensure_uci_include
  if [ "$changed" = 0 ] || ! fw_rules_live; then
    if fw4 reload >/dev/null 2>&1 && fw_rules_live; then
      cp -f "$FW_FILE" "$FW_FILE.applied" 2>/dev/null || true
      log "firewall reloaded (dnat $TENANT_NET -> $HOME_NET)"
    else
      # reload 或实况校验失败: 恢复上次成功应用的 include 并 reload,再报失败
      if [ -f "$FW_FILE.applied" ]; then cp -f "$FW_FILE.applied" "$FW_FILE" 2>/dev/null || true; fi
      fw4 reload >/dev/null 2>&1 || true
      die "$E_FW4" "fw4 reload/verify failed; rolled back firewall include"
    fi
  else
    log "firewall unchanged; no reload"
  fi
}

# ---- 安装通知型 hotplug 钩子 (只发通知触发重收敛,不改 rc.local) ----
install_hotplug() {
  mkdir -p "$(dirname "$HOTPLUG")"
  tmp="$HOTPLUG.tmp.$$"
  cat >"$tmp" <<EOF
#!/bin/sh
# wg-provision 通知钩子: 上联 ifup/ifupdate 时请求控制端口重收敛
case "\$ACTION" in
  ifup|ifupdate)
    curl -fsS --max-time 2 -X POST \\
      http://127.0.0.1:18097/api/v1/wireguard/reconcile >/dev/null 2>&1 &
    ;;
esac
exit 0
EOF
  chmod 755 "$tmp"
  if [ ! -f "$HOTPLUG" ] || ! cmp -s "$tmp" "$HOTPLUG"; then mv -f "$tmp" "$HOTPLUG"; else rm -f "$tmp"; fi
}

# ---- 只读预检 ----
do_preflight() {
  need_deps
  need_kernel_wg
  need_fw4
  need_netmap
  if probe_uplink; then log "uplink ready: $UPLINK / $UPLINK_IP"; else die "$E_DEP" "no uplink LAN IP yet"; fi
  parse_resp
  log "preflight ok; device unchanged (home=$HOME_NET tenant=$TENANT_NET)"
}

# ---- 收敛 ----
do_reconcile() {
  # 互斥锁 (mkdir 原子; 持有者退出时清理)
  if ! mkdir "$LOCKDIR" 2>/dev/null; then
    die "$E_LOCK" "another wg-provision holds the lock ($LOCKDIR)"
  fi
  trap 'rmdir "$LOCKDIR" 2>/dev/null || true' EXIT INT TERM

  need_deps
  need_kernel_wg
  need_fw4
  need_netmap
  parse_resp

  ensure_forwarding
  prepare_conf
  bringup
  if gen_firewall; then fw_changed=0; else fw_changed=1; fi
  apply_firewall "$fw_changed"
  install_hotplug

  fw_rules_live || die "$E_FW4" "firewall rules not live after apply"
  log "reconcile done"
}

case "${1:-}" in
preflight) do_preflight ;;
reconcile) do_reconcile ;;
*)
  err "usage: wg-provision preflight|reconcile"
  exit "$E_USAGE"
  ;;
esac
