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,12 @@
     PreEventUser,
     PreEventProblem_Pydantic,
 )
+from pwncore.models.round2 import (
+    R2Problem,
+    R2Ports,
+    R2Container,
+)
+
 
 __all__ = (
     "Problem",
@@ -48,4 +60,21 @@
     "MetaTeam_Pydantic",
     "PreEventProblem_Pydantic",
     "Problem_Pydantic",
+    "BaseProblem",
+    "R2Problem",
+    "R2Ports",
+    "R2Container",
+    "R2Container_Pydantic",
 )
+
+
+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,52 @@
+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",
+)
+
+
+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)

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,6 +1,9 @@
+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 pwncore.models import (
@@ -12,8 +15,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 +34,76 @@
     logging.basicConfig(level=logging.INFO)
 
 
+async def _del_cont(id: str):
+    container = await docker_client.containers.get(id)
+    await container.stop()
+    await container.delete()
+
+
+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:")

src/pwncore/routes/leaderboard.py

@@ -3,7 +3,7 @@
 from time import monotonic
 
 from fastapi import APIRouter, Request
-from tortoise.functions import Sum
+
 from tortoise.expressions import RawSQL
 
 from pwncore.models import Team
@@ -29,13 +29,13 @@ async def _do_update(self):
             await Team.all()
             .filter(solved_problem__problem__visible=True)
             .annotate(
-                tpoints=Sum(
-                    RawSQL(
-                        '"solvedproblem"."penalty" * "solvedproblem__problem"."points"'
-                    )
+                tpoints=RawSQL(
+                    'SUM("solvedproblem"."penalty" * "solvedproblem__problem"."points")'
+                    ' + "team"."points"'
                 )
             )
-            .values("name", "tpoints")
+            .group_by("id")
+            .values("name", "tpoints", "meta_team__name")
         )
         self.last_update = monotonic()
 

src/pwncore/routes/round2.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from fastapi import APIRouter, Response
+import tortoise.exceptions
+from tortoise.expressions import F
+
+from pwncore.config import config
+from pwncore.models import (
+    R2Container,
+    R2Container_Pydantic,
+    MetaTeam_Pydantic,
+    MetaTeam,
+    Team,
+)
+from pwncore.routes.auth import RequireJwt
+from pwncore.routes.ctf import Flag
+
+metadata = {
+    "name": "round2",
+    "decription": "Operations related to an attack defense round",
+}
+
+router = APIRouter(prefix="/round2", tags=["round2"])
+
+
+@router.get("/list")
+async def r2ctf_list():
+    return await R2Container_Pydantic.from_queryset(
+        R2Container.all().prefetch_related("problem", "ports")
+    )
+
+
+@router.get("/meta_lb")
+async def r2_meta_lb():
+    return await MetaTeam_Pydantic.from_queryset(MetaTeam.all())
+
+
+@router.post("/{container_id}/flag")
+async def r2_submit(container_id: int, flag: Flag, jwt: RequireJwt, response: Response):
+    team_id = jwt["team_id"]
+    try:
+        container = await R2Container.get(id=container_id).prefetch_related(
+            "meta_team", "problem"
+        )
+    except tortoise.exceptions.DoesNotExist:
+        response.status_code = 404
+        return {"msg_code": config.msg_codes["ctf_not_found"]}
+
+    if container.solved:
+        response.status_code = 401
+        return {"msg_code": config.msg_codes["ctf_solved"]}
+
+    team = await Team.get(id=team_id).prefetch_related("meta_team")
+
+    if team.meta_team is None:
+        response.status_code = 401
+        return ""
+
+    if container.flag != flag.flag:
+        return {"status": False}
+
+    if team.meta_team == container.meta_team:
+        container.solved = True
+        await container.save()
+        return {"status": True, "action": "defend"}
+    else:
+        # Valid type error but we dont use this object again
+        team.points = F("points") + round(container.problem.points / 4)  # type: ignore[unused-ignore]
+        await team.save(update_fields=["points"])
+        team.meta_team.points = F("points") + container.problem.points  # type: ignore[unused-ignore]
+        await team.meta_team.save(update_fields=["points"])
+        return {"status": True, "action": "attack"}