#!/usr/bin/env bash
#
# traffic-guard.sh — 月度出站流量预算看门狗（为 GCP Standard Tier 每结算账号每月 200 GiB 免费出站设计）
#
# 目的：在本机本月出站流量逼近免费额度前，【只停掉代理类服务】（xray / caddy / hysteria-server），
# 避免超额计费；进入新的结算月后，只把"本脚本自己停掉的服务"重新拉起来。
#
# 绝不关机：本脚本只调用 `systemctl is-active|stop|start <GUARD_SERVICES 里的服务>`，
# 永远不会 poweroff / shutdown / halt / reboot，也不会停掉 VM（否则会连 SSH 一起丢，机器失联）。
#
# 用法：
#   traffic-guard.sh [check]        # 默认；由 systemd timer 周期调用（幂等）
#   traffic-guard.sh status         # 打印 月份 / 已用 / 剩余 / 阈值 / 是否已停（只读，不改状态）
#   traffic-guard.sh set-used <GiB> # 把本月已用量校准为某个值（例如照 GCP 账单控制台的数字）
#   traffic-guard.sh resume         # 手动恢复：拉起本脚本停掉的服务并清除标记（不改本月累计）
#   加 --dry-run（或 GUARD_DRY_RUN=1）：不调用 systemctl、不写状态/标记，只打印将执行的动作
#
# 环境变量（默认值）：
#   GUARD_LIMIT_GIB=200        免费额度（用于 status 的"剩余"）
#   GUARD_STOP_AT_GIB=185      累计 >= 该值就停服务（留 15 GiB 余量给统计误差、停服前的在途流量）
#   GUARD_SERVICES="xray caddy hysteria-server"
#   GUARD_TZ=America/Los_Angeles   结算月的时区（见下"账期"）
#   GUARD_STATE_DIR=/var/lib/traffic-guard   状态目录（状态文件 state，标记文件 STOPPED）
#   GUARD_SYSFS_NET=/sys/class/net           （测试用）网卡统计目录
#   GUARD_EXCLUDE_IFACES=""    （可选）不计入的网卡名正则，如 '^(docker0|veth.*|br-.*)$'
#   GUARD_MONTH=""             （测试用）强制当前月份 YYYY-MM
#   （GUARD_INTERVAL 只由安装器 setup_traffic_guard.sh 用于生成 timer）
#
# 计量：累加所有【UP 且非 loopback】网卡 /sys/class/net/*/statistics/tx_bytes 的增量。
#   * tx_bytes 统计的是网卡发出的全部字节（含 L2/L3 头、重传、内网/VPC 流量、回环以外的隧道/容器桥
#     流量），而 GCP 只对"到互联网的出站"按载荷计费 —— 所以 tx_bytes 会【略微多算】，
#     偏向保守（宁可早停，不会漏算），这是有意接受的安全方向。
#   * 免费额度是【每结算账号】共享的：同账号下别的 VM / 项目的出站不在本机计量里。
#     本脚本只能看到本机，若同账号还有其他出站流量，请调低 GUARD_STOP_AT_GIB 或用 set-used 校准。
#   * 增量法：每次与上次各网卡读数相减；读数变小（重启/网卡重建，计数器清零）则把当前读数当增量。
#     上次检查到重启之间的那一小段流量会漏记（最多一个 GUARD_INTERVAL 的量，由 15 GiB 余量吸收）。
#   * 首次运行没有历史：本月从 0 开始计（不把"开机以来"的历史流量算进来，避免误停）。
#     如已知本月已用量，请用 `set-used` 校准。
#   * 实测校准（2026-11-07 gcoregon）：取 GCP Monitoring 的
#     compute.googleapis.com/instance/network/sent_bytes_count，与本站 tx_bytes 增量法在
#     同一时间窗口（boot 后到查询时刻）的读数分别是 4,096,497,518 / 4,137,694,034 bytes，
#     相差不到 1% —— 说明【计量手段本身够准】，误差主要来自 set-used 的初始估计。
#     查真实用量的办法（gcloud 没有 time-series 子命令，只能打 Monitoring v3 REST API，
#     delta 指标需要把返回的点逐个求和）：
#       tok=$(gcloud auth print-access-token)
#       curl -s -H "Authorization: Bearer $tok" --get \
#         https://monitoring.googleapis.com/v3/projects/<PROJECT>/timeSeries \
#         --data-urlencode 'filter=metric.type="compute.googleapis.com/instance/network/sent_bytes_count"' \
#         --data-urlencode 'interval.startTime=2026-11-01T07:00:00Z' \
#         --data-urlencode "interval.endTime=$(date -u +%Y-%m-%dT%H:%M:%SZ)"
#     注意 Monitoring 指标比计费口径【宽】（含到 Google API / 同区 VM 的流量），是上界。
#
# 账期：GCP 结算月是自然月，用量按太平洋时间（America/Los_Angeles）切月。默认按 PT 切月 ——
#   万一 GCP 实际按 UTC 切，PT 的月初比 UTC 晚 7~8 小时，我们只会"多算几小时"（安全方向）；
#   反过来若默认 UTC 而 GCP 按 PT，则会在 GCP 仍属旧月的几个小时里提前清零（不安全）。
#   需要别的时区用 GUARD_TZ。
#
# 到阈值：停掉 GUARD_SERVICES 里【当前处于 active】的服务，写标记文件 $GUARD_STATE_DIR/STOPPED，
#   日志只记一次；之后每次检查若服务仍是停的 → 什么都不做、不刷日志（幂等）。若期间重启机器导致
#   服务被开机拉起，下一次检查会再次停掉并记一条日志。
# 新账期：清除标记，只 start 本脚本自己停掉的服务（记录在 state 的 stopped_services）；
#   从没被本脚本停过的服务（例如你手动停的、或根本没装的）一概不碰。我们只 stop 不 disable，
#   所以"恢复"就是 start，不会改动任何 unit 的 enable 状态。

set -euo pipefail

# 手动执行时（非 systemd 环境）也从安装器写的配置文件里取值，否则 status 会显示内置默认、
# 与实际 timer 生效的配置不一致。逐行精确匹配 KEY="value"，不 source（避免被文件内容执行）。
guard_env_file="${GUARD_ENV_FILE:-/etc/traffic-guard.env}"
if [ -r "$guard_env_file" ]; then
  _gv() { sed -n "s/^$1=\"\\(.*\\)\"\$/\\1/p" "$guard_env_file" 2>/dev/null | head -1; }
  [ -n "${GUARD_LIMIT_GIB:-}" ] || GUARD_LIMIT_GIB="$(_gv GUARD_LIMIT_GIB)"
  [ -n "${GUARD_STOP_AT_GIB:-}" ] || GUARD_STOP_AT_GIB="$(_gv GUARD_STOP_AT_GIB)"
  [ -n "${GUARD_SERVICES:-}" ] || GUARD_SERVICES="$(_gv GUARD_SERVICES)"
  [ -n "${GUARD_TZ:-}" ] || GUARD_TZ="$(_gv GUARD_TZ)"
  [ -n "${GUARD_INTERVAL:-}" ] || GUARD_INTERVAL="$(_gv GUARD_INTERVAL)"
  unset -f _gv
fi

GUARD_LIMIT_GIB="${GUARD_LIMIT_GIB:-200}"
GUARD_STOP_AT_GIB="${GUARD_STOP_AT_GIB:-185}"
GUARD_SERVICES="${GUARD_SERVICES:-xray caddy hysteria-server}"
GUARD_TZ="${GUARD_TZ:-America/Los_Angeles}"
GUARD_STATE_DIR="${GUARD_STATE_DIR:-/var/lib/traffic-guard}"
GUARD_SYSFS_NET="${GUARD_SYSFS_NET:-/sys/class/net}"
GUARD_EXCLUDE_IFACES="${GUARD_EXCLUDE_IFACES:-}"
GUARD_MONTH="${GUARD_MONTH:-}"
GUARD_DRY_RUN="${GUARD_DRY_RUN:-}"

STATE_FILE="$GUARD_STATE_DIR/state"
MARKER_FILE="$GUARD_STATE_DIR/STOPPED"
LOCK_FILE="$GUARD_STATE_DIR/lock"

log() { printf 'traffic-guard: %s\n' "$*"; }
die() { printf 'traffic-guard: ERROR: %s\n' "$*" >&2; exit 1; }

# ---------- 参数 ----------
CMD="check"; ARG=""
while [ $# -gt 0 ]; do
  case "$1" in
    --dry-run) GUARD_DRY_RUN=1;;
    check|status|resume) CMD="$1";;
    set-used) CMD="set-used"; shift; ARG="${1:-}";;
    -h|--help) sed -n '2,/^set -euo/p' "$0" | sed '$d'; exit 0;;
    *) die "未知参数：$1（check|status|set-used <GiB>|resume|--dry-run）";;
  esac
  shift
done

# ---------- 校验 ----------
num_re='^[0-9]+(\.[0-9]+)?$'
[[ "$GUARD_LIMIT_GIB" =~ $num_re ]] || die "GUARD_LIMIT_GIB 必须是数字：$GUARD_LIMIT_GIB"
[[ "$GUARD_STOP_AT_GIB" =~ $num_re ]] || die "GUARD_STOP_AT_GIB 必须是数字：$GUARD_STOP_AT_GIB"
gib_to_bytes() { awk -v g="$1" 'BEGIN{printf "%.0f", g*1073741824}'; }
LIMIT_BYTES="$(gib_to_bytes "$GUARD_LIMIT_GIB")"
STOP_BYTES="$(gib_to_bytes "$GUARD_STOP_AT_GIB")"
[ "$STOP_BYTES" -le "$LIMIT_BYTES" ] || die "GUARD_STOP_AT_GIB($GUARD_STOP_AT_GIB) 不能大于 GUARD_LIMIT_GIB($GUARD_LIMIT_GIB)"
# 服务名只许常规 unit 名；明确拒绝 .target 与任何关机/重启类名字 —— 本脚本绝不能把机器停掉
for _svc in $GUARD_SERVICES; do
  [[ "$_svc" =~ ^[A-Za-z0-9@_.:-]+$ ]] || die "GUARD_SERVICES 含非法服务名：$_svc"
  case "$_svc" in
    *.target|poweroff*|halt*|reboot*|shutdown*|kexec*|suspend*|hibernate*)
      die "GUARD_SERVICES 不允许包含 $_svc（本脚本绝不关机/重启）";;
  esac
done
unset _svc

# ---------- systemctl 封装（dry-run 不执行） ----------
svc_active() { systemctl is-active --quiet "$1" 2>/dev/null; }
svc_stop()   { if [ -n "$GUARD_DRY_RUN" ]; then log "[dry-run] 将停止 $1"; return 0; fi; systemctl stop "$1"; }
svc_start()  { if [ -n "$GUARD_DRY_RUN" ]; then log "[dry-run] 将启动 $1"; return 0; fi; systemctl start "$1"; }

# ---------- 状态文件（key=value，不 source，逐键 awk 读取） ----------
state_get() {
  [ -f "$STATE_FILE" ] || return 0
  awk -F= -v k="$1" '$1==k{sub(/^[^=]*=/,""); print; exit}' "$STATE_FILE"
}

current_month() {
  if [ -n "$GUARD_MONTH" ]; then printf '%s' "$GUARD_MONTH"; else TZ="$GUARD_TZ" date +%Y-%m; fi
}

# 当前各网卡读数：CUR_TX[name]=tx_bytes，CUR_UP[name]=1/0（UP 且非 loopback 才记 1）
declare -A CUR_TX=() CUR_UP=()
snapshot_ifaces() {
  CUR_TX=(); CUR_UP=()
  local d name flags up tx
  for d in "$GUARD_SYSFS_NET"/*; do
    [ -d "$d" ] || continue
    name="${d##*/}"
    [ "$name" = lo ] && continue
    tx="$(cat "$d/statistics/tx_bytes" 2>/dev/null || true)"
    [[ "$tx" =~ ^[0-9]+$ ]] || continue
    up=0
    if [ -r "$d/flags" ]; then
      flags="$(cat "$d/flags")"
      # IFF_UP=0x1，IFF_LOOPBACK=0x8
      if [[ "$flags" =~ ^0x[0-9a-fA-F]+$ ]] && (( (flags & 1) == 1 && (flags & 8) == 0 )); then up=1; fi
    elif [ "$(cat "$d/operstate" 2>/dev/null || true)" = up ]; then
      up=1
    fi
    if [ -n "$GUARD_EXCLUDE_IFACES" ] && [[ "$name" =~ $GUARD_EXCLUDE_IFACES ]]; then up=0; fi
    CUR_TX[$name]="$tx"; CUR_UP[$name]="$up"
  done
}

# 计算 "状态里的累计 + 自上次以来的增量"，结果放 NEW_USED。
# 所有非 lo 网卡都登记 last（含 DOWN 的），这样网卡 UP 起来时不会把历史读数一次性算成增量。
NEW_USED=0
compute_usage() {
  local base="$1" name last cur delta
  NEW_USED="$base"
  for name in "${!CUR_TX[@]}"; do
    cur="${CUR_TX[$name]}"
    last="$(state_get "if:$name")"
    if [ -z "$last" ]; then delta="$cur"                 # 本月中途新出现的网卡：计数从它诞生起算
    elif [ "$cur" -ge "$last" ]; then delta=$((cur - last))
    else delta="$cur"; fi                                 # 读数变小 = 计数器清零（重启/网卡重建）
    [ "${CUR_UP[$name]}" = 1 ] && NEW_USED=$((NEW_USED + delta))
  done
  return 0
}

# 写状态（原子替换）。参数：month used stopped stopped_services stopped_at
write_state() {
  [ -z "$GUARD_DRY_RUN" ] || return 0
  local tmp name
  mkdir -p "$GUARD_STATE_DIR"
  tmp="$(mktemp "$GUARD_STATE_DIR/state.XXXXXX")"
  {
    echo "month=$1"
    echo "used_bytes=$2"
    echo "stopped=$3"
    echo "stopped_services=$4"
    echo "stopped_at=$5"
    for name in "${!CUR_TX[@]}"; do echo "if:$name=${CUR_TX[$name]}"; done
  } > "$tmp"
  chmod 600 "$tmp"
  mv -f "$tmp" "$STATE_FILE"
}

fmt_gib() { awk -v b="$1" 'BEGIN{printf "%.2f", b/1073741824}'; }

# 恢复：只 start 本脚本记录在案的服务，清除标记。参数：stopped_services
restore_services() {
  local svc
  for svc in $1; do
    if svc_start "$svc"; then log "已重新启动 $svc（此前由 traffic-guard 停止）"
    else log "WARN: 启动 $svc 失败，请手动检查"; fi
  done
  [ -n "$GUARD_DRY_RUN" ] || rm -f "$MARKER_FILE"
}

# ---------- check ----------
do_check() {
  local month used stopped stopped_svcs stopped_at now svc newly="" fresh=0
  month="$(current_month)"
  snapshot_ifaces
  now="$(date -u +%Y-%m-%dT%H:%M:%SZ)"
  stopped="$(state_get stopped)"; stopped_svcs="$(state_get stopped_services)"
  stopped_at="$(state_get stopped_at)"; used="$(state_get used_bytes)"

  if [ ! -f "$STATE_FILE" ]; then
    log "首次运行：本月（$month）从 0 开始计量（如已知本月已用量，请用 set-used <GiB> 校准）"
    used=0; stopped=0; stopped_svcs=""; stopped_at=""; fresh=1
  elif [ "$(state_get month)" != "$month" ]; then
    log "进入新账期 $month（上月：$(state_get month)，累计 $(fmt_gib "${used:-0}") GiB）→ 计数清零"
    if [ "$stopped" = 1 ]; then restore_services "$stopped_svcs"; fi
    used=0; stopped=0; stopped_svcs=""; stopped_at=""; fresh=1
  fi

  if [ "$fresh" = 1 ]; then
    # 新基线：不计增量，只记下当前各网卡读数
    NEW_USED=0
    write_state "$month" 0 0 "" ""
  else
    compute_usage "${used:-0}"
  fi
  used="$NEW_USED"

  if [ "$used" -ge "$STOP_BYTES" ]; then
    # 只停当前 active 的服务；已经是停的/没装的不碰、不记
    for svc in $GUARD_SERVICES; do
      if svc_active "$svc"; then
        if svc_stop "$svc"; then
          newly+="$svc "
          case " $stopped_svcs " in *" $svc "*) ;; *) stopped_svcs="${stopped_svcs:+$stopped_svcs }$svc";; esac
        else
          log "WARN: 停止 $svc 失败，下次检查重试"
        fi
      fi
    done
    if [ "$stopped" != 1 ]; then
      if [ -n "$newly" ] || [ -n "$GUARD_DRY_RUN" ]; then
        stopped=1; stopped_at="$now"
        log "!!! 本月累计出站 $(fmt_gib "$used") GiB ≥ 阈值 $GUARD_STOP_AT_GIB GiB（额度 $GUARD_LIMIT_GIB GiB）→ 已停止服务：${newly% }；下个账期自动恢复"
        if [ -z "$GUARD_DRY_RUN" ]; then
          mkdir -p "$GUARD_STATE_DIR"
          printf 'stopped_at=%s\nmonth=%s\nused_bytes=%s\nservices=%s\n' "$now" "$month" "$used" "$stopped_svcs" > "$MARKER_FILE"
        fi
      else
        # 超阈值但没有任何 active 服务可停（例如全是手动停的）：记一次状态，不记 stopped_services
        stopped=1; stopped_at="$now"
        log "!!! 本月累计出站 $(fmt_gib "$used") GiB ≥ 阈值 $GUARD_STOP_AT_GIB GiB，但 GUARD_SERVICES 中没有处于 active 的服务可停"
        [ -n "$GUARD_DRY_RUN" ] || printf 'stopped_at=%s\nmonth=%s\nused_bytes=%s\nservices=\n' "$now" "$month" "$used" > "$MARKER_FILE"
      fi
    elif [ -n "$newly" ]; then
      log "已超阈值但服务又被拉起（可能重启过），再次停止：$newly"
    fi
  fi
  write_state "$month" "$used" "${stopped:-0}" "$stopped_svcs" "$stopped_at"
}

# ---------- status ----------
do_status() {
  local month smonth used stopped stopped_svcs stopped_at note=""
  month="$(current_month)"; smonth="$(state_get month)"
  snapshot_ifaces
  if [ ! -f "$STATE_FILE" ]; then
    log "尚无状态文件（$STATE_FILE）；首次 check 会建立本月基线"
    used=0; stopped=0; stopped_svcs=""; stopped_at=""
  elif [ "$smonth" != "$month" ]; then
    note="（状态文件还停留在 $smonth，下次 check 会清零并恢复服务）"
    used=0; stopped="$(state_get stopped)"; stopped_svcs="$(state_get stopped_services)"; stopped_at="$(state_get stopped_at)"
  else
    compute_usage "$(state_get used_bytes)"; used="$NEW_USED"
    stopped="$(state_get stopped)"; stopped_svcs="$(state_get stopped_services)"; stopped_at="$(state_get stopped_at)"
  fi
  local remaining=$((LIMIT_BYTES - used))
  [ "$remaining" -ge 0 ] || remaining=0
  echo "month:      $month  (账期时区 $GUARD_TZ) $note"
  echo "used:       $used bytes = $(fmt_gib "$used") GiB"
  echo "remaining:  $(fmt_gib "$remaining") GiB  (额度 $GUARD_LIMIT_GIB GiB)"
  echo "threshold:  $GUARD_STOP_AT_GIB GiB  (到达后停止：$GUARD_SERVICES)"
  if [ "${stopped:-0}" = 1 ]; then
    echo "stopped:    yes  (自 ${stopped_at:-?}；traffic-guard 停掉的服务：${stopped_svcs:-无})"
  else
    echo "stopped:    no"
  fi
}

# ---------- set-used ----------
do_set_used() {
  [[ "$ARG" =~ $num_re ]] || die "用法：set-used <GiB>（数字）"
  local month bytes
  month="$(current_month)"; bytes="$(gib_to_bytes "$ARG")"
  snapshot_ifaces
  # 沿用已有 stopped 信息；读数基线刷新为当前（之后只累加新增量）
  local stopped stopped_svcs stopped_at
  stopped="$(state_get stopped)"; stopped_svcs="$(state_get stopped_services)"; stopped_at="$(state_get stopped_at)"
  [ "$(state_get month)" = "$month" ] || { stopped=0; stopped_svcs=""; stopped_at=""; }
  write_state "$month" "$bytes" "${stopped:-0}" "$stopped_svcs" "$stopped_at"
  log "本月（$month）已用量已设为 $ARG GiB（$bytes bytes）"
}

# ---------- resume ----------
do_resume() {
  [ "$(state_get stopped)" = 1 ] || { log "当前没有被 traffic-guard 停止的服务"; return 0; }
  local month used
  month="$(state_get month)"; used="$(state_get used_bytes)"
  snapshot_ifaces
  restore_services "$(state_get stopped_services)"
  write_state "$month" "${used:-0}" 0 "" ""
}

# ---------- 主流程 ----------
if [ "$CMD" != status ] && [ -z "$GUARD_DRY_RUN" ]; then
  mkdir -p "$GUARD_STATE_DIR"
  exec 9> "$LOCK_FILE"
  flock -n 9 || { log "另一个实例正在运行，本次跳过"; exit 0; }
fi

case "$CMD" in
  check)    do_check;;
  status)   do_status;;
  set-used) do_set_used;;
  resume)   do_resume;;
esac