#!/usr/bin/env bash
# Generic client-side service interception for OpenFUT.
#
#   openfut-switch.sh on  --server-ip <IP> --intercept-port <P> --target-port <Q> \
#                         --name <ID> [--client-ip <IP>]
#   openfut-switch.sh off --name <ID>
#   openfut-switch.sh status [--name <ID>]
#
# ONE implementation. `blaze-switch.sh` is a thin compatibility wrapper over
# this; there is deliberately no second copy of the iptables logic, because two
# scripts manipulating the same table diverge and then disagree about what is
# installed.
#
# WHAT IT DOES
#
# Redirects traffic destined for <server-ip>:<intercept-port> to a local
# <target-port>, so a replacement or observer can sit in front of a service
# without touching that service. Loopback is never matched: the rule is scoped
# to the server address, so 127.0.0.1:<intercept-port> keeps reaching the
# original process and stays usable as an oracle.
#
# TWO INDEPENDENT VIEWS, ON PURPOSE
#
# Rules are CREATED and DELETED by their comment tag. They are VERIFIED by
# parsing the kernel's own rule fields — chain, destination, dport, to-ports —
# with no reference to the comment. An earlier version of the Blaze switch
# matched `--comment "tag"` with quotes this iptables does not emit, so removal
# found nothing AND the verification used the same broken matcher, confirming a
# rollback that had not happened. A verifier must not share the failure mode of
# the thing it verifies.
set -uo pipefail

TAG_PREFIX="openfut-switch"
SUDO=""
[[ $EUID -eq 0 ]] || SUDO=sudo

die() { echo "openfut-switch: $*" >&2; exit 2; }

# ------------------------------------------------------------------ parsing
#
# Structured view of the nat table. Parses FIELDS, not comment text — this is
# the independent verifier referred to above.
rules_json() {
  $SUDO iptables-save -t nat 2>/dev/null | python3 -c '
import json, re, sys
out = []
for line in sys.stdin:
    line = line.strip()
    if not line.startswith("-A "):
        continue
    parts = line.split()
    def val(flag):
        try:
            return parts[parts.index(flag) + 1]
        except (ValueError, IndexError):
            return None
    # Comment may be quoted or bare depending on iptables version; correctness
    # never depends on which, because every check below can use the fields.
    m = re.search(r"--comment\s+(\"([^\"]*)\"|(\S+))", line)
    comment = (m.group(2) or m.group(3)) if m else None
    out.append({
        "chain": parts[1],
        "dest": (val("-d") or "").split("/")[0],
        "src": (val("-s") or "").split("/")[0],
        "dport": val("--dport"),
        "target": val("-j"),
        "to_ports": val("--to-ports"),
        "comment": comment,
        "spec": line,
    })
print(json.dumps(out))
'
}

# Rules bearing a given switch name.
rules_named() {
  rules_json | python3 -c '
import json, sys
name = sys.argv[1]
print(json.dumps([r for r in json.load(sys.stdin) if r["comment"] == name]))
' "$1"
}

# Any REDIRECT touching a port, whoever owns it. Used to spot conflicts and
# stale rules that lost or never had our tag.
redirects_on_port() {
  rules_json | python3 -c '
import json, sys
port = sys.argv[1]
print(json.dumps([r for r in json.load(sys.stdin)
                  if r["target"] == "REDIRECT" and r["dport"] == port]))
' "$1"
}

# ------------------------------------------------------------------- args
CMD="${1:-}"; shift || true
NAME=""; SERVER=""; CLIENT=""; IPORT=""; TPORT=""; LEGACY=""
while [[ $# -gt 0 ]]; do
  case "$1" in
    --name)           NAME="${2:-}"; shift 2 ;;
    --server-ip)      SERVER="${2:-}"; shift 2 ;;
    --client-ip)      CLIENT="${2:-}"; shift 2 ;;
    --intercept-port) IPORT="${2:-}"; shift 2 ;;
    --target-port)    TPORT="${2:-}"; shift 2 ;;
    # A tag this switch previously used. Rules carrying it are OURS and must
    # still be removable, otherwise renaming a switch orphans live NAT rules
    # that no tool can clean up while `off` cheerfully reports success.
    --legacy-tag)     LEGACY="${2:-}"; shift 2 ;;
    *) die "unknown argument: $1" ;;
  esac
done

tag_for() { echo "${TAG_PREFIX}:$1"; }

is_port() { [[ "$1" =~ ^[0-9]+$ ]] && (( $1 > 0 && $1 < 65536 )); }

# ------------------------------------------------------------------ status
cmd_status() {
  local tag all
  if [[ -n "$NAME" ]]; then
    tag="$(tag_for "$NAME")"
    all="$(rules_json | python3 -c '
import json,sys
tags=[t for t in sys.argv[1:] if t]
print(json.dumps([r for r in json.load(sys.stdin) if r["comment"] in tags]))
' "$tag" "$LEGACY")"
  else
    tag=""
    all="$(rules_json | python3 -c '
import json,sys
print(json.dumps([r for r in json.load(sys.stdin)
                  if (r["comment"] or "").startswith("'"$TAG_PREFIX"':")]))')"
  fi

  python3 - "$all" "$tag" <<'PY'
import json, sys
rules = json.loads(sys.argv[1])
tag = sys.argv[2]

if not rules:
    print("INACTIVE: no switch rules%s" % (f" named {tag}" if tag else ""))
    raise SystemExit(0)

# Group by the SEMANTIC identity of the redirect, not by comment text.
groups = {}
for r in rules:
    key = (r["dest"], r["dport"], r["to_ports"], r["comment"])
    groups.setdefault(key, []).append(r)

print("ACTIVE:")
problems = []
for (dest, dport, to, comment), rs in sorted(groups.items()):
    chains = ",".join(sorted(r["chain"] for r in rs))
    print(f"  {comment}: {dest}:{dport} -> :{to}   [{chains}]")
    # A healthy switch installs exactly one PREROUTING and one OUTPUT rule.
    per_chain = {}
    for r in rs:
        per_chain[r["chain"]] = per_chain.get(r["chain"], 0) + 1
    for chain, n in per_chain.items():
        if n > 1:
            problems.append(f"DUPLICATE: {n} identical rules in {chain} for {comment}")
    for want in ("PREROUTING", "OUTPUT"):
        if want not in per_chain:
            problems.append(f"INCOMPLETE: {comment} has no {want} rule")

# Several different targets for one name is inconsistent state.
by_name = {}
for (dest, dport, to, comment), _ in groups.items():
    by_name.setdefault(comment, set()).add((dest, dport, to))
for comment, variants in by_name.items():
    if len(variants) > 1:
        problems.append(f"CONFLICT: {comment} has {len(variants)} different redirects: {sorted(variants)}")

if problems:
    print()
    for p in problems:
        print(f"  !! {p}")
    raise SystemExit(1)
PY
  local rc=$?

  # Foreign or untagged redirects on the same port are reported, never touched.
  if [[ -n "$IPORT" ]]; then
    local foreign
    foreign="$(redirects_on_port "$IPORT" | python3 -c '
import json,sys
tag=sys.argv[1]
tags=set(sys.argv[1:])
o=[r for r in json.load(sys.stdin) if r["comment"] not in tags]
print("\n".join("  ?? untagged/foreign: %s" % r["spec"] for r in o))' "$(tag_for "$NAME")" "$LEGACY")"
    [[ -n "$foreign" ]] && { echo "$foreign" >&2; rc=1; }
  fi
  return $rc
}

# ---------------------------------------------------------------------- on
cmd_on() {
  [[ -n "$NAME"   ]] || die "--name is required"
  [[ -n "$SERVER" ]] || die "--server-ip is required (the backend the client dials)"
  is_port "${IPORT:-}" || die "--intercept-port must be a port"
  is_port "${TPORT:-}" || die "--target-port must be a port"
  [[ "$IPORT" != "$TPORT" ]] || die "--intercept-port and --target-port must differ"

  # REFUSE to arm at a port nothing is listening on. Arming a switch whose
  # target is dead silently breaks the client path: the redirect happens, the
  # connection is refused, and the proven Python service is bypassed for no
  # benefit. This has happened three times, always the same way — a build guard
  # correctly refuses to start the replacement, and the `on` that follows in the
  # same script arms anyway because it never checked.
  #
  # Overridable for the rare case of arming ahead of a service that is about to
  # start, but it must be deliberate rather than the default.
  if [[ "${ALLOW_DEAD_TARGET:-0}" != "1" ]]; then
    if ! ss -ltn 2>/dev/null | grep -qE "[:.]${TPORT}[[:space:]]"; then
      die "REFUSING: nothing is listening on target port $TPORT.
  Arming would break the client path — the redirect would land on a closed
  socket and the working service would be bypassed.
  Start the replacement first, or set ALLOW_DEAD_TARGET=1 if that is intended."
    fi
  fi

  local tag; tag="$(tag_for "$NAME")"

  # Never stack: start from a known state for THIS name only.
  cmd_off_quiet "$tag"
  [[ -n "$LEGACY" ]] && cmd_off_quiet "$LEGACY"

  local scope=()
  [[ -n "$CLIENT" ]] && scope=(-s "$CLIENT")

  $SUDO iptables -t nat -I PREROUTING 1 -p tcp "${scope[@]}" -d "$SERVER" --dport "$IPORT" \
      -m comment --comment "$tag" -j REDIRECT --to-ports "$TPORT" \
    || die "failed to add PREROUTING rule"
  # OUTPUT covers this host's own connections so the switch can be smoke tested
  # locally. Loopback is still unmatched: it is scoped to the server address.
  $SUDO iptables -t nat -I OUTPUT 1 -p tcp -d "$SERVER" --dport "$IPORT" \
      -m comment --comment "$tag" -j REDIRECT --to-ports "$TPORT" \
    || die "failed to add OUTPUT rule"

  # Verify from the kernel's fields, not from what we think we just ran.
  local ok
  ok="$(rules_named "$tag" | python3 -c '
import json,sys
rs=json.load(sys.stdin); dest,dport,to=sys.argv[1:4]
good=[r for r in rs if r["dest"]==dest and r["dport"]==dport and r["to_ports"]==to
      and r["target"]=="REDIRECT"]
chains={r["chain"] for r in good}
print("yes" if len(good)==2 and chains=={"PREROUTING","OUTPUT"} else "no:%d:%s"%(len(good),sorted(chains)))
' "$SERVER" "$IPORT" "$TPORT")"
  [[ "$ok" == "yes" ]] || die "rule verification failed after install ($ok)"

  echo "$NAME ON: $SERVER:$IPORT -> local :$TPORT"
  echo "  127.0.0.1:$IPORT still reaches the original service (scoped to $SERVER)"
  echo "  roll back with: $0 off --name $NAME"
}

# --------------------------------------------------------------------- off
#
# Removes ONLY rules bearing this switch's exact tag. Anything else that
# redirects the same port is reported, never deleted — precise removal, not
# broad deletion.
cmd_off_quiet() {
  local tag="$1"
  # Delete by RECONSTRUCTED FIELDS, never by re-feeding the raw `iptables-save`
  # line through the shell.
  #
  # `iptables-save` prints `--comment "tag"` WITH quotes on this version. Word-
  # splitting that back into an argv leaves the quote characters inside the
  # comment value, so iptables looks for a rule whose comment literally contains
  # `"` and finds nothing — a silent no-op delete, and the same
  # comment-formatting trap that produced the original lying rollback. Rules
  # with a bare comment deleted fine, which is exactly what made it look like it
  # worked.
  #
  # Fields are passed as argv elements, so no quoting survives to be
  # misinterpreted.
  # Unit Separator, NOT tab. Tab is an IFS *whitespace* character, so bash
  # collapses runs of it and drops empties — an absent `-s` therefore shifted
  # every later field left, producing `-s <dport> --dport <to_ports>
  # --to-ports ''`. Those deletes failed harmlessly here, but a shifted spec
  # that happened to match a real rule would delete the wrong one.
  local chain dest src dport to
  while IFS=$'\x1f' read -r chain dest src dport to; do
    [[ -n "$chain" ]] || continue
    local args=(-t nat -D "$chain" -p tcp)
    [[ -n "$src"  ]] && args+=(-s "$src")
    [[ -n "$dest" ]] && args+=(-d "$dest")
    args+=(--dport "$dport" -m comment --comment "$tag" -j REDIRECT --to-ports "$to")
    $SUDO iptables "${args[@]}" 2>/dev/null
  done < <(rules_named "$tag" | python3 -c '
import json,sys
for r in json.load(sys.stdin):
    print("\x1f".join([r["chain"], r["dest"] or "", r["src"] or "",
                       r["dport"] or "", r["to_ports"] or ""]))')
}

cmd_off() {
  [[ -n "$NAME" ]] || die "--name is required"
  local tag; tag="$(tag_for "$NAME")"

  local before legacy_before=0
  before="$(rules_named "$tag" | python3 -c 'import json,sys; print(len(json.load(sys.stdin)))')"
  if [[ -n "$LEGACY" ]]; then
    legacy_before="$(rules_named "$LEGACY" | python3 -c 'import json,sys; print(len(json.load(sys.stdin)))')"
  fi

  cmd_off_quiet "$tag"
  [[ -n "$LEGACY" ]] && cmd_off_quiet "$LEGACY"
  before=$(( before + legacy_before ))

  # Independent verification: re-read the table and check the FIELDS.
  local after
  after="$(rules_named "$tag" | python3 -c 'import json,sys; print(len(json.load(sys.stdin)))')"
  if [[ -n "$LEGACY" ]]; then
    after=$(( after + $(rules_named "$LEGACY" | python3 -c 'import json,sys; print(len(json.load(sys.stdin)))') ))
  fi
  if [[ "$after" != "0" ]]; then
    echo "FAILED: $after rule(s) named $tag still present after removing $before" >&2
    rules_named "$tag" | python3 -c 'import json,sys
for r in json.load(sys.stdin): print("  "+r["spec"])' >&2
    return 1
  fi

  # A redirect on that port owned by someone else is a conflict to report, not
  # something this switch may delete.
  if [[ -n "$IPORT" ]]; then
    local others
    others="$(redirects_on_port "$IPORT" | python3 -c 'import json,sys
rs=json.load(sys.stdin)
print("\n".join("  "+r["spec"] for r in rs))')"
    if [[ -n "$others" ]]; then
      echo "WARNING: other REDIRECT rule(s) still target port $IPORT (not ours, not removed):" >&2
      echo "$others" >&2
      return 1
    fi
  fi

  echo "$NAME OFF: removed $before rule(s), verified none remain"
}

case "$CMD" in
  on)     cmd_on ;;
  off)    cmd_off ;;
  status) cmd_status ;;
  *) sed -n '2,12p' "$0" | sed 's/^# \?//'; exit 2 ;;
esac