feat: add AttackDefProblem model

Author: tnmkr

src/pwncore/models/__init__.py

@@ -24,6 +24,9 @@
     Team_Pydantic,
     User_Pydantic,
 )
+from pwncore.models.round2 import (
+    AttackDefProblem,
+)
 from pwncore.models.pre_event import (
     PreEventProblem,
     PreEventSolvedProblem,
@@ -51,6 +54,7 @@
     "PreEventProblem_Pydantic",
     "Problem_Pydantic",
     "BaseProblem",
+    "AttackDefProblem",
 )
 
 

src/pwncore/models/powerups.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+from datetime import timedelta
+from dataclasses import dataclass
+from typing import ClassVar
+from pydantic import field_serializer
+
+from tortoise import fields
+from tortoise.contrib.pydantic.creator import pydantic_model_creator
+from tortoise.models import Model
+
+from pwncore.models import (
+    AttackDefTeam
+)
+
+__all__ = ("PowerupType", "UsedPowerup", "Powerup", "Powerup_Pydantic")
+
+@dataclass
+class PowerupType:
+    name: str
+    cost: int
+    max_uses_default: int = 1
+    duration: timedelta | None = None
+    all_powerup_types: ClassVar[dict[str, PowerupType]] = {}
+
+    def __post_init__(self):
+        self.all_powerup_types[self.name] = self
+
+class PowerupTypeField(fields.CharField):
+    def __init__(self, **kwargs):
+        super().__init__(128, **kwargs)
+        # if not issubclass(enum_type, Enum):
+            # raise ConfigurationError("{} is not a subclass of Enum!".format(enum_type))
+        # self._powerup_type = powerup_type
+
+    def to_db_value(self, powerup_type: PowerupType, instance) -> str:
+        return powerup_type.name
+
+    def to_python_value(self, value: PowerupType) -> PowerupType:
+        return value
+
+Sabotage = PowerupType(name = "Sabotage", cost = 500, duration = timedelta(seconds=5))
+Shield = PowerupType(name = "Shield", cost = 200, max_uses_default = 3, duration = timedelta(seconds=10))
+PointSiphon = PowerupType(name = "PointSiphon", cost = 150, duration = timedelta(seconds=15))
+Upgrade = PowerupType(name = "Upgrade", cost = 100)
+
+class Powerup(Model):
+    id = fields.IntField(pk=True)
+    powerup_type = PowerupTypeField()
+    attack_def_team: fields.ForeignKeyRelation[AttackDefTeam] = fields.ForeignKeyField(
+        "models.AttackDefTeam", related_name="powerups"
+    )
+    created_at = fields.DatetimeField(auto_now_add=True)
+    used_at_least_once = fields.BooleanField(default=False)
+    max_uses = fields.IntField(default=1)
+    used_count = fields.IntField(default=0)
+    used_instance = fields.ReverseRelation["UsedPowerup"]
+
+    async def save(self, *args, **kwargs):
+        if not self.max_uses:
+            self.max_uses = all_powerup_types[self.powerup_type].max_uses_default
+        await super().save(*args, **kwargs)
+
+class UsedPowerup(Model):
+    id = fields.IntField(pk=True)
+    powerup: fields.ForeignKeyRelation[Powerup] = fields.ForeignKeyField(
+        "models.Powerup", related_name="used_instances", null=False
+    )
+    used_at = fields.DatetimeField(auto_now_add=True)
+
+Powerup_Pydantic = pydantic_model_creator(Powerup)

src/pwncore/models/round2.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+from tortoise import fields
+from tortoise.models import Model
+__all__ = (
+    "AttackDefProblem",
+    "Problem",
+)
+#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"
+    )

src/pwncore/routes/admin.py

@@ -14,6 +14,8 @@
     Problem,
     Team,
     User,
+    AttackDefTeam,
+    AttackDefProblem,
 )
 from pwncore.models.ctf import SolvedProblem
 from pwncore.models.pre_event import PreEventUser
@@ -150,6 +152,19 @@ async def init_db(
     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 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/powerups.py

@@ -0,0 +1,67 @@
+
+from tortoise.contrib.pydantic.creator import pydantic_model_creator
+from fastapi import APIRouter, HTTPException, Response
+
+from pwncore.config import config
+from pwncore.models import (
+    Powerup, 
+    PowerupType, 
+    Team,
+    AttackDefTeam,
+    UsedPowerup,
+    Powerup_Pydantic,
+)
+from pwncore.routes.auth import RequireJwt
+
+metadata = {"name": "powerups", "description": "Powerups for teams"}
+router = APIRouter(prefix="/powerups", tags=["team"])
+
+@router.get("/remaining")
+async def view_remaining_powerups(jwt: RequireJwt, response: Response):
+    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"]}
+    return await Powerup_Pydantic.from_queryset(await attack_def_team.remaining_powerups())
+
+@router.get("/used")
+async def view_used_powerups(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"]}
+    
+    used_powerups = UsedPowerup.filter(powerup__attack_def_team=attack_def_team.id).prefetch_related("powerup")
+    return await used_powerups
+
+
+# @router.post("/use/{powerup_type}")
+# async def use_powerup(powerup_type: PowerupType, jwt: RequireJwt):
+    # Define logic to use available powerups and deduct points
+    # pass
+
+
+# @router.get("/about/{powerup_type}")
+# async def get_about_powerups(powerup_type: PowerupType, jwt: RequireJwt):
+    # Return the about of the powerup
+    # pass
+
+
+# async def lucky_draw(team: Team):
+#     # Lucky draw logic
+#     pass
+
+
+# async def sabotage(team: Team):
+#     # Sabotage logic
+#     pass
+
+
+# async def gamble(team: Team):
+#     # Gamble logic
+#     pass
+
+
+# async def point_siphon(team: Team):
+#     # Point Siphon logic
+#     pass