diff --git a/src/pwncore/config.py b/src/pwncore/config.py index f1b818b..68126b4 100644 --- a/src/pwncore/config.py +++ b/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 } diff --git a/src/pwncore/models/__init__.py b/src/pwncore/models/__init__.py index fbb4682..75273f7 100644 --- a/src/pwncore/models/__init__.py +++ b/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", ) diff --git a/src/pwncore/models/powerups.py b/src/pwncore/models/powerups.py new file mode 100644 index 0000000..54b30c6 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/src/pwncore/models/round2.py b/src/pwncore/models/round2.py new file mode 100644 index 0000000..f82ed0b --- /dev/null +++ b/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')) diff --git a/src/pwncore/routes/__init__.py b/src/pwncore/routes/__init__.py index d113e38..7dc086c 100644 --- a/src/pwncore/routes/__init__.py +++ b/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) diff --git a/src/pwncore/routes/admin.py b/src/pwncore/routes/admin.py index d6a5a2b..2c6c047 100644 --- a/src/pwncore/routes/admin.py +++ b/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="dd@ff.in") await PreEventUser.create(tag="23BRS1000", email="d2d@ff.in") await PreEventSolvedProblem.create(user_id="23BCE1000", problem_id="1") diff --git a/src/pwncore/routes/ctf/__init__.py b/src/pwncore/routes/ctf/__init__.py index 40ce348..21e4598 100644 --- a/src/pwncore/routes/ctf/__init__.py +++ b/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") diff --git a/src/pwncore/routes/powerups.py b/src/pwncore/routes/powerups.py new file mode 100644 index 0000000..f589dde --- /dev/null +++ b/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