attack-defense-platform/panel/gen_receiver_services.py

#!/usr/bin/env python3
"""Generate + (re)start per-team receiver systemd services.

Each team receiver becomes gemastik-receiver-teamN.service, independent of
the panel service. This fixes the bug where restarting gemastik-panel killed
all team receivers (they were child processes of the panel systemd cgroup).
"""
import json
import os
import subprocess
import sys
from pathlib import Path

TEAMS_DIR = Path("/opt/gemastik18-final/teams")
RECEIVER_VENV = "/opt/gemastik18-final/receiver/.venv/bin/python"
UNIT_DIR = Path("/etc/systemd/system")
REGISTRY_PATH = TEAMS_DIR / "challenge_registry.json"


def load_registry() -> dict:
    try:
        return json.loads(REGISTRY_PATH.read_text())
    except Exception:
        return {"sets": {}, "challenges": []}

def write_unit(idx: int, st: dict):
    port = st["ports"]["receiver"]
    recv_dir = TEAMS_DIR / f"team{idx}" / "receiver"
    env = {}
    # ports/containers for challenge classes.
    # NOTE: systemd Environment= keys must be [A-Za-z0-9_]+ — a hyphen in the
    # challenge name (gift-card) would make systemd silently drop the line, so
    # normalize the name to underscores here. main.py looks up the same key.
    # Same normalization applies to CHALLENGE_SCHEME_<NAME>.
    schemes = {c["name"]: c.get("scheme") for c in load_registry().get("challenges", [])}
    # SSH login user per challenge. The panel already chpasswds the right
    # account from this field (teams.challenge_ssh_users); the receiver must
    # report the SAME user via /credential/<name> or participants get a login
    # that cannot work. Registry is the single source of truth for both sides.
    ssh_users = {c["name"]: c.get("ssh_user", "ctfuser")
                 for c in load_registry().get("challenges", [])}
    for ch in st["ports"]:
        if ch in ("receiver", "panel"):
            continue
        key = ch.upper().replace("-", "_")
        env[f"CHALLENGE_PORT_{key}"] = str(st["ports"][ch]["chall"])
        env[f"CHALLENGE_CONTAINER_{key}"] = f"{ch}_container_team{idx}"
        if schemes.get(ch):
            env[f"CHALLENGE_SCHEME_{key}"] = schemes[ch]
        if ssh_users.get(ch):
            env[f"SSH_USER_{st['ports'][ch]['chall']}"] = ssh_users[ch]
        # SSH passwords: checker Challenge.credentials() reads PASSWORD_<self.port>
        # where self.port is the team challenge port.
        pwd = st.get("chall_passwords", {}).get(ch)
        if pwd:
            env[f"PASSWORD_{st['ports'][ch]['chall']}"] = pwd
    env["COMPOSE_LOCATION"] = str(TEAMS_DIR / f"team{idx}" / "services" / "docker-compose.yml")
    env_lines = "\n".join(f'Environment="{k}={v}"' for k, v in env.items())
    unit = f"""[Unit]
Description=Gemastik A/D receiver for team {idx}
After=docker.service
Wants=docker.service

[Service]
Type=simple
WorkingDirectory={recv_dir}
EnvironmentFile={recv_dir}/.env
{env_lines}
ExecStart={RECEIVER_VENV} -m uvicorn main:app --host 0.0.0.0 --port {port}
Restart=always
RestartSec=3

[Install]
WantedBy=multi-user.target
"""
    (UNIT_DIR / f"gemastik-receiver-team{idx}.service").write_text(unit)

def main():
    action = sys.argv[1] if len(sys.argv) > 1 else "start"
    # Regenerate receiver main.py for existing teams from the registry
    try:
        from gen_receiver_main import sync_team_receivers
        sync_team_receivers()
    except Exception as e:
        print(f"WARN: receiver main.py regen failed: {e}")
    for d in sorted(TEAMS_DIR.glob("team*")):
        sf = d / "state.json"
        if not sf.exists():
            continue
        st = json.loads(sf.read_text())
        idx = st["index"]
        write_unit(idx, st)
        subprocess.run(["systemctl", "daemon-reload"], check=False)
        subprocess.run(["systemctl", "enable", f"gemastik-receiver-team{idx}.service"], check=False)
        if action == "stop":
            subprocess.run(["systemctl", "stop", f"gemastik-receiver-team{idx}.service"], check=False)
        else:
            subprocess.run(["systemctl", "restart", f"gemastik-receiver-team{idx}.service"], check=False)
        print(f"gemastik-receiver-team{idx}: {action} (port {st['ports']['receiver']})")

if __name__ == "__main__":
    main()