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

challenges = {
    "blogpost": Blogpost(10000),
    "carbeat": Carbeat(11000),
    "cdn": CDN(12000),
    "phew": Phew(13000),
    "sheesh": Sheesh(14000),
    "warmup": 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")