diff --git a/src/pwncore/config.py b/src/pwncore/config.py index af36cae..cdee561 100644 --- a/src/pwncore/config.py +++ b/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, ) diff --git a/src/pwncore/models/__init__.py b/src/pwncore/models/__init__.py index 4edee09..b38cbf8 100644 --- a/src/pwncore/models/__init__.py +++ b/src/pwncore/models/__init__.py @@ -9,14 +9,18 @@ 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", @@ -24,6 +28,8 @@ "Container", "Ports", "User", + "PreEventSolvedProblem", + "PreEventProblem", "User_Pydantic", "Team", "Team_Pydantic", diff --git a/src/pwncore/models/ctf.py b/src/pwncore/models/ctf.py index 0798a9f..8b5f0da 100644 --- a/src/pwncore/models/ctf.py +++ b/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) diff --git a/src/pwncore/models/pre_event.py b/src/pwncore/models/pre_event.py new file mode 100644 index 0000000..50e4300 --- /dev/null +++ b/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"),) diff --git a/src/pwncore/models/user.py b/src/pwncore/models/user.py index 793e1bc..993fb65 100644 --- a/src/pwncore/models/user.py +++ b/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] diff --git a/src/pwncore/routes/__init__.py b/src/pwncore/routes/__init__.py index a12543c..1405efc 100644 --- a/src/pwncore/routes/__init__.py +++ b/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) diff --git a/src/pwncore/routes/admin.py b/src/pwncore/routes/admin.py index 9f4dc65..27cf541 100644 --- a/src/pwncore/routes/admin.py +++ b/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) diff --git a/src/pwncore/routes/auth.py b/src/pwncore/routes/auth.py index baa5cc0..8b37098 100644 --- a/src/pwncore/routes/auth.py +++ b/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) diff --git a/src/pwncore/routes/ctf/__init__.py b/src/pwncore/routes/ctf/__init__.py index 6410db5..4603c51 100644 --- a/src/pwncore/routes/ctf/__init__.py +++ b/src/pwncore/routes/ctf/__init__.py @@ -2,6 +2,7 @@ from fastapi import APIRouter, Response from pydantic import BaseModel +from tortoise.transactions import atomic from pwncore.models import ( Problem, @@ -9,11 +10,13 @@ Container, Hint, ViewedHint, - Problem_Pydantic, + BaseProblem_Pydantic, Hint_Pydantic, + Team, ) from pwncore.config import config from pwncore.routes.ctf.start import router as start_router +from pwncore.routes.ctf.pre_event import router as pre_event_router from pwncore.routes.auth import RequireJwt # Metadata at the top for instant accessibility @@ -25,6 +28,7 @@ router = APIRouter(prefix="/ctf", tags=["ctf"]) router.include_router(start_router) +router.include_router(pre_event_router) class Flag(BaseModel): @@ -33,14 +37,15 @@ class Flag(BaseModel): @router.get("/list") async def ctf_list(): - problems = await Problem_Pydantic.from_queryset(Problem.all()) + problems = await BaseProblem_Pydantic.from_queryset(Problem.all()) return problems +@atomic() @router.post("/{ctf_id}/flag") async def flag_post(ctf_id: int, flag: Flag, response: Response, jwt: RequireJwt): team_id = jwt["team_id"] - problem = await Problem.exists(id=ctf_id) + problem = await Problem.get_or_none(id=ctf_id) if not problem: response.status_code = 404 return {"msg_code": config.msg_codes["ctf_not_found"]} @@ -59,6 +64,7 @@ async def flag_post(ctf_id: int, flag: Flag, response: Response, jwt: RequireJwt return {"status": False} +@atomic() @router.get("/{ctf_id}/hint") async def hint_get(ctf_id: int, response: Response, jwt: RequireJwt): team_id = jwt["team_id"] @@ -67,6 +73,10 @@ async def hint_get(ctf_id: int, response: Response, jwt: RequireJwt): response.status_code = 404 return {"msg_code": config.msg_codes["ctf_not_found"]} + team = await Team.get(id=team_id) + if team.coins < config.hint_penalty: + return {"msg_code": config.msg_codes["insufficient_coins"]} + viewed_hints = ( await Hint.filter(problem_id=ctf_id, viewedhints__team_id=team_id) .order_by("-order") @@ -82,6 +92,9 @@ async def hint_get(ctf_id: int, response: Response, jwt: RequireJwt): else: hint = await Hint.get(problem_id=ctf_id, order=0) + team.coins -= config.hint_penalty + await team.save() + await ViewedHint.create(hint_id=hint.id, team_id=team_id) return {"text": hint.text, "order": hint.order} @@ -98,7 +111,7 @@ async def viewed_problem_hints_get(ctf_id: int, jwt: RequireJwt): @router.get("/completed") async def completed_problem_get(jwt: RequireJwt): team_id = jwt["team_id"] - problems = await Problem_Pydantic.from_queryset( + problems = await BaseProblem_Pydantic.from_queryset( Problem.filter(solvedproblems__team_id=team_id) ) return problems @@ -106,7 +119,7 @@ async def completed_problem_get(jwt: RequireJwt): @router.get("/{ctf_id}") async def ctf_get(ctf_id: int, response: Response): - problem = await Problem_Pydantic.from_queryset(Problem.filter(id=ctf_id)) + problem = await BaseProblem_Pydantic.from_queryset(Problem.filter(id=ctf_id)) if not problem: response.status_code = 404 return {"msg_code": config.msg_codes["ctf_not_found"]} diff --git a/src/pwncore/routes/ctf/pre_event.py b/src/pwncore/routes/ctf/pre_event.py new file mode 100644 index 0000000..e0a350c --- /dev/null +++ b/src/pwncore/routes/ctf/pre_event.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +from fastapi import APIRouter, Response +from pydantic import BaseModel +from tortoise.transactions import atomic + +from pwncore.models import ( + PreEventProblem, + PreEventSolvedProblem, + BaseProblem_Pydantic, +) +from pwncore.config import config + +router = APIRouter(prefix="/pre", tags=["ctf"]) + + +class PreEventFlag(BaseModel): + tag: str + flag: str + + +@router.get("/list") +async def ctf_list(): + problems = await BaseProblem_Pydantic.from_queryset(PreEventProblem.all()) + return problems + + +@atomic() +@router.post("/{ctf_id}/flag") +async def pre_event_flag_post(ctf_id: int, post_body: PreEventFlag, response: Response): + problem = await PreEventProblem.get_or_none(id=ctf_id) + if not problem: + response.status_code = 404 + return {"msg_code": config.msg_codes["ctf_not_found"]} + + user_tag = post_body.tag.strip().casefold() + + if await PreEventSolvedProblem.exists(tag=user_tag, problem_id=ctf_id): + response.status_code = 401 + return {"msg_code": config.msg_codes["ctf_solved"]} + + if problem.flag == post_body.flag: + await PreEventSolvedProblem.create(tag=user_tag, problem_id=ctf_id) + + return {"status": True} + return {"status": False} + + +@router.get("/{ctf_id}") +async def ctf_get(ctf_id: int, response: Response): + problem = await BaseProblem_Pydantic.from_queryset( + PreEventProblem.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 1d76ee6..23ad2f9 100644 --- a/src/pwncore/routes/ctf/start.py +++ b/src/pwncore/routes/ctf/start.py @@ -4,7 +4,7 @@ import uuid from tortoise.transactions import atomic -from pwncore.models import Problem, Container, Ports, Team +from pwncore.models import Problem, Container, Ports from pwncore.container import docker_client from pwncore.config import config from pwncore.routes.auth import RequireJwt @@ -25,18 +25,6 @@ async def start_docker_container(ctf_id: int, response: Response, jwt: RequireJw } } """ - if config.development: - await Problem.create( - name="Invisible-Incursion", - description="Chod de tujhe se na ho paye", - author="Meetesh Saini", - points=300, - image_name="key:latest", - image_config={"PortBindings": {"22/tcp": [{}]}}, - ) - await Team.create( - name="CID Squad" + uuid.uuid4().hex, secret_hash="veryverysecret" - ) ctf = await Problem.get_or_none(id=ctf_id) if not ctf: diff --git a/src/pwncore/routes/leaderboard.py b/src/pwncore/routes/leaderboard.py index e69de29..7a4f4b5 100644 --- a/src/pwncore/routes/leaderboard.py +++ b/src/pwncore/routes/leaderboard.py @@ -0,0 +1,21 @@ +from __future__ import annotations +from fastapi import APIRouter +from tortoise.functions import Sum + +from pwncore.models import Team + +# Metadata at the top for instant accessibility +metadata = {"name": "leaderboard", "description": "Operations on the leaderboard"} + +router = APIRouter(prefix="/leaderboard", tags=["leaderboard"]) + + +@router.get("") +async def fetch_leaderboard(): + leaderboard = dict( + await Team.all() + .annotate(team_points=Sum("solved_problem__problem__points")) + .values_list("name", "team_points") + ) + # points = sorted(points, key=lambda x: x[0], reverse=True) + return leaderboard diff --git a/src/pwncore/routes/team.py b/src/pwncore/routes/team.py index a6f3640..030b09b 100644 --- a/src/pwncore/routes/team.py +++ b/src/pwncore/routes/team.py @@ -1,6 +1,10 @@ from __future__ import annotations -from fastapi import APIRouter +from fastapi import APIRouter, Response +from pydantic import BaseModel +from tortoise.transactions import atomic + +from pwncore.config import config from pwncore.models import Team, User, Team_Pydantic, User_Pydantic from pwncore.routes.auth import RequireJwt @@ -10,6 +14,17 @@ router = APIRouter(prefix="/team", tags=["team"]) +class UserAddBody(BaseModel): + tag: str + name: str + email: str + phone_num: str + + +class UserRemoveBody(BaseModel): + tag: str + + @router.get("/list") async def team_list(): teams = await Team_Pydantic.from_queryset(Team.all()) @@ -23,3 +38,46 @@ async def team_members(jwt: RequireJwt): members = await User_Pydantic.from_queryset(User.filter(team_id=team_id)) # Incase of no members, it just returns an empty list. return members + + +@atomic() +@router.post("/add") +async def add_member(user: UserAddBody, response: Response, jwt: RequireJwt): + team_id = jwt["team_id"] + + if await User.get_or_none(tag=user.tag): + response.status_code = 403 + return {"msg_code": config.msg_codes["user_already_in_team"]} + + try: + await User.create( + # Validation for user tag (reg. no. in our case) + # needs to be done on frontend to not make the server event specific + tag=user.tag.strip().casefold(), + name=user.name, + email=user.email, + phone_num=user.phone_num, + team_id=team_id, + ) + except Exception: + response.status_code = 500 + return {"msg_code": config.msg_codes["db_error"]} + return {"msg_code": config.msg_codes["user_added"]} + + +@atomic() +@router.post("/remove") +async def remove_member(user_info: UserRemoveBody, response: Response, jwt: RequireJwt): + team_id = jwt["team_id"] + + user = await User.get_or_none(team_id=team_id, tag=user_info.tag) + if not user: + response.status_code = 403 + return {"msg_code": config.msg_codes["user_not_in_team"]} + + try: + await user.delete() + except Exception: + response.status_code = 500 + return {"msg_code": config.msg_codes["db_error"]} + return {"msg_code": config.msg_codes["user_removed"]}