diff --git a/Dockerfile b/Dockerfile index 6a4873e..c604c7a 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,7 +1,5 @@ FROM python:3.12-slim -ENV DOCKER_HOST=unix:///var/run/docker.sock - WORKDIR /app RUN pip install poetry diff --git a/src/pwncore/__init__.py b/src/pwncore/__init__.py index c9b103e..7437540 100644 --- a/src/pwncore/__init__.py +++ b/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,6 +19,8 @@ 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 @@ -25,7 +28,9 @@ async def app_lifespan(app: FastAPI): 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 ( diff --git a/src/pwncore/container.py b/src/pwncore/container.py deleted file mode 100644 index 5f398c4..0000000 --- a/src/pwncore/container.py +++ /dev/null @@ -1,17 +0,0 @@ -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()) diff --git a/src/pwncore/containerASD.py b/src/pwncore/containerASD.py new file mode 100644 index 0000000..aea2550 --- /dev/null +++ b/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()) diff --git a/src/pwncore/routes/admin.py b/src/pwncore/routes/admin.py index 0e9d25a..4d35f3c 100644 --- a/src/pwncore/routes/admin.py +++ b/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() diff --git a/src/pwncore/routes/ctf/__init__.py b/src/pwncore/routes/ctf/__init__.py index 84c08c6..40ce348 100644 --- a/src/pwncore/routes/ctf/__init__.py +++ b/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() # diff --git a/src/pwncore/routes/ctf/start.py b/src/pwncore/routes/ctf/start.py index d22c260..7a70dea 100644 --- a/src/pwncore/routes/ctf/start.py +++ b/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,11 +37,8 @@ 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: @@ -49,7 +46,9 @@ async def start_docker_container(ctf_id: int, response: Response, jwt: RequireJw 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()