From 0df880514c62ddc91f45f39f1de888758a67a0b2 Mon Sep 17 00:00:00 2001 From: mradigen <55953083+mradigen@users.noreply.github.com> Date: Thu, 16 Nov 2023 15:38:46 +0530 Subject: [PATCH] chore: Add app_lifespan and make start.py adhere with models from #7 --- src/pwncore/__init__.py | 32 +++++--- src/pwncore/config.py | 6 +- src/pwncore/container.py | 8 -- src/pwncore/db.py | 29 ------- src/pwncore/routes/ctf/__init__.py | 6 +- src/pwncore/routes/ctf/start.py | 118 ++++++++++++++--------------- 6 files changed, 86 insertions(+), 113 deletions(-) delete mode 100644 src/pwncore/db.py diff --git a/src/pwncore/__init__.py b/src/pwncore/__init__.py index 867c3dd..48e316c 100644 --- a/src/pwncore/__init__.py +++ b/src/pwncore/__init__.py @@ -1,21 +1,33 @@ from fastapi import FastAPI +from contextlib import asynccontextmanager from tortoise.contrib.fastapi import register_tortoise import pwncore.docs as docs import pwncore.routes as routes +from pwncore.container import docker_client from pwncore.config import config -app = FastAPI( - title="Pwncore", openapi_tags=docs.tags_metadata, description=docs.description -) -app.include_router(routes.router) +@asynccontextmanager +async def app_lifespan(app: FastAPI): + # Startup + register_tortoise( + app, + db_url=config.db_url, + modules={"models": ["pwncore.models"]}, + generate_schemas=True, + add_exception_handlers=True, + ) + yield + # Shutdown + await docker_client.close() -register_tortoise( - app, - db_url=config.db_url, - modules={"models": ["pwncore.db"]}, - generate_schemas=True, - add_exception_handlers=True + +app = FastAPI( + title="Pwncore", + openapi_tags=docs.tags_metadata, + description=docs.description, + lifespan=app_lifespan, ) +app.include_router(routes.router) diff --git a/src/pwncore/config.py b/src/pwncore/config.py index 34cdb6c..ac66190 100644 --- a/src/pwncore/config.py +++ b/src/pwncore/config.py @@ -22,7 +22,7 @@ "containers_team_stop": 5, "container_not_found": 6, "container_already_running": 7, - "container_limit_reached": 8 + "container_limit_reached": 8, } @@ -49,6 +49,6 @@ class Config: "containers_team_stop": 5, "container_not_found": 6, "container_already_running": 7, - "container_limit_reached": 8 - } + "container_limit_reached": 8, + }, ) diff --git a/src/pwncore/container.py b/src/pwncore/container.py index 15b5c2a..5cf5045 100644 --- a/src/pwncore/container.py +++ b/src/pwncore/container.py @@ -1,11 +1,3 @@ import aiodocker -import atexit -import asyncio docker_client = aiodocker.Docker() - - -async def docker_cleanup(): - await docker_client.close() - -atexit.register(lambda: asyncio.run(docker_cleanup())) diff --git a/src/pwncore/db.py b/src/pwncore/db.py deleted file mode 100644 index 742cc31..0000000 --- a/src/pwncore/db.py +++ /dev/null @@ -1,29 +0,0 @@ -from __future__ import annotations - -from tortoise import fields -from tortoise.fields import JSONField -from tortoise.models import Model - -# Tortoise is initiated from __main__.py since initiating -# it as a submodule creates multiple event loops. - -# No relationships yet -class Container(Model): - id = fields.TextField(pk=True) - name = fields.TextField() - ctf_id = fields.IntField() - ports: fields.ReverseRelation[Ports] - team_id = fields.IntField() - flag = fields.TextField() - -class CTF(Model): - name = fields.TextField() - # docker_config = fields.JSONField() - image_name = fields.TextField() - image_config = fields.JSONField() - -class Ports(Model): - container: fields.ForeignKeyRelation[Container] = fields.ForeignKeyField( - "models.Container", related_name="ports", on_delete=fields.OnDelete.CASCADE - ) - port = fields.IntField(pk=True) diff --git a/src/pwncore/routes/ctf/__init__.py b/src/pwncore/routes/ctf/__init__.py index 4883f9e..db0e506 100644 --- a/src/pwncore/routes/ctf/__init__.py +++ b/src/pwncore/routes/ctf/__init__.py @@ -1,7 +1,5 @@ -from __future__ import annotations - from fastapi import APIRouter -from pwncore.routes.ctf import start +from pwncore.routes.ctf.start import router as start_router # Metadata at the top for instant accessibility metadata = { @@ -11,7 +9,7 @@ } router = APIRouter(prefix="/ctf", tags=["ctf"]) -router.include_router(start.router) +router.include_router(start_router) # Routes that do not need a separate submodule for themselves diff --git a/src/pwncore/routes/ctf/start.py b/src/pwncore/routes/ctf/start.py index 8602ab8..eaa012e 100644 --- a/src/pwncore/routes/ctf/start.py +++ b/src/pwncore/routes/ctf/start.py @@ -3,17 +3,17 @@ from fastapi import APIRouter, Response import uuid from tortoise.transactions import atomic -import logging -from pwncore.db import Container, CTF, Ports +# from pwncore.db import Problem, Container, Ports +from pwncore.models import Problem, Container, Ports, Team from pwncore.container import docker_client from pwncore.config import config # temporary helper functions -def get_team_id(): - return 34 -def get_empty_ports(): - return [4444] +if config.development: + + def get_team_id(): + return 1 router = APIRouter(tags=["ctf"]) @@ -22,7 +22,6 @@ def get_empty_ports(): @atomic() @router.post("/start/{ctf_id}") async def start_docker_container(ctf_id: int, response: Response): - """ image_config contains the raw POST data that gets sent to the Docker Remote API. For now it just contains the guest ports that need to be opened on the host. @@ -33,37 +32,45 @@ async def start_docker_container(ctf_id: int, response: Response): } } """ - # Testing purposes - # await CTF.create(**{ - # "name": "Invisible-Incursion", - # "image_name": "key:latest", - # "image_config": { - # "PortBindings": { - # "22/tcp": [{}] - # } - # } - # }) - - ctf = await CTF.get_or_none(id=ctf_id) + if config.development: + await Problem.create( + **{ + "name": "Invisible-Incursion", + "description": "Chod de tujhe se na ho paye", + "author": "Meetesh Saini", + "points": 300, + "image_name": "key:latest", + "image_config": {"PortBindings": {"22/tcp": [{}]}}, + } + ) + await Team.create( + **{ + "name": "CID Squad" + uuid.uuid4().hex, + "secret_hash": "veryverysecret", + } + ) + + ctf = await Problem.get_or_none(id=ctf_id) if not ctf: response.status_code = 404 return {"msg_code": config.msg_codes["ctf_not_found"]} team_id = get_team_id() # From JWT - team_container = await Container.get_or_none(team_id=team_id, ctf_id=ctf_id) + team_container = await Container.get_or_none(team=team_id, problem=ctf_id) if team_container: - db_ports = await team_container.ports.all().values('port') # Get ports from DB - ports = [ db_port['port'] for db_port in db_ports ] # Create a list out of it + db_ports = await team_container.ports.all().values("port") # Get ports from DB + ports = [db_port["port"] for db_port in db_ports] # Create a list out of it return { "msg_code": config.msg_codes["container_already_running"], "ports": ports, - "ctf_id": team_container.ctf_id + "ctf_id": team_container.problem_id, } - if await Container.filter(team_id=team_id).count() >= config.max_containers_per_team: - return { - "msg_code": config.msg_codes["container_limit_reached"] - } + if ( + await Container.filter(team_id=team_id).count() + >= config.max_containers_per_team + ): + return {"msg_code": config.msg_codes["container_limit_reached"]} # Start a new container container_name = f"{team_id}_{ctf_id}_{uuid.uuid4().hex}" @@ -79,54 +86,52 @@ async def start_docker_container(ctf_id: int, response: Response): "AttachStderr": False, "Tty": False, "OpenStdin": False, - **ctf.image_config - } + **ctf.image_config, + }, ) - await ( - await container.exec(["/bin/bash", "/root/gen_flag", container_flag]) - ).start(detach=True) + await (await container.exec(["/bin/bash", "/root/gen_flag", container_flag])).start( + detach=True + ) try: - db_container = await Container.create(**{ - "id" : container.id, - "name" : container_name, - "team_id" : team_id, - "ctf_id" : ctf_id, - "flag" : container_flag - }) + db_container = await Container.create( + **{ + "docker_id": container.id, + "team_id": team_id, + "problem_id": ctf_id, + "flag": container_flag, + } + ) # Get ports and save them ports = [] # List to return back to frontend - for guest_port in ctf.image_config['PortBindings']: + for guest_port in ctf.image_config["PortBindings"]: # Docker assigns the port to the IPv4 and IPv6 addresses # Since we only require IPv4, we select the zeroth item # from the returned list. port = int((await container.port(guest_port))[0]["HostPort"]) ports.append(port) - host_port = await Ports.create(port=port, container=db_container) + await Ports.create(port=port, container=db_container) - except Exception: + except Exception as e: # Stop the container if failed to make a DB record await container.stop() await container.delete() response.status_code = 500 - return { - "msg_code": config.msg_codes["db_error"] - } + return {"msg_code": config.msg_codes["db_error"]} return { "msg_code": config.msg_codes["container_start"], "ports": ports, - "ctf_id": ctf_id + "ctf_id": ctf_id, } @atomic() @router.post("/stopall") async def stopall_docker_container(response: Response): - team_id = get_team_id() # From JWT containers = await Container.filter(team_id=team_id).values() @@ -137,12 +142,10 @@ async def stopall_docker_container(response: Response): await Container.filter(team_id=team_id).delete() except Exception: response.status_code = 500 - return { - "msg_code": config.msg_codes["db_error"] - } + return {"msg_code": config.msg_codes["db_error"]} for db_container in containers: - container = await docker_client.containers.get(db_container.id) + container = await docker_client.containers.get(db_container["docker_id"]) await container.stop() await container.delete() @@ -152,28 +155,25 @@ async def stopall_docker_container(response: Response): @atomic() @router.post("/stop/{ctf_id}") async def stop_docker_container(ctf_id: int, response: Response): - - ctf = await CTF.get_or_none(id=ctf_id) + ctf = await Problem.get_or_none(id=ctf_id) if not ctf: response.status_code = 404 return {"msg_code": config.msg_codes["ctf_not_found"]} team_id = get_team_id() - team_container = await Container.get_or_none(team_id=team_id, ctf_id=ctf_id) + team_container = await Container.get_or_none(team_id=team_id, problem_id=ctf_id) if not team_container: return {"msg_code": config.msg_codes["container_not_found"]} # We first try to delete the record from the DB # Then we stop the container try: - await Container.filter(team_id=team_id, ctf_id=ctf_id).delete() + await Container.filter(team_id=team_id, problem_id=ctf_id).delete() except Exception: response.status_code = 500 - return { - "msg_code": config.msg_codes["db_error"] - } + return {"msg_code": config.msg_codes["db_error"]} - container = await docker_client.containers.get(team_container.id) + container = await docker_client.containers.get(team_container.docker_id) await container.stop() await container.delete()