Add AttackDefTeam

Added new model with two new routes under '/ctf/round2/list' and
'/team/round2/powerups' to list a team's ctfs and used powerups
respectively.

Also added test rows in admin init for new models.

Author: tnmkr

src/pwncore/config.py

@@ -1,5 +1,7 @@
 import os
 from dataclasses import dataclass
+from enum import Enum
+from datetime import timedelta
 
 """
 Sample messages:
@@ -40,8 +42,14 @@
     "insufficient_coins": 22,
     "user_or_email_exists": 23,
     "users_not_found": 24,
+    "attack_def_team_not_found": 25
 }
 
+class PowerUpType(str, Enum):
+    SHIELD = "SHIELD"
+    POINT_SIPHON = "POINT_SIPHON"
+    SABOTAGE = "SABOTAGE"
+    UPGRADE = "UPGRADE"
 
 @dataclass
 class Config:
@@ -55,7 +63,7 @@ class Config:
     jwt_valid_duration: int
     hint_penalty: int
     max_members_per_team: int
-
+    powerups: dict[PowerUpType, dict[str, any]]
 
 config = Config(
     development=True,
@@ -71,4 +79,21 @@ class Config:
     msg_codes=msg_codes,
     hint_penalty=50,
     max_members_per_team=3,
+    powerups={
+        PowerUpType.UPGRADE: {
+            "cost": 100,
+        },
+        PowerUpType.SHIELD: {
+            "duration": timedelta(seconds=10),
+            "cost": 200,
+        },
+        PowerUpType.POINT_SIPHON: {
+            "duration": timedelta(seconds=15),
+            "cost": 150,
+        },
+        PowerUpType.SABOTAGE: {
+            "duration": timedelta(seconds=5),
+            "cost": 500,
+        },
+    }
 )

src/pwncore/models/__init__.py

@@ -25,6 +25,12 @@
     Team_Pydantic,
     User_Pydantic,
 )
+from pwncore.models.round2 import (
+    AttackDefProblem,
+    AttackDefTeam,
+    ActivatedPowerups,
+    ActivatedPowerups_Pydantic
+)
 from pwncore.models.pre_event import (
     PreEventProblem,
     PreEventSolvedProblem,
@@ -52,6 +58,10 @@
     "PreEventProblem_Pydantic",
     "Problem_Pydantic",
     "BaseProblem",
+    "AttackDefProblem",
+    "AttackDefTeam",
+    "ActivatedPowerups",
+    "ActivatedPowerups_Pydantic"
 )
 
 

src/pwncore/models/round2.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+from datetime import datetime, timezone
+from tortoise import fields
+from tortoise.models import Model
+from tortoise.contrib.pydantic import pydantic_model_creator
+from pwncore.models import (
+    Problem
+)
+from pwncore.config import config, PowerUpType
+
+__all__ = (
+    "AttackDefProblem",
+    "AttackDefTeam",
+    "ActivatedPowerups"
+)
+#TODO: Actually implement this.
+# For now this is a dummy class for testing AttackDefTeam.
+class AttackDefProblem(Model):
+    problem: fields.OneToOneRelation[Problem] = fields.OneToOneField(
+        "models.Problem"
+    )
+    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)
+
+
+class ActivatedPowerups(Model):
+    id = fields.IntField(pk=True)
+    used_by: fields.OneToOneRelation[AttackDefTeam] = fields.OneToOneField(
+        "models.AttackDefTeam", related_name="powerups_used", null=False
+    )
+    used_on: fields.OneToOneNullableRelation[AttackDefTeam] = fields.OneToOneField(
+        "models.AttackDefTeam", related_name="powerups_defended"
+    )
+    powerup_type = fields.data.CharEnumField(enum_type=PowerUpType, null=False)
+    started_at = fields.data.DatetimeField(auto_now_add=True)
+
+    def is_active(self) -> bool:
+        duration = config.powerups[self.powerup_type]["duration"]
+        if ((self.started_at + duration) < datetime.now(timezone.utc)):
+            return False
+
+        return True
+
+    class PydanticMeta:
+        arbitrary_types_allowed = True
+        computed = ("is_active",)
+
+ActivatedPowerups_Pydantic = pydantic_model_creator(ActivatedPowerups)

src/pwncore/routes/admin.py

@@ -6,14 +6,17 @@
 from tortoise.transactions import atomic, in_transaction
 
 import pwncore.containerASD as containerASD
-from pwncore.config import config
+from pwncore.config import config, PowerUpType
 from pwncore.models import (
     Hint,
     PreEventProblem,
     PreEventSolvedProblem,
     Problem,
     Team,
     User,
+    AttackDefTeam,
+    AttackDefProblem,
+    ActivatedPowerups
 )
 from pwncore.models.ctf import SolvedProblem
 from pwncore.models.pre_event import PreEventUser
@@ -146,10 +149,37 @@ 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(
+        name="Triple B battery", secret_hash=bcrypt.hash("chotiwali2"), coins=20
+    )
+    await AttackDefTeam.create(
+        team_id=(await Team.get(name="Triple A battery")).id
+    )
+    await AttackDefTeam.create(
+        team_id=(await Team.get(name="Triple B battery")).id
+    )
+    await AttackDefProblem.create(
+        problem_id=(await Problem.get(name="GitGood2")).id,
+        attack_def_team_id=(await AttackDefTeam.get(team_id=triple_b_battery.id)).id
+    )
+    await ActivatedPowerups.create(
+        used_by_id=(await AttackDefTeam.get(team_id=triple_b_battery.id)).id,
+        used_on_id = (await AttackDefTeam.get(team_id=triple_a_battery.id)).id,
+        powerup_type = PowerUpType.SABOTAGE
+    )
+
     await PreEventUser.create(tag="23BCE1000", email="[email protected]")
     await PreEventUser.create(tag="23BRS1000", email="[email protected]")
     await PreEventSolvedProblem.create(user_id="23BCE1000", problem_id="1")

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:

src/pwncore/routes/team.py

@@ -5,7 +5,16 @@
 from tortoise.transactions import atomic
 
 from pwncore.config import config
-from pwncore.models import Team, User, Team_Pydantic, User_Pydantic, Container
+from pwncore.models import (
+    Team,
+    User,
+    Team_Pydantic,
+    User_Pydantic,
+    Container,
+    ActivatedPowerups,
+    ActivatedPowerups_Pydantic,
+    AttackDefTeam
+)
 from pwncore.routes.auth import RequireJwt
 
 # from pwncore.routes.leaderboard import gcache
@@ -116,3 +125,13 @@ async def get_team_containers(response: Response, jwt: RequireJwt):
         )
 
     return result
+
+@router.get("/round2/powerups")
+async def get_team_powerups(response: Response, jwt:RequireJwt):
+    team_id = jwt["team_id"]
+    team = await Team.get(id=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"]}
+    activated_powerups = await ActivatedPowerups_Pydantic.from_queryset(ActivatedPowerups.filter(used_by_id=attack_def_team.id))
+    return activated_powerups