Merge pull request #6 from parrothacker1/dev

Add jwt based Team authentication

pyproject.toml

@@ -21,6 +21,9 @@ fastapi = "^0.104.1"
 uvicorn = "^0.24.0"
 tortoise-orm = {version = "^0.20.0", extras = ["asyncpg", "accel"]}
 aiodocker = "^0.21.0"
+passlib = "^1.7.4"
+pyjwt = "^2.8.0"
+types-passlib = "^1.7.7.13"
 
 [tool.poetry.group.dev.dependencies]
 mypy = "^1.6.1"

src/pwncore/config.py

@@ -28,6 +28,10 @@
     "team_not_found": 10,
     "user_not_found": 11,
     "ctf_solved": 12,
+    "signup_success": 13,
+    "wrong_password": 14,
+    "login_success": 15,
+    "team_exists": 17,
 }
 
 
@@ -39,6 +43,8 @@ class Config:
     docker_url: str | None
     flag: str
     max_containers_per_team: int
+    jwt_secret: str
+    jwt_valid_duration: int
 
 
 config = Config(
@@ -47,5 +53,7 @@ class Config:
     docker_url=None,  # None for default system docker
     flag="C0D",
     max_containers_per_team=3,
+    jwt_secret="mysecret",
+    jwt_valid_duration=12,  # In hours
     msg_codes=msg_codes,
 )

src/pwncore/models/user.py

@@ -21,7 +21,7 @@ class User(Model):
     # abstractly just represents one person, expand this
     # field for Identity providers
     tag = fields.CharField(128, unique=True)
-    name = fields.CharField(255)
+    name = fields.TextField()
     email = fields.TextField()
     phone_num = fields.CharField(15)
 
@@ -40,6 +40,9 @@ async def save(self, *args, **kwargs):
 
 
 class Team(Model):
+    id = fields.IntField(
+        pk=True
+    )  # team.id raises Team does not have id, so explicitly adding it
     name = fields.CharField(255, unique=True)
     secret_hash = fields.TextField()
 

src/pwncore/routes/__init__.py

@@ -1,10 +1,14 @@
 from fastapi import APIRouter
 
-from pwncore.routes import ctf, team
+from pwncore.routes import ctf, team, auth, admin
+from pwncore.config import config
 
 # Main router (all routes go under /api)
 router = APIRouter(prefix="/api")
 
 # Include all the subroutes
+router.include_router(auth.router)
 router.include_router(ctf.router)
 router.include_router(team.router)
+if config.development:
+    router.include_router(admin.router)

src/pwncore/routes/admin.py

@@ -0,0 +1,79 @@
+from fastapi import APIRouter
+
+from pwncore.models import Team, Problem, Hint, User
+
+metadata = {
+    "name": "admin",
+    "description": "Admin routes (currently only running when development on)",
+}
+
+router = APIRouter(prefix="/admin", tags=["admin"])
+
+
+@router.get("/create")
+async def init_db():
+    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 Problem.create(
+        name="In-Plain-Sight",
+        description="A curious image with hidden secrets?",
+        author="KreativeThinker",
+        points=300,
+        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 User.create(
+        tag="23BRS1000",
+        name="abc",
+        team_id=2,
+        phone_num=1111111111,
+        email="[email protected]",
+    )
+    await User.create(
+        tag="23BCE1000",
+        name="def",
+        team_id=2,
+        phone_num=2222222222,
+        email="[email protected]",
+    )
+    await User.create(
+        tag="23BAI1000",
+        name="ghi",
+        team_id=2,
+        phone_num=3333333333,
+        email="[email protected]",
+    )
+    await User.create(
+        tag="23BRS2000",
+        name="ABC",
+        team_id=1,
+        phone_num=4444444444,
+        email="[email protected]",
+    )
+    await User.create(
+        tag="23BCE2000",
+        name="DEF",
+        team_id=1,
+        phone_num=5555555555,
+        email="[email protected]",
+    )
+    await User.create(
+        tag="23BAI2000",
+        name="GHI",
+        team_id=1,
+        phone_num=6666666666,
+        email="[email protected]",
+    )
+    await Hint.create(order=0, problem_id=1, text="This is the first hint")
+    await Hint.create(order=1, problem_id=1, text="This is the second hint")
+    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")

src/pwncore/routes/auth.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+import datetime
+import typing as t
+
+import jwt
+from fastapi import APIRouter, Header, Response, HTTPException, Depends
+from passlib.hash import bcrypt
+from pydantic import BaseModel
+from tortoise.transactions import atomic
+
+from pwncore.models import Team
+from pwncore.config import config
+
+# Metadata at the top for instant accessibility
+metadata = {
+    "name": "auth",
+    "description": "Authentication using a JWT using a single access token.",
+}
+
+router = APIRouter(prefix="/auth", tags=["auth"])
+
+
+class SignupBody(BaseModel):
+    name: str
+    password: str
+    # members: list[str]
+
+
+class LoginBody(BaseModel):
+    name: str
+    password: str
+
+
+@atomic()
+@router.post("/signup")
+async def signup_team(team: SignupBody, response: Response):
+    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
+        return {"msg_code": config.msg_codes["db_error"]}
+    return {"msg_code": config.msg_codes["signup_success"]}
+
+
+@router.post("/login")
+async def team_login(team_data: LoginBody, response: Response):
+    # TODO: Simplified logic since we're not doing refresh tokens.
+
+    team = await Team.get_or_none(name=team_data.name)
+    if team is None:
+        response.status_code = 404
+        return {"msg_code": config.msg_codes["team_not_found"]}
+    if not bcrypt.verify(team_data.password, team.secret_hash):
+        response.status_code = 401
+        return {"msg_code": config.msg_codes["wrong_password"]}
+
+    current_time = datetime.datetime.utcnow()
+    expiration_time = current_time + datetime.timedelta(hours=config.jwt_valid_duration)
+    token_payload = {"team_id": team.id, "exp": expiration_time}
+    token = jwt.encode(token_payload, config.jwt_secret, algorithm="HS256")
+
+    # Returning token to be sent as an authorization header "Bearer <TOKEN>"
+    return {
+        "msg_code": config.msg_codes["login_success"],
+        "access_token": token,
+        "token_type": "bearer",
+    }
+
+
+# Custom JWT processing (since FastAPI's implentation deals with refresh tokens)
+# Supressing B008 in order to be able to use Header() in arguments
+def get_jwt(*, authorization: t.Annotated[str, Header()]) -> JwtInfo:
+    try:
+        token = authorization.split(" ")[1]  # Remove Bearer
+        decoded_token: JwtInfo = jwt.decode(
+            token, config.jwt_secret, algorithms=["HS256"]
+        )
+    except Exception:  # Will filter for invalid signature/expired tokens
+        raise HTTPException(status_code=401)
+    return decoded_token
+
+
+# Using a pre-assigned variable everywhere else to follow flake8's B008
+JwtInfo = t.TypedDict("JwtInfo", {"team_id": int, "exp": int})
+RequireJwt = t.Annotated[JwtInfo, Depends(get_jwt)]