157 lines
4.8 KiB
Python
157 lines
4.8 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.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")
|