Add AttackDefTeam

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

@@ -24,6 +24,17 @@
     Team_Pydantic,
     User_Pydantic,
 )
+from pwncore.models.round2 import (
+    AttackDefProblem,
+    AttackDefTeam
+)
+
+from pwncore.models.powerups import (
+    Powerup,
+    PowerupType,
+    UsedPowerup,
+    Powerup_Pydantic,
+)
 from pwncore.models.pre_event import (
     PreEventProblem,
     PreEventSolvedProblem,
@@ -51,6 +62,11 @@
     "PreEventProblem_Pydantic",
     "Problem_Pydantic",
     "BaseProblem",
+    "AttackDefProblem",
+    "AttackDefTeam",
+    "Powerup",
+    "UsedPowerup",
+    "Powerup_Pydantic",
 )
 
 

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,28 @@
+from __future__ import annotations
+from tortoise import fields
+from tortoise.models import Model
+from tortoise.expressions import F
+
+__all__ = (
+    "AttackDefProblem",
+    "AttackDefTeam",
+    "Problem",
+    "Powerup",
+)
+#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)
+    async def remaining_powerups(self):
+        return self.powerups.filter(used_count__lt=F('max_uses'))

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, powerups
 
 # from pwncore.config import config
 
@@ -13,5 +13,6 @@
 router.include_router(team.router)
 router.include_router(leaderboard.router)
 router.include_router(admin.router)
+router.include_router(powerups.router)
 # if config.development:
 # router.include_router(admin.router)

src/pwncore/routes/admin.py

@@ -14,9 +14,17 @@
     Problem,
     Team,
     User,
+    AttackDefTeam,
+    AttackDefProblem,
+    Powerup,
+    UsedPowerup
 )
 from pwncore.models.ctf import SolvedProblem
 from pwncore.models.pre_event import PreEventUser
+from pwncore.models.powerups import (
+    Sabotage,
+    Shield
+)
 
 metadata = {
     "name": "admin",
@@ -146,10 +154,43 @@ 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 Powerup.create(
+        attack_def_team = (await AttackDefTeam.get(team__id=triple_b_battery.id)),
+        powerup_type = Sabotage
+    )
+    shield = await Powerup.create(
+        attack_def_team = (await AttackDefTeam.get(team__id=triple_b_battery.id)),
+        powerup_type = Shield
+    )
+    await UsedPowerup.create(
+        powerup = shield
+    )
+
     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:
@@ -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")

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