Files
attack-defense-platform/receiver/main.py
T
2025-10-25 14:48:20 +07:00

108 lines
3.3 KiB
Python

from fastapi import Depends, FastAPI, HTTPException
from pydantic import BaseModel
from fastapi.security import HTTPBasic, HTTPBasicCredentials
from config import get_settings
from challenges.Blogspot import Blogspot
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
app = FastAPI()
security = HTTPBasic()
settings = get_settings()
challenges = {
"blogpost": Blogspot(10000),
"carbeat": Carbeat(11000),
"cdn": CDN(12000),
"phew": Phew(13000),
"sheesh": Sheesh(14000),
"warmup": Warmup(15000),
}
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")