attack-defense-platform/panel/gen_receiver_main.py

#!/usr/bin/env python3
"""Build the per-team receiver main.py from the challenge registry.

The receiver's `challenges` dict is generated from ENABLED registry entries so
the SLA checker pool exactly matches the distributed challenges.  Imports come
from three packages:
  - challenges.<Name>        (gemastik18 native checkers)
  - challenges.xvi.<Name>    (GEMASTIK XVI checkers)
  - challenges.xvii.checkers (GEMASTIK XVII generic checkers)
"""
from __future__ import annotations

import json
import shutil
from pathlib import Path

from teams import TEAMS_DIR, load_registry


def checker_class_for(ch: dict) -> str:
    """Return Python import path + class name for a registry challenge."""
    name = ch["name"]
    s = ch["set"]
    # gemastik18 native challenges have their own checker class (capitalized)
    if s == "gemastik18":
        cls = {
            "blogpost": "Blogpost",
            "carbeat": "Carbeat",
            "cdn": "CDN",
            "phew": "Phew",
            "sheesh": "Sheesh",
            "warmup": "Warmup",
        }.get(name)
        if cls:
            return f"from challenges.{cls} import {cls}", cls
        raise KeyError(f"no native checker for {name}")
    if s == "xvi":
        cls = {
            "art": "Art",
            "xl": "XL",
            "gemas-notes": "GemasNotes",
            "pasta": "Pasta",
            "burvesigner": "Burvesigner",
            "hirnfick": "Hirnfick",
            "gemas-fetcher": "GemasFetcher",
            "s3": "S3",
            "crawlback": "Crawlback",
            "back-to-basic": "BackToBasic",
        }.get(name)
        if cls:
            return f"from challenges.xvi.{cls} import {cls}", cls
        raise KeyError(f"no xvi checker for {name}")
    if s == "xvii":
        cls = {
            "anti-alchemy": "AntiAlchemy",
            "asmr": "Asmr",
            "bit-canvas": "BitCanvas",
            "fjb": "Fjb",
            "gift-card": "GiftCard",
            "gift-voucher": "GiftVoucher",
            "gleam-drive": "GleamDrive",
            "go-green": "GoGreen",
            "kode-viewer": "KodeViewer",
            "more-less": "MoreLess",
            "tempest-poc": "TempestPoc",
            "ticketer": "Ticketer",
        }.get(name)
        if cls:
            return f"from challenges.xvii.checkers import {cls}", cls
        raise KeyError(f"no xvii checker for {name}")
    raise KeyError(f"unknown set {s}")


def render_receiver_main(enabled: list[dict], port_defaults: dict) -> str:
    imports = []
    entries = []
    for ch in enabled:
        name = ch["name"]
        imp, cls = checker_class_for(ch)
        imports.append(imp)
        default = port_defaults.get(name, 10000)
        entries.append(f'    "{name}": {cls}(_ch_port("{name}", {default})),')
    return f"""from fastapi import Depends, FastAPI, HTTPException
from pydantic import BaseModel
from fastapi.security import HTTPBasic, HTTPBasicCredentials
from config import get_settings

{chr(10).join(imports)}

import os
import asyncio
import logging

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()
security = HTTPBasic()
settings = get_settings()

def _envkey(name: str) -> str:
    # systemd Environment= keys can't contain hyphens; gen_receiver_services
    # writes CHALLENGE_PORT_GIFT_CARD for the challenge "gift-card".
    return name.upper().replace("-", "_")

def _ch_port(name: str, default: int) -> int:
    # read from .env manually (pydantic settings has fixed fields)
    key = _envkey(name)
    val = os.environ.get(f"CHALLENGE_PORT_{{key}}")
    if not val:
        try:
            with open(os.path.join(os.path.dirname(__file__), ".env")) as f:
                for line in f:
                    if line.startswith(f"CHALLENGE_PORT_{{key}}="):
                        val = line.strip().split("=", 1)[1]
        except Exception:
            pass
    return int(val) if val else default

def _ch_container(name: str, default: str) -> str:
    key = _envkey(name)
    val = os.environ.get(f"CHALLENGE_CONTAINER_{{key}}")
    if not val:
        try:
            with open(os.path.join(os.path.dirname(__file__), ".env")) as f:
                for line in f:
                    if line.startswith(f"CHALLENGE_CONTAINER_{{key}}="):
                        val = line.strip().split("=", 1)[1]
        except Exception:
            pass
    return val or default

challenges = {{
{chr(10).join(entries)}
}}

async def run_challenge_checks():
    \"\"\"Run check function on all challenges at startup\"\"\"
    logger.info("\\n" + "="*60)
    logger.info("Running challenge checks...")
    logger.info("="*60 + "\\n")
    
    results = {{}}
    
    for name, challenge in challenges.items():
        logger.info(f"\\n[{{name}}] Starting check...")
        try:
            # Give service time between checks
            await asyncio.sleep(2)
            
            result = challenge.check()
            results[name] = result
            
            if result:
                logger.info(f"[{{name}}] ✓ Check PASSED")
            else:
                logger.warning(f"[{{name}}] ✗ Check FAILED")
        except Exception as e:
            logger.error(f"[{{name}}] ✗ Check ERROR: {{e}}")
            results[name] = False
    
    # Print summary
    logger.info("\\n" + "="*60)
    logger.info("Challenge Check Summary:")
    logger.info("="*60)
    passed = sum(1 for r in results.values() if r)
    total = len(results)
    for name, result in results.items():
        status = "✓ PASS" if result else "✗ FAIL"
        logger.info(f"  {{name:20}} {{status}}")
    logger.info(f"\\nTotal: {{passed}}/{{total}} passed")
    logger.info("="*60 + "\\n")
    
    return results

@app.on_event("startup")
async def startup_event():
    \"\"\"Run challenge checks on application startup\"\"\"
    asyncio.create_task(run_challenge_checks())

class Flag(BaseModel):
    flag: str
    challenge: str

class History(BaseModel):
    log: str

@app.get("/")
def read_root():
    return {{"service": "receiver-service"}}


@app.get("/restart/{{challenge}}")
def restart(challenge: str, credentials: HTTPBasicCredentials = Depends(security)):
    validate(credentials, challenge)
    os.system(f"docker compose -f {{settings.COMPOSE_LOCATION}} restart {{challenge}}")
    return {{"message": "Challenge restarted"}}


@app.get("/rollback/{{challenge}}")
def rollback(challenge: str, credentials: HTTPBasicCredentials = Depends(security)):
    validate(credentials, challenge)
    os.system(f"docker compose -f {{settings.COMPOSE_LOCATION}} up -d --force-recreate {{challenge}}")
    return {{"message": "Challenge restarted"}}


@app.get("/activate/{{challenge}}")
def activate(challenge: str, credentials: HTTPBasicCredentials = Depends(security)):
    validate(credentials, challenge)
    os.system(f"docker compose -f {{settings.COMPOSE_LOCATION}} up -d {{challenge}}")
    return {{"message": "Challenge activated"}}


@app.get("/deactivate/{{challenge}}")
def deactive(challenge: str, credentials: HTTPBasicCredentials = Depends(security)):
    validate(credentials, challenge)
    os.system(f"docker compose -f {{settings.COMPOSE_LOCATION}} down {{challenge}}")
    return {{"message": "Challenge deactivated"}}


@app.get("/credential/{{challenge}}")
def credential(challenge: str, credentials: HTTPBasicCredentials = Depends(security)):
    validate(credentials, challenge)
    return challenges[challenge].credentials()


@app.post("/flag")
def receive(data: Flag, credentials: HTTPBasicCredentials = Depends(security)):
    validate(credentials, data.challenge)
    challenge = challenges[data.challenge]
    if challenge.distribute(data.flag):
        return {{"message": "Flag received"}}

    raise HTTPException(status_code=500, detail="Error receiving flag")


@app.get("/check/{{challenge}}")
def check(challenge: str, credentials: HTTPBasicCredentials = Depends(security)):
    validate(credentials, challenge)
    return {{"success": challenges[challenge].check()}}

@app.post("/history")
def history(data: History):
    with open('history/command.txt', 'a') as f:
        f.write(data.log + '\\n')
    return {{"message": "Command received"}}

def is_admin(credentials):
    if credentials.username != settings.ADMIN_USERNAME or credentials.password != settings.ADMIN_PASSWORD:
        return False
    return True


def validate(credentials, challenge):
    if not is_admin(credentials):
        raise HTTPException(status_code=401, detail="Invalid credentials")

    if challenge not in challenges:
        raise HTTPException(status_code=400, detail="Invalid challenge")
"""


def _sync_checker_packages(recv_dir) -> None:
    """Mirror the shared receiver's challenge checkers into a team receiver.

    Two reasons this must run on every sync, not just at create_team time:
      1. Teams created before the XVI/XVII import have no xvi/xvii subpackages,
         so a regenerated main.py that imports them would crash-loop the receiver.
      2. Native checkers (Blogpost/Phew/...) get bug fixes (e.g. mandatory
         subprocess timeouts); a team copy made earlier keeps the stale version.
    """
    shared_challenges = Path("/opt/gemastik18-final/receiver/challenges")
    dst_root = recv_dir / "challenges"
    dst_root.mkdir(parents=True, exist_ok=True)
    for src_file in sorted(shared_challenges.glob("*.py")):
        if src_file.name == "__init__.py":
            continue
        shutil.copy2(src_file, dst_root / src_file.name)
    for pkg in ("xvi", "xvii"):
        src = shared_challenges / pkg
        if not src.exists():
            continue
        dst = dst_root / pkg
        if dst.exists():
            shutil.rmtree(dst)
        shutil.copytree(src, dst, ignore=shutil.ignore_patterns("__pycache__", "*.pyc"))


def sync_team_receivers() -> None:
    """Regenerate main.py for every existing team from its state + registry."""
    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"]
        enabled = [c for c in load_registry()["challenges"] if c.get("enabled")]
        _sync_checker_packages(d / "receiver")
        text = render_receiver_main(enabled, {c["name"]: st["ports"][c["name"]]["chall"] for c in enabled})
        (d / "receiver" / "main.py").write_text(text)
        print(f"team{idx}: receiver main.py regenerated ({len(enabled)} challenges)")


if __name__ == "__main__":
    sync_team_receivers()