Merge pull request #32 from WizzyGeek/dev-round2

Super Hot Fix: Any PRs after this are a disgrace

src/pwncore/__init__.py

@@ -53,6 +53,7 @@ async def app_lifespan(app: FastAPI):
 
 if config.development:
     origins.append("http://localhost:5173")
+    origins.append("http://localhost:4173")
 
 app.add_middleware(
     CORSMiddleware,

src/pwncore/config.py

@@ -67,6 +67,6 @@ class Config:
     jwt_secret="mysecret",
     jwt_valid_duration=12,  # In hours
     msg_codes=msg_codes,
-    hint_penalty=10,
+    hint_penalty=50,
     max_members_per_team=3,
 )

src/pwncore/models/round2.py

@@ -38,7 +38,7 @@ class R2Container(Model):
     ports: fields.ReverseRelation[R2Ports]
 
     class PydanticMeta:
-        exclude = ["docker_id", "flag", "meta_team"]
+        exclude = ["docker_id", "flag", "meta_team", "r2attackrecords"]
 
 
 class R2Ports(Model):

src/pwncore/routes/__init__.py

@@ -1,7 +1,8 @@
 from fastapi import APIRouter
 
 from pwncore.routes import ctf, team, auth, admin, leaderboard, round2
-from pwncore.config import config
+
+# from pwncore.config import config
 
 # Main router (all routes go under /api)
 router = APIRouter(prefix="/api")
@@ -12,5 +13,6 @@
 router.include_router(team.router)
 router.include_router(leaderboard.router)
 router.include_router(round2.router)
-if config.development:
-    router.include_router(admin.router)
+router.include_router(admin.router)
+# if config.development:
+# router.include_router(admin.router)

src/pwncore/routes/admin.py

@@ -4,9 +4,11 @@
 import logging
 import uuid
 
-from fastapi import APIRouter, Response
+from fastapi import APIRouter, Request, Response
 from passlib.hash import bcrypt
-from tortoise.transactions import atomic
+from tortoise.transactions import atomic, in_transaction
+from tortoise.expressions import RawSQL, Q
+from tortoise.functions import Sum
 
 from pwncore.models import (
     Team,
@@ -35,6 +37,22 @@
 if config.development:
     logging.basicConfig(level=logging.INFO)
 
+ADMIN_HASH = "$2b$12$K2LsLGS/Mahksh0V6xZYKOviNEHMv3Of5f1zhyF6CWJ8rJIcKnSqu"
+NAMES = [
+    "Mimas",
+    "Enceladus",
+    "Tethys",
+    "Dione",
+    "Rhea",
+    "Titan",
+    "Hyperion",
+    "Iapetus",
+    "Phoebe",
+    "Janus",
+    "Epimetheus",
+    "Pan",
+]
+
 
 async def _del_cont(id: str):
     container = await docker_client.containers.get(id)
@@ -85,53 +103,104 @@ async def _create_container(prob: R2Problem, mteam: MetaTeam):
 
 
 @router.get("/round2")
-async def round2(response: Response):
+async def round2(response: Response, req: Request):
+    if not bcrypt.verify((await req.body()).strip(), ADMIN_HASH):
+        response.status_code = 401
+        return
+
     containers = await Container.all()
 
-    async with asyncio.TaskGroup() as tg:
-        for container in containers:
-            tg.create_task(_del_cont(container.docker_id))
+    async with in_transaction():
+        async with asyncio.TaskGroup() as tg:
+            for container in containers:
+                tg.create_task(_del_cont(container.docker_id))
 
-    try:
-        await Container.all().delete()
-    except Exception:
-        response.status_code = 500
-        logging.exception("Error while initing round2")
-        return {"msg_code": config.msg_codes["db_error"]}
+        try:
+            await Container.all().delete()
+        except Exception:
+            response.status_code = 500
+            logging.exception("Error while initing round2")
+            return {"msg_code": config.msg_codes["db_error"]}
 
-    problems = await R2Problem.all()
-    mteams = await MetaTeam.all()
+        await MetaTeam.all().delete()
+        await MetaTeam.bulk_create(
+            [MetaTeam(name=NAMES[i], id=i + 1) for i in range(12)]
+        )
 
-    async with asyncio.TaskGroup() as tg:
-        for pm in itertools.product(problems, mteams):
-            tg.create_task(_create_container(*pm))
+        teams = (
+            await Team.all()
+            .filter(Q(solved_problem__problem__visible=True) | Q(points__gte=0))
+            .annotate(
+                tpoints=RawSQL(
+                    'COALESCE((SUM("solvedproblem"."penalty" * '
+                    '"solvedproblem__problem"."points")'
+                    ' + "team"."points"), 0)'
+                )
+            )
+            .annotate(
+                tpoints2=Sum(
+                    RawSQL(
+                        '"solvedproblem"."penalty" * "solvedproblem__problem"."points"'
+                    )
+                )
+            )
+            .order_by("-tpoints")
+        )
+
+        for i in range(12):
+            for team in teams[i::12]:
+                team.meta_team_id = i + 1  # type: ignore[attr-defined]
+                # print(mts[i].pk, mts)
+                await team.save(update_fields=["meta_team_id"])
+
+        # await Team.bulk_update(teams, fields=["meta_team_id"])
+
+        problems = await R2Problem.all()
+        mteams = await MetaTeam.all()
+
+        async with asyncio.TaskGroup() as tg:
+            for pm in itertools.product(problems, mteams):
+                tg.create_task(_create_container(*pm))
 
 
 @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
+async def calculate_team_coins(
+    response: Response, req: Request
+):  # Inefficient, anyways will be used only once
+    if not bcrypt.verify((await req.body()).strip(), ADMIN_HASH):
+        response.status_code = 401
+        return
+    async with in_transaction():
+        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(user_id__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()
+            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():
+async def init_db(
+    response: Response, req: Request
+):  # Inefficient, anyways will be used only once
+    if not bcrypt.verify((await req.body()).strip(), ADMIN_HASH):
+        response.status_code = 401
+        return
     await Problem.create(
         name="Invisible-Incursion",
         description="Chod de tujhe se na ho paye",
@@ -191,10 +260,10 @@ async def init_db():
     await Team.create(
         name="Triple A battery", secret_hash=bcrypt.hash("chotiwali"), coins=20
     )
-    await PreEventUser.create(tag="23bce1000", email="[email protected]")
-    await PreEventUser.create(tag="23brs1000", email="[email protected]")
-    await PreEventSolvedProblem.create(user_id="23bce1000", problem_id="1")
-    await PreEventSolvedProblem.create(user_id="23brs1000", problem_id="1")
+    await PreEventUser.create(tag="23BCE1000", email="[email protected]")
+    await PreEventUser.create(tag="23BRS1000", email="[email protected]")
+    await PreEventSolvedProblem.create(user_id="23BCE1000", problem_id="1")
+    await PreEventSolvedProblem.create(user_id="23BRS1000", problem_id="1")
     # await PreEventSolvedProblem.create(
     #     tag="23BAI1000",
     #     problem_id="2"

src/pwncore/routes/ctf/__init__.py

@@ -56,15 +56,27 @@ class Flag(BaseModel):
 @router.get("/completed")
 async def completed_problem_get(jwt: RequireJwt):
     team_id = jwt["team_id"]
+    ViewedHint.filter(team_id=team_id).annotate()
     problems = await Problem_Pydantic.from_queryset(
         Problem.filter(solvedproblems__team_id=team_id, visible=True)
     )
     return problems
 
 
 @router.get("/list")
-async def ctf_list():
+async def ctf_list(jwt: RequireJwt):
+    team_id = jwt["team_id"]
     problems = await Problem_Pydantic.from_queryset(Problem.filter(visible=True))
+    acc: dict[int, float] = defaultdict(lambda: 1.0)
+    for k, v in map(
+        lambda x: (x.hint.problem_id, HINTPENALTY[x.hint.order]),  # type: ignore[attr-defined]
+        await ViewedHint.filter(team_id=team_id, with_points=True).prefetch_related(
+            "hint"
+        ),
+    ):
+        acc[k] -= v / 100
+    for i in problems:
+        i.points = int(acc[i.id] * i.points)  # type: ignore[attr-defined]
     return problems
 
 

src/pwncore/routes/leaderboard.py

@@ -4,13 +4,12 @@
 
 from fastapi import APIRouter, Request
 
-from tortoise.expressions import RawSQL
+from tortoise.expressions import RawSQL, Q
 
 from pwncore.models import Team
 
 # Metadata at the top for instant accessibility
-metadata = {"name": "leaderboard",
-            "description": "Operations on the leaderboard"}
+metadata = {"name": "leaderboard", "description": "Operations on the leaderboard"}
 
 router = APIRouter(prefix="/leaderboard", tags=["leaderboard"])
 
@@ -28,17 +27,19 @@ def __init__(self, period: float) -> None:
     async def _do_update(self):
         self.data = (
             await Team.all()
-            .filter(solved_problem__problem__visible=True)
+            .filter(Q(solved_problem__problem__visible=True) | Q(points__gte=0))
             .annotate(
                 tpoints=RawSQL(
-                    'SUM("solvedproblem"."penalty" * "solvedproblem__problem"."points")'
-                    ' + "team"."points"'
+                    'COALESCE((SUM("solvedproblem"."penalty" * '
+                    '"solvedproblem__problem"."points")'
+                    ' + "team"."points"), 0)'
                 )
             )
             .group_by("id", "meta_team__name")
             .order_by("-tpoints")
             .values("name", "tpoints", "meta_team__name")
         )
+
         self.last_update = monotonic()
 
     async def get_lb(self, req: Request):

src/pwncore/routes/round2.py

@@ -2,7 +2,7 @@
 
 from fastapi import APIRouter, Response
 import tortoise.exceptions
-from tortoise.expressions import F
+from tortoise.expressions import F, Q, Subquery
 from tortoise.transactions import in_transaction
 
 from pwncore.config import config
@@ -26,9 +26,19 @@
 
 
 @router.get("/list")
-async def r2ctf_list():
+async def r2ctf_list(jwt: RequireJwt):
     return await R2Container_Pydantic.from_queryset(
-        R2Container.all().prefetch_related("problem", "ports")
+        R2Container.all()
+        .filter(
+            ~Q(
+                id__in=Subquery(
+                    R2AttackRecord.filter(
+                        meta_team_id=(await Team.get(id=jwt["team_id"])).meta_team_id  # type: ignore[attr-defined]
+                    ).values("container_id")
+                )
+            )
+        )
+        .prefetch_related("problem", "ports")
     )
 
 

src/pwncore/routes/team.py

@@ -7,6 +7,7 @@
 from pwncore.config import config
 from pwncore.models import Team, User, Team_Pydantic, User_Pydantic, Container
 from pwncore.routes.auth import RequireJwt
+
 # from pwncore.routes.leaderboard import gcache
 
 # Metadata at the top for instant accessibility
@@ -45,14 +46,10 @@ async def team_members(jwt: RequireJwt):
 async def get_self_team(jwt: RequireJwt):
     team_id = jwt["team_id"]
 
-    team_model = await Team.get(id=team_id).prefetch_related("members")
-    team = dict(await Team_Pydantic.from_tortoise_orm(team_model))
-
-    # Get members
-    team["members"] = [
-        await User_Pydantic.from_tortoise_orm(member) for member in team_model.members
-    ]
+    team_model = await Team.get(id=team_id).prefetch_related("meta_team")
 
+    team = dict(await Team_Pydantic.from_tortoise_orm(team_model))
+    team["meta_team__name"] = team_model.meta_team.name  # type: ignore[arg-type, union-attr]
     # Get points from leaderboard
     # would be better is cache stores the values in a dict indexed by team id
     # for leaderboard_team in gcache.data: