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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
SOURCE="$ROOT_DIR/scripts/refresh-deployment-contract.sh"
PUBLIC_KEY_FILE="${DEPLOY_REFRESH_PUBLIC_KEY_FILE:-}"
MARKER="# hyb-oa-managed-deployment-contract-refresh"
MODE="${1:-install}"
INSTALL_TEST_MODE=0
ROOT_PREFIX=""

fail() {
  echo "[deployment-contract-refresh-installer][ERROR] $1" >&2
  exit 1
}

if [[ "$(id -u)" != "0" ]]; then
  test_root="${HYB_CONTRACT_REFRESH_INSTALL_TEST_ROOT:-}"
  if [[ "$test_root" =~ ^/tmp/[A-Za-z0-9._/-]+$ && "$test_root" != */ \
    && "$test_root" != *"//"* && "/$test_root/" != *"/../"* ]]; then
    INSTALL_TEST_MODE=1
    ROOT_PREFIX="$test_root"
  else
    fail "root privileges are required"
  fi
fi
TARGET="${ROOT_PREFIX}/usr/local/sbin/hyb-oa-deployment-contract-refresh"
AUTHORIZED_KEYS="${ROOT_PREFIX}/root/.ssh/authorized_keys"
SSH_DIRECTORY="${ROOT_PREFIX}/root/.ssh"
[[ -f "$SOURCE" && ! -L "$SOURCE" ]] || fail "reviewed refresh source is missing"

remove_managed_entry() {
  python3 - "$AUTHORIZED_KEYS" "$MARKER" "$TARGET" "$INSTALL_TEST_MODE" <<'PY'
import os, sys, tempfile
path, marker, target, test_mode = sys.argv[1:]
if not os.path.exists(path):
    raise SystemExit(0)
with open(path, encoding="utf-8") as handle:
    lines = handle.read().splitlines()
result = []
index = 0
forced_prefix = f'command="{target}",no-agent-forwarding,no-port-forwarding,no-X11-forwarding,no-pty,no-user-rc '
while index < len(lines):
    line = lines[index]
    if line == marker:
        if index + 1 >= len(lines) or not lines[index + 1].startswith(forced_prefix):
            raise SystemExit("managed forced-command marker is malformed")
        index += 2
        continue
    result.append(line)
    index += 1
directory = os.path.dirname(path)
fd, temp_path = tempfile.mkstemp(prefix=".authorized_keys.", dir=directory, text=True)
try:
    with os.fdopen(fd, "w", encoding="utf-8") as handle:
        if result:
            handle.write("\n".join(result) + "\n")
        handle.flush()
        os.fsync(handle.fileno())
    os.chmod(temp_path, 0o600)
    if test_mode != "1":
        os.chown(temp_path, 0, 0)
    os.replace(temp_path, path)
finally:
    if os.path.exists(temp_path):
        os.unlink(temp_path)
PY
}

install_entry() {
  [[ -n "$PUBLIC_KEY_FILE" && -f "$PUBLIC_KEY_FILE" && ! -L "$PUBLIC_KEY_FILE" ]] \
    || fail "DEPLOY_REFRESH_PUBLIC_KEY_FILE must name a regular public-key file"
  key_line="$(sed -e '/^[[:space:]]*#/d' -e '/^[[:space:]]*$/d' "$PUBLIC_KEY_FILE")"
  [[ "$(printf '%s\n' "$key_line" | wc -l)" -eq 1 ]] \
    || fail "public-key file must contain exactly one key"
  read -r key_type key_blob key_comment <<<"$key_line"
  [[ "$key_type" =~ ^(ssh-ed25519|ecdsa-sha2-nistp(256|384|521)|ssh-rsa)$ \
    && "$key_blob" =~ ^[A-Za-z0-9+/]+={0,3}$ ]] \
    || fail "public key format is invalid"

  if [[ "$INSTALL_TEST_MODE" == "0" ]]; then
    install -d -o root -g root -m 0700 "$SSH_DIRECTORY"
  else
    install -d -m 0700 "$SSH_DIRECTORY"
  fi
  [[ ! -L "$SSH_DIRECTORY" && ! -L "$AUTHORIZED_KEYS" ]] \
    || fail "root SSH paths must not be symbolic links"
  touch "$AUTHORIZED_KEYS"
  if [[ "$INSTALL_TEST_MODE" == "0" ]]; then
    chown root:root "$AUTHORIZED_KEYS"
  fi
  chmod 0600 "$AUTHORIZED_KEYS"

  PUBLIC_KEY_BLOB="$key_blob" python3 - "$AUTHORIZED_KEYS" "$MARKER" "$TARGET" "$key_type" "$key_line" "$INSTALL_TEST_MODE" <<'PY'
import os, re, sys, tempfile
path, marker, target, key_type, key_line, test_mode = sys.argv[1:]
key_blob = os.environ["PUBLIC_KEY_BLOB"]
with open(path, encoding="utf-8") as handle:
    lines = handle.read().splitlines()
unmanaged = []
kept = []
pattern = re.compile(r"(?:^|\s)(?:ssh-ed25519|ecdsa-sha2-nistp(?:256|384|521)|ssh-rsa)\s+([A-Za-z0-9+/=]+)(?:\s|$)")
forced_prefix = f'command="{target}",no-agent-forwarding,no-port-forwarding,no-X11-forwarding,no-pty,no-user-rc '
index = 0
while index < len(lines):
    line = lines[index]
    if line == marker:
        if index + 1 >= len(lines) or not lines[index + 1].startswith(forced_prefix):
            raise SystemExit("managed forced-command marker is malformed")
        index += 2
        continue
    match = pattern.search(line)
    if match and match.group(1) == key_blob:
        unmanaged.append(line)
    kept.append(line)
    index += 1
if unmanaged:
    raise SystemExit("dedicated key already has unmanaged or unrestricted root authorization")
forced = f'command="{target}",no-agent-forwarding,no-port-forwarding,no-X11-forwarding,no-pty,no-user-rc {key_line}'
kept.extend([marker, forced])
directory = os.path.dirname(path)
fd, temp_path = tempfile.mkstemp(prefix=".authorized_keys.", dir=directory, text=True)
try:
    with os.fdopen(fd, "w", encoding="utf-8") as handle:
        handle.write("\n".join(kept) + "\n")
        handle.flush()
        os.fsync(handle.fileno())
    os.chmod(temp_path, 0o600)
    if test_mode != "1":
        os.chown(temp_path, 0, 0)
    os.replace(temp_path, path)
finally:
    if os.path.exists(temp_path):
        os.unlink(temp_path)
PY
}

case "$MODE" in
  install)
    if [[ "$INSTALL_TEST_MODE" == "0" ]]; then
      install -d -o root -g root -m 0755 "$(dirname "$TARGET")"
      install -o root -g root -m 0755 "$SOURCE" "$TARGET"
      [[ -f "$TARGET" && ! -L "$TARGET" \
        && "$(stat -c %u:%g:%a "$TARGET")" == "0:0:755" ]] \
        || fail "installed target ownership or mode is invalid"
      probe_env=()
    else
      install -d -m 0755 "$(dirname "$TARGET")"
      install -m 0755 "$SOURCE" "$TARGET"
      [[ -f "$TARGET" && ! -L "$TARGET" && "$(stat -c %a "$TARGET")" == "755" ]] \
        || fail "installed target mode is invalid"
      probe_env=(env HYB_CONTRACT_REFRESH_TEST_ROOT="$ROOT_PREFIX")
    fi
    "${probe_env[@]}" env SSH_ORIGINAL_COMMAND=probe "$TARGET" >/dev/null \
      || fail "installed target probe failed"
    install_entry
    echo "[deployment-contract-refresh-installer][PASS] forced command installed"
    ;;
  uninstall)
    remove_managed_entry
    echo "[deployment-contract-refresh-installer][PASS] authorization removed; target retained"
    ;;
  *)
    fail "unsupported mode"
    ;;
esac