chore: Cleanup and format

lugvitc · Nov 17, 2023 · 8cd63c6 · 8cd63c6
1 parent 2125999
commit 8cd63c6
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 32 deletions.
diff --git a/src/pwncore/__init__.py b/src/pwncore/__init__.py
@@ -13,10 +13,7 @@
 @asynccontextmanager
 async def app_lifespan(app: FastAPI):
     # Startup
-    await Tortoise.init(
-        db_url=config.db_url,
-        modules={"models": ["pwncore.models"]}
-    )
+    await Tortoise.init(db_url=config.db_url, modules={"models": ["pwncore.models"]})
     await Tortoise.generate_schemas()
 
     yield
@@ -29,7 +26,8 @@ async def app_lifespan(app: FastAPI):
         await container.stop()
         await container.delete()
 
-    await Tortoise.close_connections()  # Deprecated, not sure how to use connections.close_all()
+    # close_connections is deprecated, not sure how to use connections.close_all()
+    await Tortoise.close_connections()
     await docker_client.close()
 
 

diff --git a/src/pwncore/config.py b/src/pwncore/config.py
@@ -10,7 +10,8 @@
 "containers_team_stop": "All team containers stopped.",
 "container_not_found": "You have no running containers for this CTF.",
 "container_already_running": "Your team already has a running container for this CTF.",
-"container_limit_reached": "Your team already has reached the maximum number of containers limit, please stop other unused containers."
+"container_limit_reached": "Your team already has reached the maximum number"
+                           " of containers limit, please stop other unused containers."
 """
 
 msg_codes = {
@@ -31,13 +32,15 @@ class Config:
     development: bool
     msg_codes: dict
     db_url: str
+    docker_url: str | None
     flag: str
     max_containers_per_team: int
 
 
 config = Config(
     development=True,
     db_url="sqlite://:memory:",
+    docker_url=None,  # None for default system docker
     flag="C0D",
     max_containers_per_team=3,
     msg_codes={

diff --git a/src/pwncore/container.py b/src/pwncore/container.py
@@ -1,3 +1,4 @@
 import aiodocker
+from pwncore.config import config
 
-docker_client = aiodocker.Docker()
+docker_client = aiodocker.Docker(url=config.docker_url)
diff --git a/src/pwncore/routes/ctf/start.py b/src/pwncore/routes/ctf/start.py
@@ -4,14 +4,15 @@
 import uuid
 from tortoise.transactions import atomic
 
-# 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
 if config.development:
-    def get_team_id(): return 1
+
+    def get_team_id():
+        return 1
 
 
 router = APIRouter(tags=["ctf"])
@@ -32,20 +33,15 @@ async def start_docker_container(ctf_id: int, response: Response):
     """
     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": [{}]}},
-            }
+            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",
-            }
+            name="CID Squad" + uuid.uuid4().hex, secret_hash="veryverysecret"
         )
 
     ctf = await Problem.get_or_none(id=ctf_id)
@@ -61,13 +57,10 @@ async def start_docker_container(ctf_id: int, response: Response):
         return {
             "msg_code": config.msg_codes["container_already_running"],
             "ports": ports,
-            "ctf_id": team_container.problem_id,
+            "ctf_id": ctf_id,
         }
 
-    if (
-        await Container.filter(team_id=team_id).count()
-        >= config.max_containers_per_team
-    ):
+    if await Container.filter(team_id=team_id).count() >= config.max_containers_per_team:  # noqa: B950
         return {"msg_code": config.msg_codes["container_limit_reached"]}
 
     # Start a new container
@@ -79,6 +72,7 @@ async def start_docker_container(ctf_id: int, response: Response):
         name=container_name,
         config={
             "Image": ctf.image_name,
+            # Detach stuff
             "AttachStdin": False,
             "AttachStdout": False,
             "AttachStderr": False,
@@ -94,12 +88,10 @@ async def start_docker_container(ctf_id: int, response: Response):
 
     try:
         db_container = await Container.create(
-            **{
-                "docker_id": container.id,
-                "team_id": team_id,
-                "problem_id": ctf_id,
-                "flag": container_flag,
-            }
+            docker_id=container.id,
+            team_id=team_id,
+            problem_id=ctf_id,
+            flag=container_flag,
         )
 
         # Get ports and save them