Merge pull request #10 from KreativeThinker/master

Hints, Problems and Flag functions

Merging this before #6 since resolving potential conflicts with #6 is easier given it has only 8 commits.

src/pwncore/__init__.py

@@ -5,6 +5,7 @@
 
 import pwncore.docs as docs
 import pwncore.routes as routes
+
 from pwncore.container import docker_client
 from pwncore.config import config
 from pwncore.models import Container
@@ -22,9 +23,14 @@ async def app_lifespan(app: FastAPI):
     containers = await Container.all().values()
     await Container.all().delete()
     for db_container in containers:
-        container = await docker_client.containers.get(db_container["docker_id"])
-        await container.stop()
-        await container.delete()
+        try:
+            container = await docker_client.containers.get(db_container["docker_id"])
+            await container.stop()
+            await container.delete()
+        except (
+            Exception
+        ):  # Raises DockerError if container does not exist, just pass for now.
+            pass
 
     # close_connections is deprecated, not sure how to use connections.close_all()
     await Tortoise.close_connections()

src/pwncore/config.py

@@ -24,6 +24,10 @@
     "container_not_found": 6,
     "container_already_running": 7,
     "container_limit_reached": 8,
+    "hint_limit_reached": 9,
+    "team_not_found": 10,
+    "user_not_found": 11,
+    "ctf_solved": 12,
 }
 
 
@@ -43,15 +47,5 @@ class Config:
     docker_url=None,  # None for default system docker
     flag="C0D",
     max_containers_per_team=3,
-    msg_codes={
-        "db_error": 0,
-        "port_limit_reached": 1,
-        "ctf_not_found": 2,
-        "container_start": 3,
-        "container_stop": 4,
-        "containers_team_stop": 5,
-        "container_not_found": 6,
-        "container_already_running": 7,
-        "container_limit_reached": 8,
-    },
+    msg_codes=msg_codes,
 )

src/pwncore/models/__init__.py

@@ -4,16 +4,27 @@
 """
 
 from pwncore.models.container import Container, Ports
-from pwncore.models.ctf import Problem, SolvedProblem, Hint, ViewedHint
-from pwncore.models.user import User, Team
+from pwncore.models.ctf import (
+    Problem,
+    SolvedProblem,
+    Hint,
+    ViewedHint,
+    Problem_Pydantic,
+    Hint_Pydantic,
+)
+from pwncore.models.user import User, Team, Team_Pydantic, User_Pydantic
 
 __all__ = (
-    "Container",
-    "Ports",
     "Problem",
-    "SolvedProblem",
+    "Problem_Pydantic",
     "Hint",
+    "Hint_Pydantic",
+    "SolvedProblem",
     "ViewedHint",
+    "Container",
+    "Ports",
     "User",
+    "User_Pydantic",
     "Team",
+    "Team_Pydantic",
 )

src/pwncore/models/container.py

@@ -18,9 +18,7 @@ class Container(Model):
     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"
-    )
+    team: fields.ForeignKeyRelation[Team] = fields.ForeignKeyField("models.Team")
     flag = fields.TextField()
 
     ports: fields.ReverseRelation[Ports]

src/pwncore/models/ctf.py

@@ -1,15 +1,19 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from tortoise.models import Model
 from tortoise import fields
+from tortoise.contrib.pydantic import pydantic_model_creator
 
-if TYPE_CHECKING:
-    from tortoise.fields import Field
-    from pwncore.models.user import Team
+from pwncore.models.user import Team
 
-__all__ = ("Problem", "Hint", "SolvedProblem", "ViewedHint")
+__all__ = (
+    "Problem",
+    "Hint",
+    "SolvedProblem",
+    "ViewedHint",
+    "Problem_Pydantic",
+    "Hint_Pydantic",
+)
 
 
 class Problem(Model):
@@ -19,15 +23,19 @@ class Problem(Model):
     author = fields.TextField()
 
     image_name = fields.TextField()
-    image_config: Field[dict[str, list]] = fields.JSONField()  # type: ignore[assignment]
+    image_config: fields.Field[dict[str, list]] = fields.JSONField()  # type: ignore[assignment]
 
     hints: fields.ReverseRelation[Hint]
 
+    class PydanticMeta:
+        exclude = ["image_name", "image_config"]
+
 
 class Hint(Model):
+    id = fields.IntField(pk=True)
     order = fields.SmallIntField()  # 0, 1, 2
     problem: fields.ForeignKeyRelation[Problem] = fields.ForeignKeyField(
-        "models.Problem", related_name="hints"
+        "models.Problem"
     )
     text = fields.TextField()
 
@@ -48,7 +56,13 @@ class Meta:
 
 class ViewedHint(Model):
     team: fields.ForeignKeyRelation[Team] = fields.ForeignKeyField("models.Team")
-    hint: fields.ForeignKeyRelation[Hint] = fields.ForeignKeyField("models.Hint")
+    hint: fields.ForeignKeyRelation[Hint] = fields.ForeignKeyField(
+        "models.Hint",
+    )
 
     class Meta:
         unique_together = (("team", "hint"),)
+
+
+Problem_Pydantic = pydantic_model_creator(Problem)
+Hint_Pydantic = pydantic_model_creator(Hint)

src/pwncore/models/user.py

@@ -1,16 +1,19 @@
 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
+from tortoise.contrib.pydantic import pydantic_model_creator
 
-if TYPE_CHECKING:
-    from pwncore.models.container import Container
+from pwncore.models.container import Container
 
-__all__ = ("User", "Team")
+__all__ = (
+    "User",
+    "Team",
+    "User_Pydantic",
+    "Team_Pydantic",
+)
 
 
 class User(Model):
@@ -29,7 +32,7 @@ class User(Model):
     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:
+        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")
@@ -42,3 +45,10 @@ class Team(Model):
 
     members: fields.ReverseRelation[User]
     containers: fields.ReverseRelation[Container]
+
+    class PydanticMeta:
+        exclude = ["secret_hash"]
+
+
+Team_Pydantic = pydantic_model_creator(Team)
+User_Pydantic = pydantic_model_creator(User)

src/pwncore/routes/__init__.py

@@ -8,9 +8,3 @@
 # Include all the subroutes
 router.include_router(ctf.router)
 router.include_router(team.router)
-
-
-# Miscellaneous routes
-@router.get("/asd")
-async def a():
-    return {"ASD": "asd"}

src/pwncore/routes/ctf/__init__.py

@@ -1,4 +1,18 @@
-from fastapi import APIRouter
+from __future__ import annotations
+
+from fastapi import APIRouter, Response
+
+from pwncore.models import (
+    Problem,
+    SolvedProblem,
+    Container,
+    Hint,
+    ViewedHint,
+    Problem_Pydantic,
+    Hint_Pydantic,
+)
+from pwncore.config import config
+from pwncore.routes.team import get_team_id
 from pwncore.routes.ctf.start import router as start_router
 
 # Metadata at the top for instant accessibility
@@ -10,19 +24,89 @@
 
 router = APIRouter(prefix="/ctf", tags=["ctf"])
 router.include_router(start_router)
+
 # Routes that do not need a separate submodule for themselves
 
 
+# For testing purposes. Replace function with POST method
+def get_flag():
+    return "pwncore{flag_1}"
+
+
 @router.get("/list")
 async def ctf_list():
-    # Get list of ctfs
-    return [
-        {"name": "Password Juggling", "ctf_id": 2243},
-        {"name": "hexane", "ctf_id": 2242},
-    ]
+    problems = await Problem_Pydantic.from_queryset(Problem.all())
+    return problems
+
+
+# For testing purposes only. flag to be passed in body of POST function.
+@router.get("/{ctf_id}/flag")
+async def flag_get(ctf_id: int, response: Response):
+    problem = await Problem.exists(id=ctf_id)
+    if not problem:
+        response.status_code = 404
+        return {"msg_code": config.msg_codes["ctf_not_found"]}
+
+    status = await SolvedProblem.exists(team_id=get_team_id(), problem_id=ctf_id)
+    if status:
+        response.status_code = 401
+        return {"msg_code": config.msg_codes["ctf_solved"]}
+
+    check_solved = await Container.exists(
+        team_id=get_team_id(), flag=get_flag(), problem_id=ctf_id
+    )
+    if check_solved:
+        await SolvedProblem.create(team_id=get_team_id(), problem_id=ctf_id)
+        return {"status": True}
+    return {"status": False}
+
+
+@router.get("/{ctf_id}/hint")
+async def hint_get(ctf_id: int, response: Response):
+    problem = await Problem.exists(id=ctf_id)
+    if not problem:
+        response.status_code = 404
+        return {"msg_code": config.msg_codes["ctf_not_found"]}
+
+    viewed_hints = (
+        await Hint.filter(problem_id=ctf_id, viewedhints__team_id=get_team_id())
+        .order_by("-order")
+        .first()
+    )
+    if viewed_hints:
+        if not await Hint.exists(problem_id=ctf_id, order=viewed_hints.order + 1):
+            response.status_code = 403
+            return {"msg_code": config.msg_codes["hint_limit_reached"]}
+
+        hint = await Hint.get(problem_id=ctf_id, order=viewed_hints.order + 1)
+
+    else:
+        hint = await Hint.get(problem_id=ctf_id, order=0)
+
+    await ViewedHint.create(hint_id=hint.id, team_id=get_team_id())
+    return {"text": hint.text, "order": hint.order}
+
+
+@router.get("/{ctf_id}/viewed_hints")
+async def viewed_problem_hints_get(ctf_id: int):
+    viewed_hints = await Hint_Pydantic.from_queryset(
+        Hint.filter(problem_id=ctf_id, viewedhints__team_id=get_team_id())
+    )
+    return viewed_hints
+
+
+@router.get("/completed")
+async def completed_problem_get():
+    problems = await Problem_Pydantic.from_queryset(
+        Problem.filter(solvedproblems__team_id=get_team_id())
+    )
+    return problems
 
 
 @router.get("/{ctf_id}")
-async def ctf_get(ctf_id: int):
-    # Get ctf from ctf_id
-    return {"status": "logged in!"}
+async def ctf_get(ctf_id: int, response: Response):
+    problem = await Problem_Pydantic.from_queryset(Problem.filter(id=ctf_id))
+    if not problem:
+        response.status_code = 404
+        return {"msg_code": config.msg_codes["ctf_not_found"]}
+    return problem

src/pwncore/routes/ctf/start.py

@@ -60,7 +60,7 @@ async def start_docker_container(ctf_id: int, response: Response):
             "ctf_id": ctf_id,
         }
 
-    if await Container.filter(team_id=team_id).count() >= config.max_containers_per_team:  # noqa: B950
+    if (await Container.filter(team_id=team_id).count() >= config.max_containers_per_team):  # noqa: B950
         return {"msg_code": config.msg_codes["container_limit_reached"]}
 
     # Start a new container

src/pwncore/routes/team.py

@@ -1,17 +1,23 @@
 from __future__ import annotations
 
 from fastapi import APIRouter
+from pwncore.models import Team, User, Team_Pydantic, User_Pydantic
 
 # Metadata at the top for instant accessibility
 metadata = {"name": "team", "description": "Operations with teams"}
 
 router = APIRouter(prefix="/team", tags=["team"])
 
 
+# Retrieve team_id from cookies
+def get_team_id():
+    return 1
+
+
 @router.get("/list")
 async def team_list():
-    # Do login verification here
-    return [{"team_name": "CID Squad"}, {"team_name": "Astra"}]
+    teams = await Team_Pydantic.from_queryset(Team.all())
+    return teams
 
 
 @router.get("/login")
@@ -20,7 +26,9 @@ async def team_login():
     return {"status": "logged in!"}
 
 
-@router.get("/members/{team_id}")
-async def team_members(team_id: int):
-    # Get team members from team_id
-    return [{"name": "ABC", "user_id": 3432}, {"name": "DEF", "user_id": 3422}]
+# Unable to test as adding users returns an error
+@router.get("/members")
+async def team_members():
+    members = await User_Pydantic.from_queryset(User.filter(team_id=get_team_id()))
+    # Incase of no members, it just returns an empty list.
+    return members