Merge pull request #12 from mradigen/master

Add coins to teams and user add/remove route

src/pwncore/config.py

@@ -32,6 +32,11 @@
     "wrong_password": 14,
     "login_success": 15,
     "team_exists": 17,
+    "user_added": 18,
+    "user_removed": 19,
+    "user_already_in_team": 20,
+    "user_not_in_team": 21,
+    "insufficient_coins": 22,
 }
 
 
@@ -45,15 +50,21 @@ class Config:
     max_containers_per_team: int
     jwt_secret: str
     jwt_valid_duration: int
+    hint_penalty: int
+    max_members_per_team: int
 
 
 config = Config(
     development=True,
     db_url="sqlite://:memory:",
-    docker_url=None,  # None for default system docker
+    # docker_url=None,  # None for default system docker
+    # Or set it to an arbitrary URL for testing without Docker
+    docker_url="http://google.com",
     flag="C0D",
     max_containers_per_team=3,
     jwt_secret="mysecret",
     jwt_valid_duration=12,  # In hours
     msg_codes=msg_codes,
+    hint_penalty=10,
+    max_members_per_team=3,
 )

src/pwncore/models/__init__.py

@@ -9,21 +9,27 @@
     SolvedProblem,
     Hint,
     ViewedHint,
-    Problem_Pydantic,
     Hint_Pydantic,
+    BaseProblem_Pydantic,
 )
 from pwncore.models.user import User, Team, Team_Pydantic, User_Pydantic
+from pwncore.models.pre_event import (
+    PreEventProblem,
+    PreEventSolvedProblem,
+)
 
 __all__ = (
     "Problem",
-    "Problem_Pydantic",
+    "BaseProblem_Pydantic",
     "Hint",
     "Hint_Pydantic",
     "SolvedProblem",
     "ViewedHint",
     "Container",
     "Ports",
     "User",
+    "PreEventSolvedProblem",
+    "PreEventProblem",
     "User_Pydantic",
     "Team",
     "Team_Pydantic",

src/pwncore/models/ctf.py

@@ -11,25 +11,27 @@
     "Hint",
     "SolvedProblem",
     "ViewedHint",
-    "Problem_Pydantic",
+    "BaseProblem_Pydantic",
     "Hint_Pydantic",
 )
 
 
-class Problem(Model):
+class BaseProblem(Model):
     name = fields.TextField()
     description = fields.TextField()
+    # both tables inherit points, for pre-event points means coins
     points = fields.IntField()
     author = fields.TextField()
 
+
+class Problem(BaseProblem):
     image_name = fields.TextField()
-    image_config: fields.Field[dict[str, list]] = fields.JSONField()  # type: ignore[assignment]
+    image_config: fields.Field[dict[str, list]] = fields.JSONField(
+        null=True
+    )  # type: ignore[assignment]
 
     hints: fields.ReverseRelation[Hint]
 
-    class PydanticMeta:
-        exclude = ["image_name", "image_config"]
-
 
 class Hint(Model):
     id = fields.IntField(pk=True)
@@ -44,7 +46,9 @@ class Meta:
 
 
 class SolvedProblem(Model):
-    team: fields.ForeignKeyRelation[Team] = fields.ForeignKeyField("models.Team")
+    team: fields.ForeignKeyRelation[Team] = fields.ForeignKeyField(
+        "models.Team", related_name="solved_problem"
+    )
     problem: fields.ForeignKeyRelation[Problem] = fields.ForeignKeyField(
         "models.Problem"
     )
@@ -64,5 +68,5 @@ class Meta:
         unique_together = (("team", "hint"),)
 
 
-Problem_Pydantic = pydantic_model_creator(Problem)
+BaseProblem_Pydantic = pydantic_model_creator(BaseProblem)
 Hint_Pydantic = pydantic_model_creator(Hint)

src/pwncore/models/pre_event.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from tortoise.models import Model
+from tortoise import fields
+
+from pwncore.models.ctf import BaseProblem
+
+__all__ = (
+    "PreEventSolvedProblem",
+    "PreEventProblem",
+)
+
+
+class PreEventProblem(BaseProblem):
+    flag = fields.TextField()
+    url = fields.TextField()  # Link to CTF file/repo/website
+
+
+class PreEventSolvedProblem(Model):
+    tag = fields.CharField(128)
+    problem: fields.ForeignKeyRelation[PreEventProblem] = fields.ForeignKeyField(
+        "models.PreEventProblem"
+    )
+    solved_at = fields.DatetimeField(auto_now_add=True)
+
+    class Meta:
+        unique_together = (("tag", "problem"),)

src/pwncore/models/user.py

@@ -7,6 +7,7 @@
 from tortoise.contrib.pydantic import pydantic_model_creator
 
 from pwncore.models.container import Container
+from pwncore.config import config
 
 __all__ = (
     "User",
@@ -34,8 +35,11 @@ async def save(self, *args, **kwargs):
         # Reason why we dont use pre_save: overhead, ugly
         if self.team is not None and hasattr(self.team, "members"):
             count = await self.team.members.filter(~Q(id=self.pk)).count()
-            if count >= 3:
-                raise IntegrityError("3 or more users already exist for the team")
+            if count >= config.max_members_per_team:
+                raise IntegrityError(
+                    f"{config.max_members_per_team}"
+                    " or more users already exist for the team"
+                )
         return await super().save(*args, **kwargs)
 
 
@@ -45,6 +49,7 @@ class Team(Model):
     )  # team.id raises Team does not have id, so explicitly adding it
     name = fields.CharField(255, unique=True)
     secret_hash = fields.TextField()
+    coins = fields.IntField(default=0)
 
     members: fields.ReverseRelation[User]
     containers: fields.ReverseRelation[Container]

src/pwncore/routes/__init__.py

@@ -1,6 +1,6 @@
 from fastapi import APIRouter
 
-from pwncore.routes import ctf, team, auth, admin
+from pwncore.routes import ctf, team, auth, admin, leaderboard
 from pwncore.config import config
 
 # Main router (all routes go under /api)
@@ -10,5 +10,6 @@
 router.include_router(auth.router)
 router.include_router(ctf.router)
 router.include_router(team.router)
+router.include_router(leaderboard.router)
 if config.development:
     router.include_router(admin.router)

src/pwncore/routes/admin.py

@@ -1,14 +1,52 @@
+import logging
 from fastapi import APIRouter
+from passlib.hash import bcrypt
 
-from pwncore.models import Team, Problem, Hint, User
+from pwncore.models import (
+    Team,
+    Problem,
+    Hint,
+    User,
+    PreEventSolvedProblem,
+    PreEventProblem,
+)
+from pwncore.config import config
+from pwncore.models.ctf import SolvedProblem
 
 metadata = {
     "name": "admin",
     "description": "Admin routes (currently only running when development on)",
 }
 
+# TODO: Make this protected
 router = APIRouter(prefix="/admin", tags=["admin"])
 
+if config.development:
+    logging.basicConfig(level=logging.INFO)
+
+
+@router.get("/union")
+async def calculate_team_coins():  # Inefficient, anyways will be used only once
+    logging.info("Calculating team points form pre-event CTFs:")
+    team_ids = await Team.filter().values_list("id", flat=True)
+    for team_id in team_ids:
+        member_tags = await User.filter(team_id=team_id).values_list("tag", flat=True)
+
+        if not member_tags:
+            return 0
+
+        problems_solved = set(
+            await PreEventSolvedProblem.filter(tag__in=member_tags).values_list(
+                "problem_id", flat=True
+            )
+        )
+
+        team = await Team.get(id=team_id)
+        for ctf_id in problems_solved:
+            team.coins += (await PreEventProblem.get(id=ctf_id)).points
+        logging.info(f"{team.id}) {team.name}: {team.coins}")
+        await team.save()
+
 
 @router.get("/create")
 async def init_db():
@@ -20,6 +58,22 @@ async def init_db():
         image_name="key:latest",
         image_config={"PortBindings": {"22/tcp": [{}]}},
     )
+    await PreEventProblem.create(
+        name="Static_test",
+        description="Chod de tujhe se na ho paye",
+        author="Meetesh Saini",
+        points=20,
+        flag="asd",
+        url="lugvitc.org",
+    )
+    await PreEventProblem.create(
+        name="New Static Test",
+        description="AJJSBFISHDBFHSD",
+        author="Meetesh Saini",
+        points=21,
+        flag="asdf",
+        url="lugvitc.org",
+    )
     await Problem.create(
         name="In-Plain-Sight",
         description="A curious image with hidden secrets?",
@@ -28,8 +82,24 @@ async def init_db():
         image_name="key:latest",
         image_config={"PortBindings": {"22/tcp": [{}]}},
     )
-    await Team.create(name="CID Squad", secret_hash="veryverysecret")
-    await Team.create(name="Triple A battery", secret_hash="chotiwali")
+    await Problem.create(
+        name="GitGood",
+        description="How to master the art of solving CTFs? Git good nub.",
+        author="Aadivishnu and Shoubhit",
+        points=300,
+        image_name="test:latest",
+        image_config={"PortBindings": {"22/tcp": [{}], "5000/tcp": [{}]}},
+    )
+    await Team.create(name="CID Squad", secret_hash=bcrypt.hash("veryverysecret"))
+    await Team.create(
+        name="Triple A battery", secret_hash=bcrypt.hash("chotiwali"), coins=20
+    )
+    await PreEventSolvedProblem.create(tag="23BCE1000", problem_id="1")
+    await PreEventSolvedProblem.create(tag="23BRS1000", problem_id="1")
+    # await PreEventSolvedProblem.create(
+    #     tag="23BAI1000",
+    #     problem_id="2"
+    # )
     await User.create(
         tag="23BRS1000",
         name="abc",
@@ -77,3 +147,6 @@ async def init_db():
     await Hint.create(order=2, problem_id=1, text="This is the third hint")
     await Hint.create(order=0, problem_id=2, text="This is the first hint")
     await Hint.create(order=1, problem_id=2, text="This is the second hint")
+    await SolvedProblem.create(team_id=2, problem_id=1)
+    await SolvedProblem.create(team_id=2, problem_id=2)
+    await SolvedProblem.create(team_id=1, problem_id=2)

src/pwncore/routes/auth.py

@@ -21,27 +21,20 @@
 router = APIRouter(prefix="/auth", tags=["auth"])
 
 
-class SignupBody(BaseModel):
-    name: str
-    password: str
-    # members: list[str]
-
-
-class LoginBody(BaseModel):
+class AuthBody(BaseModel):
     name: str
     password: str
 
 
 @atomic()
 @router.post("/signup")
-async def signup_team(team: SignupBody, response: Response):
+async def signup_team(team: AuthBody, response: Response):
+    team.name = team.name.strip()
     try:
         if await Team.exists(name=team.name):
             response.status_code = 406
             return {"msg_code": config.msg_codes["team_exists"]}
 
-        # TODO: Add users details
-
         await Team.create(name=team.name, secret_hash=bcrypt.hash(team.password))
     except Exception:
         response.status_code = 500
@@ -50,7 +43,7 @@ async def signup_team(team: SignupBody, response: Response):
 
 
 @router.post("/login")
-async def team_login(team_data: LoginBody, response: Response):
+async def team_login(team_data: AuthBody, response: Response):
     # TODO: Simplified logic since we're not doing refresh tokens.
 
     team = await Team.get_or_none(name=team_data.name)