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

NODE_ID=BS-SG-001
EXPECTED_HOSTNAME=guanghu-lang-sg-001
JD_SOURCE_CIDR=111.228.0.139/32
KEY_ID=jd_ops_to_bs_sg_001
PUBLIC_KEY='ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIKJdOqsqD6UFrt3jFRYks07fKBHHsfD+NdMRQlpl+w3+ jd-ops-to-bs-sg-001-20260718'
AUTHORIZED_KEYS=/root/.ssh/authorized_keys
RESTRICTED_LINE="from=\"${JD_SOURCE_CIDR}\",no-agent-forwarding,no-port-forwarding,no-X11-forwarding,no-pty ${PUBLIC_KEY}"
BACKUP_ROOT=/var/backups/guanghu/node-enrollment/${NODE_ID}
RECEIPT_DIR=/var/lib/guanghu/node-enrollment

if [[ ${EUID} -ne 0 ]]; then
  echo "Run this fixed enrollment as root." >&2
  exit 1
fi

if [[ ${1:-} == --rollback ]]; then
  if [[ ! -f $AUTHORIZED_KEYS ]]; then
    echo "authorized_keys does not exist; nothing to roll back"
    exit 0
  fi
  temp=$(mktemp "${AUTHORIZED_KEYS}.rollback.XXXXXX")
  trap 'rm -f "$temp"' EXIT
  grep -vF -- "$RESTRICTED_LINE" "$AUTHORIZED_KEYS" > "$temp" || true
  install -m 0600 -o root -g root "$temp" "$AUTHORIZED_KEYS"
  echo "Removed only ${KEY_ID}; other authorized keys were preserved."
  exit 0
fi

actual_hostname=$(hostname)
if [[ $actual_hostname != "$EXPECTED_HOSTNAME" ]]; then
  echo "Node identity mismatch: expected $EXPECTED_HOSTNAME, got $actual_hostname" >&2
  exit 2
fi

status=$(curl -fsS --max-time 5 http://127.0.0.1:3911/status)
server_id=$(python3 -c 'import json,sys; print(json.load(sys.stdin).get("server_id", ""))' <<<"$status")
if [[ $server_id != "$NODE_ID" ]]; then
  echo "Gatekeeper identity mismatch: expected $NODE_ID, got ${server_id:-missing}" >&2
  exit 2
fi

install -d -m 0700 /root/.ssh
touch "$AUTHORIZED_KEYS"
chmod 0600 "$AUTHORIZED_KEYS"

if grep -qxF -- "$RESTRICTED_LINE" "$AUTHORIZED_KEYS"; then
  changed=false
else
  timestamp=$(date -u +%Y%m%dT%H%M%SZ)
  backup_dir="$BACKUP_ROOT/$timestamp"
  install -d -m 0700 "$backup_dir"
  cp -a "$AUTHORIZED_KEYS" "$backup_dir/authorized_keys"
  printf '%s\n' "$RESTRICTED_LINE" >> "$AUTHORIZED_KEYS"
  changed=true
fi

install -d -m 0750 "$RECEIPT_DIR"
python3 - "$changed" > "$RECEIPT_DIR/latest.json" <<'PY'
import json, sys
from datetime import datetime, timezone
print(json.dumps({
    "schema": "guanghu.node-enrollment-receipt/v1",
    "node_id": "BS-SG-001",
    "controller": "JD-OPS-CENTER",
    "key_id": "jd_ops_to_bs_sg_001",
    "source_restricted": True,
    "port_forwarding": False,
    "agent_forwarding": False,
    "pty": False,
    "changed": sys.argv[1] == "true",
    "completed_at": datetime.now(timezone.utc).isoformat(),
}, ensure_ascii=False, indent=2))
PY
chmod 0640 "$RECEIPT_DIR/latest.json"

echo "BS-SG-001 enrollment complete. JD source restriction and non-forwarding guards are active."
echo "Rollback: bash $0 --rollback"