Merge pull request #5 from mradigen/master

Add `/api/ctf/start` and `/api/ctf/stop` routes

README.md

@@ -2,17 +2,18 @@
 
 ## Tech Stack:
 - Framework: **FastAPI**
-- Database: **PostgreSQL (ORM: SQLAlchemy)**
+- Database: **PostgreSQL (ORM: Tortoise)**
 - Server: **Uvicorn**
 - Test: **TestClient (in FastAPI) / Tox**
 - Containerization: **Docker**
 - CI/CD: **Github Actions** and ship to **Github packages**
 
 ## Setup:
 ```sh
+pip install poetry
 python -m venv .venv                # Create a python virtual environment
 source .venv/bin/activate           # Activate it (This command will differ for Windows)
-pip instal -r requirements.txt      # Install the dependencies
+poetry install                      # Install the dependencies
 ```
 
 ## Run:
@@ -39,13 +40,11 @@ All individual routes (`/team/*`, `/ctf/*`) are then put behind `/api` in the `r
 
 In case a certain route has multiple complex tasks, they can be separated as a submodule. For example, the route `/api/ctf/start` will perform a lot of tasks (interacting with docker etc.), and hence has a separate file for it.
 
+`src/`:
 ```
-app.py                      # Main file
 docs.py                     # Takes metadata from each route and compiles it for FastAPI
-config.py                   # Environment variables, could use .env instead
-db.py                       # Database schemas and connector (may do a separate directory if complexity exceeds)
-
-Dockerfile
+config.py                   # Configuration variables
+db.py                       # Database schemas and connector
 
 routes/
     L team.py
@@ -54,10 +53,8 @@ routes/
         L __init__.py       # Rest of the ctf routes go here
     L admin.py
     L leaderboard.py
+    L team.py
     L __init__.py           # Main router under `/api`, any misc routes go here
-helpers/
-    L container.py          # Specific helper functions
-    L __init__.py           # Contains general helper functions
 tests/
 ```
 

pyproject.toml

@@ -20,6 +20,7 @@ python = "^3.11"
 fastapi = "^0.104.1"
 uvicorn = "^0.24.0"
 tortoise-orm = {version = "^0.20.0", extras = ["asyncpg", "accel"]}
+aiodocker = "^0.21.0"
 
 [tool.poetry.group.dev.dependencies]
 mypy = "^1.6.1"

src/pwncore/__init__.py

@@ -1,10 +1,40 @@
 from fastapi import FastAPI
+from contextlib import asynccontextmanager
+
+from tortoise import Tortoise
 
 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
+
+
+@asynccontextmanager
+async def app_lifespan(app: FastAPI):
+    # Startup
+    await Tortoise.init(db_url=config.db_url, modules={"models": ["pwncore.models"]})
+    await Tortoise.generate_schemas()
+
+    yield
+    # Shutdown
+    # Stop and remove all running containers
+    containers = await Container.all().values()
+    await Container.all().delete()
+    for db_container in containers:
+        container = await docker_client.containers.get(db_container["docker_id"])
+        await container.stop()
+        await container.delete()
+
+    # close_connections is deprecated, not sure how to use connections.close_all()
+    await Tortoise.close_connections()
+    await docker_client.close()
+
 
 app = FastAPI(
-    title="Pwncore", openapi_tags=docs.tags_metadata, description=docs.description
+    title="Pwncore",
+    openapi_tags=docs.tags_metadata,
+    description=docs.description,
+    lifespan=app_lifespan,
 )
-
 app.include_router(routes.router)

src/pwncore/config.py

@@ -1,22 +1,57 @@
-from __future__ import annotations
-
-import typing as t
-
-__all__ = (
-    "Config",
-    "BaseConfig",
-)
-
-
-class Config(t.Protocol):
+from dataclasses import dataclass
+
+"""
+Sample messages:
+"db_error": "An error occurred, please try again.",
+"port_limit_reached": "Server ran out of ports 💀",
+"ctf_not_found": "CTF does not exist.",
+"container_start": "Container started.",
+"container_stop": "Container stopped.",
+"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."
+"""
+
+msg_codes = {
+    "db_error": 0,
+    "port_limit_reached": 1,
+    "ctf_not_found": 2,
+    "container_start": 3,
+    "container_stop": 4,
+    "containers_team_stop": 5,
+    "container_not_found": 6,
+    "container_already_running": 7,
+    "container_limit_reached": 8,
+}
+
+
+@dataclass
+class Config:
     development: bool
-
-
-class BaseConfig(Config):
-    __slots__ = ("development",)
-
-    def __init__(self, development: bool) -> None:
-        self.development = development
-
-
-DEV_CONFIG: t.Final[BaseConfig] = BaseConfig(True)
+    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={
+        "db_error": 0,
+        "port_limit_reached": 1,
+        "ctf_not_found": 2,
+        "container_start": 3,
+        "container_stop": 4,
+        "containers_team_stop": 5,
+        "container_not_found": 6,
+        "container_already_running": 7,
+        "container_limit_reached": 8,
+    },
+)

src/pwncore/container.py

@@ -0,0 +1,4 @@
+import aiodocker
+from pwncore.config import config
+
+docker_client = aiodocker.Docker(url=config.docker_url)

src/pwncore/routes/ctf/__init__.py

@@ -1,6 +1,5 @@
-from __future__ import annotations
-
 from fastapi import APIRouter
+from pwncore.routes.ctf.start import router as start_router
 
 # Metadata at the top for instant accessibility
 metadata = {
@@ -10,7 +9,7 @@
 }
 
 router = APIRouter(prefix="/ctf", tags=["ctf"])
-
+router.include_router(start_router)
 # Routes that do not need a separate submodule for themselves
 
 

src/pwncore/routes/ctf/start.py

@@ -1,20 +1,170 @@
 from __future__ import annotations
 
-from pwncore.routes.ctf import router
+from fastapi import APIRouter, Response
+import uuid
+from tortoise.transactions import atomic
 
+from pwncore.models import Problem, Container, Ports, Team
+from pwncore.container import docker_client
+from pwncore.config import config
 
-@router.get("/start/{ctf_id}")
-async def start_the_docker_container(
-    ctf_id: int,
-):  # The function name is inferred for the summary
-    # This is a regular single-line comment.
-    # Will not be displayed in the documentation.
-    """
-    This is a multi-line comment, and will be displayed
-    in the documentation when the route is expanded.
+# temporary helper functions
+if config.development:
+
+    def get_team_id():
+        return 1
+
+
+router = APIRouter(tags=["ctf"])
 
-    The cool thing is that Markdown works here!
-    # See, Markdown works!
-    _Pretty_ **cool** right?
+
+@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.
+    image_config:
+    {
+        "PortBindings": {
+            "22/tcp": [{}]      # Let docker randomly assign ports
+        }
+    }
     """
-    return {"status": "CTF started"}
+    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=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
+        return {
+            "msg_code": config.msg_codes["container_already_running"],
+            "ports": ports,
+            "ctf_id": ctf_id,
+        }
+
+    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
+    container_name = f"{team_id}_{ctf_id}_{uuid.uuid4().hex}"
+    container_flag = f"{config.flag}{{{uuid.uuid4().hex}}}"
+
+    # Run
+    container = await docker_client.containers.run(
+        name=container_name,
+        config={
+            "Image": ctf.image_name,
+            # Detach stuff
+            "AttachStdin": False,
+            "AttachStdout": False,
+            "AttachStderr": False,
+            "Tty": False,
+            "OpenStdin": False,
+            **ctf.image_config,
+        },
+    )
+
+    await (await container.exec(["/bin/bash", "/root/gen_flag", container_flag])).start(
+        detach=True
+    )
+
+    try:
+        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"]:
+            # 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)
+            await Ports.create(port=port, container=db_container)
+
+    except Exception:
+        # 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["container_start"],
+        "ports": ports,
+        "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()
+
+    # We first try to delete the record from the DB
+    # Then we stop the container
+    try:
+        await Container.filter(team_id=team_id).delete()
+    except Exception:
+        response.status_code = 500
+        return {"msg_code": config.msg_codes["db_error"]}
+
+    for db_container in containers:
+        container = await docker_client.containers.get(db_container["docker_id"])
+        await container.stop()
+        await container.delete()
+
+    return {"msg_code": config.msg_codes["containers_team_stop"]}
+
+
+@atomic()
+@router.post("/stop/{ctf_id}")
+async def stop_docker_container(ctf_id: int, response: Response):
+    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, 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, problem_id=ctf_id).delete()
+    except Exception:
+        response.status_code = 500
+        return {"msg_code": config.msg_codes["db_error"]}
+
+    container = await docker_client.containers.get(team_container.docker_id)
+    await container.stop()
+    await container.delete()
+
+    return {"msg_code": config.msg_codes["container_stop"]}