diff --git a/src/pwncore/__init__.py b/src/pwncore/__init__.py index 91461cd..a869c27 100644 --- a/src/pwncore/__init__.py +++ b/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() diff --git a/src/pwncore/config.py b/src/pwncore/config.py index dbb47fa..fbe56a8 100644 --- a/src/pwncore/config.py +++ b/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, ) diff --git a/src/pwncore/models/__init__.py b/src/pwncore/models/__init__.py index 8e8f78c..4edee09 100644 --- a/src/pwncore/models/__init__.py +++ b/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", ) diff --git a/src/pwncore/models/container.py b/src/pwncore/models/container.py index 47db53b..0153691 100644 --- a/src/pwncore/models/container.py +++ b/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] diff --git a/src/pwncore/models/ctf.py b/src/pwncore/models/ctf.py index d37aee7..0798a9f 100644 --- a/src/pwncore/models/ctf.py +++ b/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) diff --git a/src/pwncore/models/user.py b/src/pwncore/models/user.py index ebe5497..5ebebfc 100644 --- a/src/pwncore/models/user.py +++ b/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) diff --git a/src/pwncore/routes/__init__.py b/src/pwncore/routes/__init__.py index 2c6095a..91d36aa 100644 --- a/src/pwncore/routes/__init__.py +++ b/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"} diff --git a/src/pwncore/routes/ctf/__init__.py b/src/pwncore/routes/ctf/__init__.py index db0e506..8f8d342 100644 --- a/src/pwncore/routes/ctf/__init__.py +++ b/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 diff --git a/src/pwncore/routes/ctf/start.py b/src/pwncore/routes/ctf/start.py index 6395d72..6bdd5b4 100644 --- a/src/pwncore/routes/ctf/start.py +++ b/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 diff --git a/src/pwncore/routes/team.py b/src/pwncore/routes/team.py index 05c5b6c..4a890ee 100644 --- a/src/pwncore/routes/team.py +++ b/src/pwncore/routes/team.py @@ -1,6 +1,7 @@ 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"} @@ -8,10 +9,15 @@ 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