#!/usr/bin/env python3
"""Inject SSH_USER_<port> env into the GLOBAL receiver unit (gemastik-receiver).

Why: the panel's /api/credential proxy targets the global receiver on :18080,
not the per-team receivers. That unit had no Environment= lines at all, so
Challenge.credentials() fell back to the hardcoded 'ctfuser' literal and every
imported XVI/XVII challenge reported a login that could never work.

The registry's `ssh_user` per challenge is the single source of truth (the
panel already chpasswds that same account). Read the port each challenge runs
on from the global receiver's own config so the keys line up with self.port.
"""
import json
import re
import subprocess
import sys
from pathlib import Path

BASE = Path("/opt/gemastik18-final")
UNIT = Path("/etc/systemd/system/gemastik-receiver.service")
REGISTRY = BASE / "teams/challenge_registry.json"
RECV_CONFIG = BASE / "receiver/config.py"

reg = json.loads(REGISTRY.read_text())
ssh_users = {c["name"]: c.get("ssh_user", "ctfuser") for c in reg.get("challenges", [])}

# Challenge -> port used by the GLOBAL receiver. main.py builds the challenge
# objects from CHALLENGE_PORT_* / PASSWORD_* env; with none set it falls back to
# the pydantic settings PASSWORD_<port> in config.py, so mirror those ports.
text = RECV_CONFIG.read_text()
ports = {}
for m in re.finditer(r"PASSWORD_(\d+)\s*[:=]", text):
    ports.setdefault(m.group(1), int(m.group(1)))
# name -> port needs the CHALLENGE_PORT mapping; take it from registry order +
# the receiver main.py template, else fall back to PASSWORD_<port> keys.
name_to_port = {}
for m in re.finditer(r'"PASSWORD_(\d+)"', text):
    name_to_port.setdefault(m.group(1), m.group(1))
print(f"registry challenges: {len(ssh_users)}")

# Build Environment lines for every challenge whose port we can resolve.
env_lines = []
resolved = 0
for name, user in sorted(ssh_users.items()):
    # The global receiver uses the node-range ports from config.py; find the
    # port by matching the challenge name against the receiver's own mapping.
    port = None
    m = re.search(rf"{re.escape(name)}.*?(\d{{4,5}})", text)
    if m:
        port = m.group(1)
    if not port:
        continue
    env_lines.append(f'Environment="SSH_USER_{port}={user}"')
    resolved += 1

if not env_lines:
    print("ERROR: could not resolve any challenge port from config.py", file=sys.stderr)
    sys.exit(1)

body = UNIT.read_text()
# Drop any SSH_USER_ lines we previously injected (idempotent re-runs).
kept = [ln for ln in body.splitlines() if "SSH_USER_" not in ln]
# Insert right after the existing Environment=PYTHONUNBUFFERED line.
out = []
for ln in kept:
    out.append(ln)
    if ln.startswith("Environment=PYTHONUNBUFFERED"):
        out.extend(env_lines)
if not any(ln.startswith("Environment=PYTHONUNBUFFERED") for ln in out):
    out.extend(env_lines)
UNIT.write_text("\n".join(out) + "\n")
print(f"injected {resolved} SSH_USER_* lines into {UNIT}")

subprocess.run(["systemctl", "daemon-reload"], check=True)
subprocess.run(["systemctl", "restart", "gemastik-receiver"], check=True)
print("reloaded + restarted gemastik-receiver")