attack-defense-platform/receiver/main.py

from fastapi import Depends, FastAPI, HTTPException
from pydantic import BaseModel
from fastapi.security import HTTPBasic, HTTPBasicCredentials
from config import get_settings

from challenges.Blogpost import Blogpost
from challenges.Carbeat import Carbeat
from challenges.CDN import CDN
from challenges.Phew import Phew
from challenges.Sheesh import Sheesh
from challenges.Warmup import Warmup

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 _ch_port(name: str, default: int) -> int:
    # read from .env manually (pydantic settings has fixed fields)
    val = os.environ.get(f"CHALLENGE_PORT_{name.upper()}")
    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_{name.upper()}="):
                        val = line.strip().split("=", 1)[1]
        except Exception:
            pass
    return int(val) if val else default

def _ch_container(name: str, default: str) -> str:
    val = os.environ.get(f"CHALLENGE_CONTAINER_{name.upper()}")
    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_{name.upper()}="):
                        val = line.strip().split("=", 1)[1]
        except Exception:
            pass
    return val or default

challenges = {
    "blogpost": Blogpost(_ch_port("blogpost", 10000)),
    "carbeat": Carbeat(_ch_port("carbeat", 11000)),
    "cdn": CDN(_ch_port("cdn", 12000)),
    "phew": Phew(_ch_port("phew", 13000)),
    "sheesh": Sheesh(_ch_port("sheesh", 14000)),
    "warmup": Warmup(_ch_port("warmup", 15000)),
}

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")