Add round 2 routes along with the admin interface

src/pwncore/models/__init__.py

@@ -3,6 +3,11 @@
 Contains all Pydantic and Tortoise ORM models
 """
 
+import typing
+
+import tortoise
+from tortoise.contrib.pydantic.creator import pydantic_model_creator
+
 from pwncore.models.container import Container, Ports
 from pwncore.models.ctf import (
     Problem,
@@ -12,6 +17,7 @@
     Hint_Pydantic,
     BaseProblem_Pydantic,
     Problem_Pydantic,
+    BaseProblem,
 )
 from pwncore.models.user import (
     User,
@@ -27,6 +33,13 @@
     PreEventUser,
     PreEventProblem_Pydantic,
 )
+from pwncore.models.round2 import (
+    R2Problem,
+    R2Ports,
+    R2Container,
+    R2AttackRecord,
+)
+
 
 __all__ = (
     "Problem",
@@ -48,4 +61,22 @@
     "MetaTeam_Pydantic",
     "PreEventProblem_Pydantic",
     "Problem_Pydantic",
+    "BaseProblem",
+    "R2Problem",
+    "R2Ports",
+    "R2Container",
+    "R2Container_Pydantic",
+    "R2AttackRecord",
 )
+
+
+def get_annotations(cls, method=None):
+    return typing.get_type_hints(method or cls)
+
+
+tortoise.contrib.pydantic.utils.get_annotations = get_annotations  # type: ignore[unused-ignore]
+tortoise.contrib.pydantic.creator.get_annotations = get_annotations  # type: ignore[unused-ignore]
+
+
+tortoise.Tortoise.init_models(["pwncore.models.round2"], "models")
+R2Container_Pydantic = pydantic_model_creator(R2Container)

src/pwncore/models/ctf.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from math import tanh
+from typing import Any
 
 from tortoise.models import Model
 from tortoise import fields
@@ -16,6 +17,7 @@
     "BaseProblem_Pydantic",
     "Hint_Pydantic",
     "Problem_Pydantic",
+    "BaseProblem",
 )
 
 
@@ -29,7 +31,7 @@ class BaseProblem(Model):
 
 class Problem(BaseProblem):
     image_name = fields.TextField()
-    image_config: fields.Field[dict[str, list]] = fields.JSONField(
+    image_config: fields.Field[dict[str, Any]] = fields.JSONField(
         null=True
     )  # type: ignore[assignment]
 

src/pwncore/models/round2.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+from typing import Any
+
+from tortoise import fields
+from tortoise.models import Model
+
+from pwncore.models.ctf import BaseProblem
+from pwncore.models.user import MetaTeam
+
+__all__ = ("R2Problem", "R2Container", "R2Ports", "R2AttackRecord")
+
+
+class R2Problem(BaseProblem):
+    image_name = fields.TextField()
+    image_config: fields.Field[dict[str, Any]] = fields.JSONField(
+        null=True
+    )  # type: ignore[assignment]
+
+    class PydanticMeta:
+        exclude = [
+            "image_name",
+            "image_config",
+        ]
+
+
+class R2Container(Model):
+    docker_id = fields.CharField(128, unique=True)
+    problem: fields.ForeignKeyRelation[R2Problem] = fields.ForeignKeyField(
+        "models.R2Problem", on_delete=fields.OnDelete.NO_ACTION
+    )
+    meta_team: fields.ForeignKeyRelation[MetaTeam] = fields.ForeignKeyField(
+        "models.MetaTeam", on_delete=fields.OnDelete.NO_ACTION
+    )
+    flag = fields.TextField()
+    solved = fields.BooleanField(default=False)
+
+    ports: fields.ReverseRelation[R2Ports]
+
+    class PydanticMeta:
+        exclude = ["docker_id", "flag", "meta_team"]
+
+
+class R2Ports(Model):
+    container: fields.ForeignKeyRelation[R2Container] = fields.ForeignKeyField(
+        "models.R2Container", related_name="ports", on_delete=fields.OnDelete.CASCADE
+    )
+    port = fields.IntField(pk=True)
+
+
+class R2AttackRecord(Model):
+    container: fields.ForeignKeyRelation[R2Container] = fields.ForeignKeyField(
+        "models.R2Container", on_delete=fields.OnDelete.CASCADE
+    )
+
+    meta_team: fields.ForeignKeyRelation[MetaTeam] = fields.ForeignKeyField(
+        "models.MetaTeam", on_delete=fields.OnDelete.NO_ACTION
+    )

src/pwncore/models/user.py

@@ -52,6 +52,7 @@ class Team(Model):
     name = fields.CharField(255, unique=True)
     secret_hash = fields.TextField()
     coins = fields.IntField(default=0)
+    points = fields.IntField(default=0)
 
     members: fields.ReverseRelation[User]
     containers: fields.ReverseRelation[Container]
@@ -67,6 +68,7 @@ class PydanticMeta:
 class MetaTeam(Model):
     id = fields.IntField(pk=True)
     name = fields.CharField(255, unique=True)
+    points = fields.IntField(default=0)
 
     teams: fields.ReverseRelation[Team]
 

src/pwncore/routes/__init__.py

@@ -1,6 +1,6 @@
 from fastapi import APIRouter
 
-from pwncore.routes import ctf, team, auth, admin, leaderboard
+from pwncore.routes import ctf, team, auth, admin, leaderboard, round2
 from pwncore.config import config
 
 # Main router (all routes go under /api)
@@ -11,5 +11,6 @@
 router.include_router(ctf.router)
 router.include_router(team.router)
 router.include_router(leaderboard.router)
+router.include_router(round2.router)
 if config.development:
     router.include_router(admin.router)

src/pwncore/routes/admin.py

@@ -1,7 +1,12 @@
+import asyncio
 from datetime import date
+import itertools
 import logging
-from fastapi import APIRouter
+import uuid
+
+from fastapi import APIRouter, Response
 from passlib.hash import bcrypt
+from tortoise.transactions import atomic
 
 from pwncore.models import (
     Team,
@@ -12,8 +17,12 @@
     PreEventProblem,
 )
 from pwncore.config import config
+from pwncore.models.container import Container
 from pwncore.models.ctf import SolvedProblem
 from pwncore.models.pre_event import PreEventUser
+from pwncore.container import docker_client
+from pwncore.models.round2 import R2Container, R2Ports, R2Problem
+from pwncore.models.user import MetaTeam
 
 metadata = {
     "name": "admin",
@@ -27,6 +36,77 @@
     logging.basicConfig(level=logging.INFO)
 
 
+async def _del_cont(id: str):
+    container = await docker_client.containers.get(id)
+    await container.stop()
+    await container.delete()
+
+
+@atomic()
+async def _create_container(prob: R2Problem, mteam: MetaTeam):
+    try:
+        container = await docker_client.containers.run(
+            name=f"{mteam.name}_{prob.pk}_{uuid.uuid4().hex}",
+            config={
+                "Image": prob.image_name,
+                # Detach stuff
+                "AttachStdin": False,
+                "AttachStdout": False,
+                "AttachStderr": False,
+                "Tty": False,
+                "OpenStdin": False,
+                **prob.image_config,
+            },
+        )
+
+        container_flag = f"{config.flag}{{{uuid.uuid4().hex}}}"
+
+        await (
+            await container.exec(["/bin/bash", "/root/gen_flag", container_flag])
+        ).start(detach=True)
+
+        db_container = await R2Container.create(
+            docker_id=container.id, problem=prob, meta_tam=mteam, flag=container_flag
+        )
+
+        ports = []
+        for guest_port in prob.image_config["PortBindings"]:
+            port = int((await container.port(guest_port))[0]["HostPort"])
+            ports.append(R2Ports(port=port, container=db_container))
+
+        await R2Ports.bulk_create(ports)
+    except Exception as err:
+        try:
+            await container.stop()
+            await container.delete()
+        except Exception:
+            pass
+        logging.exception("Error while starting", exc_info=err)
+
+
+@router.get("/round2")
+async def round2(response: Response):
+    containers = await Container.all()
+
+    async with asyncio.TaskGroup() as tg:
+        for container in containers:
+            tg.create_task(_del_cont(container.docker_id))
+
+    try:
+        await Container.all().delete()
+    except Exception:
+        response.status_code = 500
+        logging.exception("Error while initing round2")
+        return {"msg_code": config.msg_codes["db_error"]}
+
+    problems = await R2Problem.all()
+    mteams = await MetaTeam.all()
+
+    async with asyncio.TaskGroup() as tg:
+        for pm in itertools.product(problems, mteams):
+            tg.create_task(_create_container(*pm))
+
+
 @router.get("/union")
 async def calculate_team_coins():  # Inefficient, anyways will be used only once
     logging.info("Calculating team points form pre-event CTFs:")