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]