Merge pull request #3 from mradigen/lugvitc-master

Please don't mess anything up

pyproject.toml

@@ -19,7 +19,7 @@ from = "src"
 python = "^3.11"
 fastapi = "^0.104.1"
 uvicorn = "^0.24.0"
-tortoise-orm = "^0.20.0"
+tortoise-orm = {version = "^0.20.0", extras = ["asyncpg", "accel"]}
 aiodocker = "^0.21.0"
 
 [tool.poetry.group.dev.dependencies]

src/pwncore/models/__init__.py

@@ -0,0 +1,19 @@
+"""pwncore.models
+
+Contains all Pydantic and Tortoise ORM models
+"""
+
+from pwncore.models.container import Container, Ports
+from pwncore.models.ctf import Problem, SolvedProblem, Hint, ViewedHint
+from pwncore.models.user import User, Team
+
+__all__ = (
+    "Container",
+    "Ports",
+    "Problem",
+    "SolvedProblem",
+    "Hint",
+    "ViewedHint",
+    "User",
+    "Team",
+)

src/pwncore/models/container.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+import typing as t
+
+from tortoise.models import Model
+from tortoise import fields
+
+if t.TYPE_CHECKING:
+    from pwncore.models.ctf import Problem
+    from pwncore.models.user import Team
+
+__all__ = ("Container", "Ports")
+
+
+# Note: These are all type annotated, dont worry
+class Container(Model):
+    docker_id = fields.CharField(128, unique=True)
+    problem: fields.ForeignKeyRelation[Problem] = fields.ForeignKeyField(
+        "models.Problem", on_delete=fields.OnDelete.NO_ACTION
+    )
+    team: fields.ForeignKeyRelation[Team] = fields.ForeignKeyField(
+        "models.Team", related_name="containers"
+    )
+    flag = fields.TextField()
+
+    ports: fields.ReverseRelation[Ports]
+
+
+class Ports(Model):
+    # FUTURE PROOFING: ADD domain
+    container: fields.ForeignKeyRelation[Container] = fields.ForeignKeyField(
+        "models.Container", related_name="ports", on_delete=fields.OnDelete.CASCADE
+    )
+    port = fields.IntField(pk=True)

src/pwncore/models/ctf.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from tortoise.models import Model
+from tortoise import fields
+
+if TYPE_CHECKING:
+    from tortoise.fields import Field
+    from pwncore.models.user import Team
+
+__all__ = ("Problem", "Hint", "SolvedProblem", "ViewedHint")
+
+
+class Problem(Model):
+    name = fields.TextField()
+    description = fields.TextField()
+    points = fields.IntField()
+    author = fields.TextField()
+
+    image_name = fields.TextField()
+    image_config: Field[dict[str, list]] = fields.JSONField()  # type: ignore[assignment]
+
+    hints: fields.ReverseRelation[Hint]
+
+
+class Hint(Model):
+    order = fields.SmallIntField()  # 0, 1, 2
+    problem: fields.ForeignKeyRelation[Problem] = fields.ForeignKeyField(
+        "models.Problem", related_name="hints"
+    )
+    text = fields.TextField()
+
+    class Meta:
+        ordering = ("order",)
+
+
+class SolvedProblem(Model):
+    team: fields.ForeignKeyRelation[Team] = fields.ForeignKeyField("models.Team")
+    problem: fields.ForeignKeyRelation[Problem] = fields.ForeignKeyField(
+        "models.Problem"
+    )
+    solved_at = fields.DatetimeField(auto_now_add=True)
+
+    class Meta:
+        unique_together = (("team", "problem"),)
+
+
+class ViewedHint(Model):
+    team: fields.ForeignKeyRelation[Team] = fields.ForeignKeyField("models.Team")
+    hint: fields.ForeignKeyRelation[Hint] = fields.ForeignKeyField("models.Hint")
+
+    class Meta:
+        unique_together = (("team", "hint"),)

src/pwncore/models/user.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from tortoise import fields
+from tortoise.exceptions import IntegrityError
+from tortoise.models import Model
+from tortoise.expressions import Q
+
+if TYPE_CHECKING:
+    from pwncore.models.container import Container
+
+__all__ = ("User", "Team")
+
+
+class User(Model):
+    # Registration numbers and other identity tags
+    # abstractly just represents one person, expand this
+    # field for Identity providers
+    tag = fields.CharField(128, unique=True)
+    name = fields.CharField(255)
+    email = fields.TextField()
+    phone_num = fields.CharField(15)
+
+    team: fields.ForeignKeyNullableRelation[Team] = fields.ForeignKeyField(
+        "models.Team", "members", null=True, on_delete=fields.OnDelete.SET_NULL
+    )
+
+    async def save(self, *args, **kwargs):
+        # TODO: Insert/Update in one query
+        # Reason why we dont use pre_save: overhead, ugly
+        if self.team is not None:
+            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")
+        return await super().save(*args, **kwargs)
+
+
+class Team(Model):
+    name = fields.CharField(255, unique=True)
+    secret_hash = fields.TextField()
+
+    members: fields.ReverseRelation[User]
+    containers: fields.ReverseRelation[Container]