chore: Add app_lifespan and make start.py adhere with models from lug…

…vitc#7

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)

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,
+ },
 )

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

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
 
 

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