feat: add AttackDefTeam model and routes

Author: tnmkr

src/pwncore/config.py

@@ -40,6 +40,7 @@
     "insufficient_coins": 22,
     "user_or_email_exists": 23,
     "users_not_found": 24,
+    "attack_def_team_not_found": 25
 }
 
 

src/pwncore/models/__init__.py

@@ -26,7 +26,9 @@
 )
 from pwncore.models.round2 import (
     AttackDefProblem,
+    AttackDefTeam
 )
+
 from pwncore.models.pre_event import (
     PreEventProblem,
     PreEventSolvedProblem,
@@ -55,6 +57,7 @@
     "Problem_Pydantic",
     "BaseProblem",
     "AttackDefProblem",
+    "AttackDefTeam",
 )
 
 

src/pwncore/models/round2.py

@@ -3,6 +3,7 @@
 from tortoise.models import Model
 __all__ = (
     "AttackDefProblem",
+    "AttackDefTeam",
     "Problem",
 )
 #TODO: Actually implement this.
@@ -14,3 +15,9 @@ class AttackDefProblem(Model):
     attack_def_team: fields.ForeignKeyRelation[AttackDefTeam] = fields.ForeignKeyField(
         "models.AttackDefTeam", related_name="assigned_attack_def_problem"
     )
+
+class AttackDefTeam(Model):
+    team: fields.OneToOneRelation[Team] = fields.OneToOneField(
+        "models.Team", null=False
+    )
+    assigned_attack_def_problem: fields.ReverseRelation(AttackDefProblem)

src/pwncore/routes/admin.py

@@ -148,8 +148,16 @@ async def init_db(
         image_name="reg.lugvitc.net/key:latest",
         # image_config={"PortBindings": {"22/tcp": [{}]}},
     )
+    await Problem.create(
+        name="GitGood2",
+        description="How to master the art of solving CTFs? Git good nub.",
+        author="Aadivishnu and Shoubhit",
+        points=300,
+        image_name="reg.lugvitc.net/key:latest",
+        # image_config={"PortBindings": {"22/tcp": [{}]}},
+    )
     await Team.create(name="CID Squad", secret_hash=bcrypt.hash("veryverysecret"))
-    await Team.create(
+    triple_a_battery = await Team.create(
         name="Triple A battery", secret_hash=bcrypt.hash("chotiwali"), coins=20
     )
     triple_b_battery = await Team.create(

src/pwncore/routes/ctf/__init__.py

@@ -18,6 +18,8 @@
     SolvedProblem,
     Team,
     ViewedHint,
+    AttackDefProblem,
+    AttackDefTeam
 )
 from pwncore.models.ctf import Problem_Pydantic
 from pwncore.routes.auth import RequireJwt
@@ -62,11 +64,7 @@ async def completed_problem_get(jwt: RequireJwt):
     )
     return problems
 
-
-@router.get("/list")
-async def ctf_list(jwt: RequireJwt):
-    team_id = jwt["team_id"]
-    problems = await Problem_Pydantic.from_queryset(Problem.filter(visible=True))
+async def calculate_problem_points(problems: [Problem_Pydantic], team_id: int) -> [Problem_Pydantic]:
     acc: dict[int, float] = defaultdict(lambda: 1.0)
     for k, v in map(
         lambda x: (x.hint.problem_id, HINTPENALTY[x.hint.order]),  # type: ignore[attr-defined]
@@ -79,6 +77,20 @@ async def ctf_list(jwt: RequireJwt):
         i.points = int(acc[i.id] * i.points)  # type: ignore[attr-defined]
     return problems
 
+@router.get("/list")
+async def ctf_list(jwt: RequireJwt):
+    team_id = jwt["team_id"]
+    problems = await Problem_Pydantic.from_queryset(Problem.filter(visible=True))
+    return await calculate_problem_points(problems, team_id)
+
+@router.get("/round2/list")
+async def ctf_list(jwt: RequireJwt):
+    team_id = jwt["team_id"]
+    attack_def_team = await AttackDefTeam.get(team_id=team_id)
+    if (attack_def_team is None):
+        return {"msg_code": config.msg_codes["attack_def_team_not_found"]}
+    problems = await Problem_Pydantic.from_queryset(Problem.filter(attackdefproblem__attack_def_team__id=attack_def_team.id))
+    return await calculate_problem_points(problems, team_id)
 
 async def update_points(req: Request, ctf_id: int):
     try:
@@ -88,6 +100,9 @@ async def update_points(req: Request, ctf_id: int):
     except Exception:
         logger.exception("An error occured while updating points")
 
+@router.get("/round2/list_all")
+async def ctf_list():
+    return await Problem_Pydantic.from_queryset(Problem.all())
 
 @atomic()
 @router.post("/{ctf_id}/flag")