Merge pull request #46 from KreativeThinker/master

fix: ASD ~Aadivishnu

Dockerfile

@@ -1,7 +1,5 @@
 FROM python:3.12-slim
 
-ENV DOCKER_HOST=unix:///var/run/docker.sock
-
 WORKDIR /app
 
 RUN pip install poetry

src/pwncore/__init__.py

@@ -1,13 +1,14 @@
 from contextlib import asynccontextmanager
 
+import aiodocker
+from aiodocker import docker
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from tortoise import Tortoise
 
+import pwncore.containerASD as containerASD
 import pwncore.docs as docs
 import pwncore.routes as routes
-
-from pwncore.container import docker_client
 from pwncore.config import config
 from pwncore.models import Container
 
@@ -18,14 +19,18 @@ async def app_lifespan(app: FastAPI):
     await Tortoise.init(db_url=config.db_url, modules={"models": ["pwncore.models"]})
     await Tortoise.generate_schemas()
 
+    containerASD.docker_client = aiodocker.Docker(url=config.docker_url)
+
     yield
     # Shutdown
     # Stop and remove all running containers
     containers = await Container.all().values()
     await Container.all().delete()
     for db_container in containers:
         try:
-            container = await docker_client.containers.get(db_container["docker_id"])
+            container = await containerASD.docker_client.containers.get(
+                db_container["docker_id"]
+            )
             await container.kill()
             await container.delete()
         except (

src/pwncore/containerASD.py

@@ -0,0 +1,17 @@
+import asyncio
+
+import aiodocker
+
+from pwncore.config import config
+
+docker_client: aiodocker.Docker = None  # type: ignore[assignment]
+
+
+# if not hasattr(sys, "_is_a_test"):
+#     docker_client = aiodocker.Docker(url=config.docker_url)
+#
+# async def create_docker_client():
+#    return aiodocker.Docker(url=config.docker_url)
+
+
+# docker_client = asyncio.run(create_docker_client())

src/pwncore/routes/admin.py

@@ -1,22 +1,22 @@
-from datetime import date
 import logging
+from datetime import date
 
 from fastapi import APIRouter, Request, Response
 from passlib.hash import bcrypt
 from tortoise.transactions import atomic, in_transaction
 
+import pwncore.containerASD as containerASD
+from pwncore.config import config
 from pwncore.models import (
-    Team,
-    Problem,
     Hint,
-    User,
-    PreEventSolvedProblem,
     PreEventProblem,
+    PreEventSolvedProblem,
+    Problem,
+    Team,
+    User,
 )
-from pwncore.config import config
 from pwncore.models.ctf import SolvedProblem
 from pwncore.models.pre_event import PreEventUser
-from pwncore.container import docker_client
 
 metadata = {
     "name": "admin",
@@ -47,7 +47,7 @@
 
 
 async def _del_cont(id: str):
-    container = await docker_client.containers.get(id)
+    container = await containerASD.docker_client.containers.get(id)
     await container.kill()
     await container.delete()
 

src/pwncore/routes/ctf/__init__.py

@@ -1,28 +1,28 @@
 from __future__ import annotations
 
 from asyncio import create_task
-from logging import getLogger
 from collections import defaultdict
+from logging import getLogger
 
 from fastapi import APIRouter, Request, Response
 from pydantic import BaseModel
 from tortoise.transactions import atomic
 
+import pwncore.containerASD as containerASD
+from pwncore.config import config
 from pwncore.models import (
-    Problem,
-    SolvedProblem,
     Container,
     Hint,
-    ViewedHint,
     Hint_Pydantic,
+    Problem,
+    SolvedProblem,
     Team,
+    ViewedHint,
 )
-from pwncore.config import config
 from pwncore.models.ctf import Problem_Pydantic
-from pwncore.routes.ctf.start import router as start_router
-from pwncore.routes.ctf.pre_event import router as pre_event_router
 from pwncore.routes.auth import RequireJwt
-from pwncore.container import docker_client
+from pwncore.routes.ctf.pre_event import router as pre_event_router
+from pwncore.routes.ctf.start import router as start_router
 
 # Metadata at the top for instant accessibility
 metadata = {
@@ -124,7 +124,9 @@ async def flag_post(
             response.status_code = 500
             return {"msg_code": config.msg_codes["db_error"]}
 
-        container = await docker_client.containers.get(team_container.docker_id)
+        container = await containerASD.docker_client.containers.get(
+            team_container.docker_id
+        )
         await container.kill()
         await container.delete()
         #

src/pwncore/routes/ctf/start.py

@@ -1,14 +1,14 @@
 from __future__ import annotations
 
+import uuid
 from logging import getLogger
 
 from fastapi import APIRouter, Response
-import uuid
 from tortoise.transactions import in_transaction
 
-from pwncore.models import Problem, Container, Ports
-from pwncore.container import docker_client
+import pwncore.containerASD as containerASD
 from pwncore.config import config
+from pwncore.models import Container, Ports, Problem
 from pwncore.routes.auth import RequireJwt
 
 router = APIRouter(tags=["ctf"])
@@ -37,19 +37,18 @@ async def start_docker_container(ctf_id: int, response: Response, jwt: RequireJw
         team_container = await Container.filter(team=team_id, problem=ctf_id)
         if team_container:
             a, b = team_container[0], team_container[1:]
-            db_ports = await a.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 a.ports.all().values("port")  # Get ports from DB
+            ports = [db_port["port"] for db_port in db_ports]  # Create a list out of it
 
             for db_container in b:
                 try:
                     await db_container.delete()
                 except Exception:
                     pass
 
-                container = await docker_client.containers.get(db_container.docker_id)
+                container = await containerASD.docker_client.containers.get(
+                    db_container.docker_id
+                )
                 await container.kill()
                 await container.delete()
 
@@ -71,7 +70,7 @@ async def start_docker_container(ctf_id: int, response: Response, jwt: RequireJw
         container_flag = f"{config.flag}{{{uuid.uuid4().hex}}}"
 
         # Run
-        container = await docker_client.containers.run(
+        container = await containerASD.docker_client.containers.run(
             name=container_name,
             config={
                 "Image": ctf.image_name,
@@ -81,7 +80,8 @@ async def start_docker_container(ctf_id: int, response: Response, jwt: RequireJw
                 "AttachStderr": False,
                 "Tty": False,
                 "OpenStdin": False,
-                **ctf.image_config,
+                "HostConfig": {"PublishAllPorts": True},
+                # **ctf.image_config,
             },
         )
 
@@ -139,7 +139,9 @@ async def stopall_docker_container(response: Response, jwt: RequireJwt):
             return {"msg_code": config.msg_codes["db_error"]}
 
         for db_container in containers:
-            container = await docker_client.containers.get(db_container["docker_id"])
+            container = await containerASD.docker_client.containers.get(
+                db_container["docker_id"]
+            )
             await container.kill()
             await container.delete()
 
@@ -169,7 +171,9 @@ async def stop_docker_container(ctf_id: int, response: Response, jwt: RequireJwt
             response.status_code = 500
             return {"msg_code": config.msg_codes["db_error"]}
 
-        container = await docker_client.containers.get(team_container.docker_id)
+        container = await containerASD.docker_client.containers.get(
+            team_container.docker_id
+        )
         await container.kill()
         await container.delete()