Merge pull request #80 from libkush/master

feat: added staticfs integration

Author: KreativeThinker

src/pwncore/__init__.py

@@ -1,6 +1,8 @@
 from contextlib import asynccontextmanager
 
+import shutil
 import aiodocker
+from logging import getLogger
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from tortoise import Tortoise
@@ -11,6 +13,7 @@
 from pwncore.config import config
 from pwncore.models import Container
 
+logger = getLogger(__name__)
 
 @asynccontextmanager
 async def app_lifespan(app: FastAPI):
@@ -36,6 +39,10 @@ async def app_lifespan(app: FastAPI):
             Exception
         ):  # Raises DockerError if container does not exist, just pass for now.
             pass
+    try:
+        shutil.rmtree(config.staticfs_data_dir)
+    except Exception as err:
+        logger.exception("Failed to delete static files", exc_info=err)
 
     # close_connections is deprecated, not sure how to use connections.close_all()
     await Tortoise.close_connections()

src/pwncore/config.py

@@ -60,6 +60,9 @@ class Config:
     jwt_valid_duration: int
     hint_penalty: int
     max_members_per_team: int
+    staticfs_url: str
+    staticfs_data_dir: str
+    staticfs_jwt_secret: str
     admin_hash: str  
 
 config = Config(
@@ -76,6 +79,9 @@ class Config:
     msg_codes=msg_codes,
     hint_penalty=50,
     max_members_per_team=3,
+    staticfs_url="http://localhost:8080",
+    staticfs_data_dir=os.environ.get("STATIC_DATA_DIR", "/data"),
+    staticfs_jwt_secret="PyMioVKFXHymQd+n7q5geOsT6fSYh3gDVw3GqilW+5U="
     admin_hash=admin_hash_value,
 )
 

src/pwncore/models/container.py

@@ -20,7 +20,8 @@ class Container(Model):
     )
     team: fields.ForeignKeyRelation[Team] = fields.ForeignKeyField("models.Team")
     flag = fields.TextField()
-
+    
+    token = fields.TextField()
     ports: fields.ReverseRelation[Ports]
 
 

src/pwncore/models/ctf.py

@@ -35,6 +35,7 @@ class Problem(BaseProblem):
     # image_config: fields.Field[dict[str, Any]] = fields.JSONField(
     #     null=True
     # )  # type: ignore[assignment]
+    static = fields.BooleanField(default=False)
 
     mi = fields.IntField(default=50)
     ma = fields.IntField(default=500)
@@ -44,7 +45,7 @@ class Problem(BaseProblem):
     hints: fields.ReverseRelation[Hint]
 
     class PydanticMeta:
-        exclude = ["image_name", "mi", "ma", "visible"]
+        exclude = ["image_name", "static", "mi", "ma", "visible"]
 
     async def _solves(self) -> int:
         return await SolvedProblem.filter(problem=self).count()

src/pwncore/routes/ctf/__init__.py

@@ -3,6 +3,7 @@
 from asyncio import create_task
 from collections import defaultdict
 from logging import getLogger
+import shutil
 
 from fastapi import APIRouter, Request, Response
 from pydantic import BaseModel
@@ -124,16 +125,21 @@ async def flag_post(
             response.status_code = 500
             return {"msg_code": config.msg_codes["db_error"]}
 
-        container = await containerASD.docker_client.containers.get(
-            team_container.docker_id
-        )
-        await container.kill()
-        await container.delete()
-        #
+        if problem.static:
+            shutil.rmtree(
+                f"{config.staticfs_data_dir}/{team_id}/{team_container.docker_id}"
+            )
+        else:
+            container = await containerASD.docker_client.containers.get(
+                team_container.docker_id
+            )
+            await container.kill()
+            await container.delete()
 
         await SolvedProblem.create(team_id=team_id, problem_id=ctf_id, penalty=pnlt)
         create_task(update_points(req, ctf_id))
         return {"status": True}
+
     return {"status": False}
 
 

src/pwncore/routes/ctf/start.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
+import os
 import uuid
+import shutil
 from logging import getLogger
-
+import jwt as jwtlib
 from fastapi import APIRouter, Response
 from tortoise.transactions import in_transaction
 
@@ -43,24 +45,29 @@ async def start_docker_container(ctf_id: int, response: Response, jwt: RequireJw
         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
-
+            ports = [db_port["port"] for db_port in db_ports]  # Create a list out of it
+            static_url = f"{config.staticfs_url}/{a.token}" if ctf.static else None
             for db_container in b:
                 try:
                     await db_container.delete()
                 except Exception:
                     pass
-
-                container = await containerASD.docker_client.containers.get(
-                    db_container.docker_id
-                )
-                await container.kill()
-                await container.delete()
+                # containers won't exist for static ctfs
+                if ctf.static:
+                    staticLocation = f"{config.staticfs_data_dir}/{team_id}/{db_container.docker_id}"
+                    if os.path.exists(staticLocation):
+                        shutil.rmtree(staticLocation) 
+                else:
+                    container = await containerASD.docker_client.containers.get(
+                        db_container.docker_id
+                    )
+                    await container.kill()
+                    await container.delete()
 
             return {
                 "msg_code": config.msg_codes["container_already_running"],
                 "ports": ports,
+                "static_url": static_url,
                 "ctf_id": ctf_id,
             }
 
@@ -75,6 +82,53 @@ async def start_docker_container(ctf_id: int, response: Response, jwt: RequireJw
         container_name = f"{team_id}_{ctf_id}_{uuid.uuid4().hex}"
         container_flag = f"{config.flag}{{{uuid.uuid4().hex}}}"
 
+        if ctf.static:
+            container_id = uuid.uuid4().hex
+            payload = {
+                "id": str(team_id),
+                "containerId": container_id,
+            }
+            token = jwtlib.encode(payload, config.staticfs_jwt_secret, algorithm="HS256")
+            container = await containerASD.docker_client.containers.run(
+                name=container_name,
+                config={
+                    "Image": ctf.image_name,
+                    "AttachStdin": False,
+                    "AttachStdout": False,
+                    "AttachStderr": False,
+                    "Tty": False,
+                    "OpenStdin": False,
+                    "HostConfig": {
+                        "AutoRemove": True,
+                        "Binds": [f"{config.staticfs_data_dir}/{team_id}/{container_id}:/dist"],
+                    },
+                    "Cmd": ["/root/gen_flag", container_flag],
+                },
+            )
+            try:
+                async with in_transaction():
+                    db_container = await Container.create(
+                        docker_id=container_id,
+                        team_id=team_id,
+                        problem_id=ctf_id,
+                        flag=container_flag,
+                        token=token,
+                    )
+            except Exception as err:
+                # Stop the container if failed to make a DB record
+                await container.kill()
+                await container.delete()
+                logger.exception("Error while starting", exc_info=err)
+
+                response.status_code = 500
+                return {"msg_code": config.msg_codes["db_error"]}
+            return {
+                "msg_code": config.msg_codes["container_start"],
+                "ports": [],
+                "static_url": f"{config.staticfs_url}/{token}",
+                "ctf_id": ctf_id,
+            }
+
         # Run
         container = await containerASD.docker_client.containers.run(
             name=container_name,
@@ -91,9 +145,9 @@ async def start_docker_container(ctf_id: int, response: Response, jwt: RequireJw
             },
         )
 
-        await (
-            await container.exec(["/root/gen_flag", container_flag])
-        ).start(detach=True)
+        await (await container.exec(["/root/gen_flag", container_flag])).start(
+            detach=True
+        )
 
         try:
             async with in_transaction():
@@ -126,6 +180,7 @@ async def start_docker_container(ctf_id: int, response: Response, jwt: RequireJw
         return {
             "msg_code": config.msg_codes["container_start"],
             "ports": ports,
+            "static_url": None,
             "ctf_id": ctf_id,
         }
 
@@ -144,7 +199,11 @@ async def stopall_docker_container(response: Response, jwt: RequireJwt):
         except Exception:
             response.status_code = 500
             return {"msg_code": config.msg_codes["db_error"]}
-
+        
+        team_path = f"{config.staticfs_data_dir}/{team_id}"
+        if os.path.exists(team_path):
+            shutil.rmtree(team_path)
+            
         for db_container in containers:
             container = await containerASD.docker_client.containers.get(
                 db_container["docker_id"]
@@ -178,6 +237,12 @@ async def stop_docker_container(ctf_id: int, response: Response, jwt: RequireJwt
             response.status_code = 500
             return {"msg_code": config.msg_codes["db_error"]}
 
+        if ctf.static:
+            shutil.rmtree(
+                f"{config.staticfs_data_dir}/{team_id}/{team_container.docker_id}"
+            )
+            return {"msg_code": config.msg_codes["container_stop"]}
+
         container = await containerASD.docker_client.containers.get(
             team_container.docker_id
         )