From 5c87d2ce4b8dda54f6002215240c2a0dd26570c1 Mon Sep 17 00:00:00 2001 From: parrothacker1 Date: Thu, 23 Nov 2023 00:42:20 +0530 Subject: [PATCH 01/13] Half complete --- src/pwncore/config.py | 2 + src/pwncore/models/user.py | 19 ++++++--- src/pwncore/routes/team.py | 87 ++++++++++++++++++++++++++++++++------ tox.ini | 1 + 4 files changed, 90 insertions(+), 19 deletions(-) diff --git a/src/pwncore/config.py b/src/pwncore/config.py index dbb47fa..80ab2bd 100644 --- a/src/pwncore/config.py +++ b/src/pwncore/config.py @@ -35,6 +35,7 @@ class Config: docker_url: str | None flag: str max_containers_per_team: int + jwt_secret : str config = Config( @@ -43,6 +44,7 @@ class Config: docker_url=None, # None for default system docker flag="C0D", max_containers_per_team=3, + jwt_secret='mysecret', msg_codes={ "db_error": 0, "port_limit_reached": 1, diff --git a/src/pwncore/models/user.py b/src/pwncore/models/user.py index ebe5497..6ddd319 100644 --- a/src/pwncore/models/user.py +++ b/src/pwncore/models/user.py @@ -6,6 +6,7 @@ from tortoise.exceptions import IntegrityError from tortoise.models import Model from tortoise.expressions import Q +from passlib.hash import bcrypt if TYPE_CHECKING: from pwncore.models.container import Container @@ -17,10 +18,10 @@ class User(Model): # Registration numbers and other identity tags # abstractly just represents one person, expand this # field for Identity providers - tag = fields.CharField(128, unique=True) - name = fields.CharField(255) - email = fields.TextField() - phone_num = fields.CharField(15) + tag : fields.CharField = fields.CharField(128, unique=True) + name : str = fields.TextField() + email : str = fields.TextField() + phone_num : fields.CharField = fields.CharField(15) team: fields.ForeignKeyNullableRelation[Team] = fields.ForeignKeyField( "models.Team", "members", null=True, on_delete=fields.OnDelete.SET_NULL @@ -37,8 +38,14 @@ async def save(self, *args, **kwargs): class Team(Model): - name = fields.CharField(255, unique=True) - secret_hash = fields.TextField() + name : fields.CharField = fields.CharField(255, unique=True) + password : fields.CharField = fields.CharField(255, unique=True) + current_points : int = fields.IntField(null=True) + current_stage : int = fields.IntField(null=True) + last_timestamp : fields.DatetimeField = fields.DatetimeField(null=True) members: fields.ReverseRelation[User] containers: fields.ReverseRelation[Container] + + async def check_password(self, password : str) -> bool: + return bcrypt.verify(self.password, password) diff --git a/src/pwncore/routes/team.py b/src/pwncore/routes/team.py index 05c5b6c..96badcd 100644 --- a/src/pwncore/routes/team.py +++ b/src/pwncore/routes/team.py @@ -1,26 +1,87 @@ from __future__ import annotations -from fastapi import APIRouter +import jwt +from fastapi import APIRouter, HTTPException, status, Depends +from fastapi.security import OAuth2PasswordBearer +from tortoise.contrib.pydantic import pydantic_model_creator +from pwncore.models import User, Team +from pwncore.config import config +import datetime +from passlib.hash import bcrypt # Metadata at the top for instant accessibility metadata = {"name": "team", "description": "Operations with teams"} router = APIRouter(prefix="/team", tags=["team"]) +Team_pydantic = pydantic_model_creator(Team, name='Team') +TeamIn_pydantic = pydantic_model_creator(Team, name='TeamIn', exclude_readonly=True) -@router.get("/list") -async def team_list(): - # Do login verification here - return [{"team_name": "CID Squad"}, {"team_name": "Astra"}] +oauth2_scheme = OAuth2PasswordBearer(tokenUrl='token') -@router.get("/login") -async def team_login(): - # Do login verification here - return {"status": "logged in!"} +async def get_current_team(token : str = Depends(oauth2_scheme)): + try: + payload = jwt.decode(token, config.jwt_secret, algorithms=['HS256']) + team = await Team.get(id=payload.get('id')) + except: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Invalid team name or password" + ) + return await Team_pydantic.from_tortoise_orm(team) -@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}] + +@router.post('/token') +async def generate_token(team_data : TeamIn_pydantic): + team = await Team.get_or_none(name=team_data.name) + if team is not None and await team.check_password(team_data.password): + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Invalid team name or password" + ) + + team_obj = await Team_pydantic.from_tortoise_orm(team) + expiration_time = datetime.datetime.utcnow() + datetime.timedelta(hours=2) + payload = team_obj.dict() + payload['exp'] = expiration_time + token = jwt.encode(payload, config.jwt_secret, algorithm='HS256') + + return { + "access_token" : token, + "token_type" : "bearer", + "expires_at" : expiration_time + } + + +@router.post('/signup', response_model=Team_pydantic) +async def signup_team(team : TeamIn_pydantic): + team_obj = await Team.create( + name=team.name, + password=bcrypt.hash(team.password) + ) + return await Team_pydantic.from_tortoise_orm(team_obj) + + +@router.post('/login') +async def team_login( + team_data : TeamIn_pydantic, + token : str = Depends(oauth2_scheme), + team : Team_pydantic = Depends(get_current_team) +): + + issued_at = datetime.datetime.utcfromtimestamp(token['iat']) + current_time = datetime.datetime.utcnow() + idle_time = current_time - issued_at + + if idle_time < datetime.timedelta(hours=2): + expiration_time = current_time + datetime.timedelta(hours=2) + token_payload = jwt.decode(token, config.jwt_secret, algorithms=['HS256']) + token_payload['exp'] = expiration_time + token = jwt.encode(token_payload, config.jwt_secret, algorithm='HS256') + + return { + "access_token": token, + "token_type": "bearer" + } diff --git a/tox.ini b/tox.ini index 4b1f0e8..0352b29 100644 --- a/tox.ini +++ b/tox.ini @@ -20,6 +20,7 @@ deps = fastapi httpx pytest + passlib commands = pytest [testenv:type] From a8ca6282a7b1a3eefe20585383e0f352d0f35691 Mon Sep 17 00:00:00 2001 From: parrothacker1 Date: Thu, 23 Nov 2023 19:11:43 +0530 Subject: [PATCH 02/13] Rewrote login function. --- src/pwncore/routes/team.py | 34 ++++++++++++++++++---------------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/src/pwncore/routes/team.py b/src/pwncore/routes/team.py index 96badcd..807b904 100644 --- a/src/pwncore/routes/team.py +++ b/src/pwncore/routes/team.py @@ -18,7 +18,7 @@ TeamIn_pydantic = pydantic_model_creator(Team, name='TeamIn', exclude_readonly=True) -oauth2_scheme = OAuth2PasswordBearer(tokenUrl='token') +oauth2_scheme = OAuth2PasswordBearer(tokenUrl='login') async def get_current_team(token : str = Depends(oauth2_scheme)): @@ -33,7 +33,6 @@ async def get_current_team(token : str = Depends(oauth2_scheme)): return await Team_pydantic.from_tortoise_orm(team) -@router.post('/token') async def generate_token(team_data : TeamIn_pydantic): team = await Team.get_or_none(name=team_data.name) if team is not None and await team.check_password(team_data.password): @@ -71,17 +70,20 @@ async def team_login( team : Team_pydantic = Depends(get_current_team) ): - issued_at = datetime.datetime.utcfromtimestamp(token['iat']) - current_time = datetime.datetime.utcnow() - idle_time = current_time - issued_at - - if idle_time < datetime.timedelta(hours=2): - expiration_time = current_time + datetime.timedelta(hours=2) - token_payload = jwt.decode(token, config.jwt_secret, algorithms=['HS256']) - token_payload['exp'] = expiration_time - token = jwt.encode(token_payload, config.jwt_secret, algorithm='HS256') - - return { - "access_token": token, - "token_type": "bearer" - } + if team_data.name in (jwt.decode(token, config.jwt_secret, algorithms=['HS256'])): + issued_at = datetime.datetime.utcfromtimestamp(token['iat']) + current_time = datetime.datetime.utcnow() + idle_time = current_time - issued_at + + if idle_time < datetime.timedelta(hours=2): + expiration_time = current_time + datetime.timedelta(hours=2) + token_payload = jwt.decode(token, config.jwt_secret, algorithms=['HS256']) + token_payload['exp'] = expiration_time + token = jwt.encode(token_payload, config.jwt_secret, algorithm='HS256') + + return { + "access_token": token, + "token_type": "bearer" + } + else: + return generate_token(team_data) From 04d91d6278c20d07707f592f1e56f834a1bb87b1 Mon Sep 17 00:00:00 2001 From: parrothacker1 Date: Thu, 23 Nov 2023 23:01:45 +0530 Subject: [PATCH 03/13] Authorization shifted to auth.py --- poetry.lock | 19 +++++++- pyproject.toml | 1 + src/pwncore/config.py | 23 ++++------ src/pwncore/models/user.py | 21 +++------ src/pwncore/routes/__init__.py | 9 +--- src/pwncore/routes/auth.py | 81 +++++++++++++++++++++++++++++++++ src/pwncore/routes/ctf/start.py | 22 ++++----- src/pwncore/routes/team.py | 74 ------------------------------ 8 files changed, 127 insertions(+), 123 deletions(-) create mode 100644 src/pwncore/routes/auth.py diff --git a/poetry.lock b/poetry.lock index 0d7051b..8260a2a 100644 --- a/poetry.lock +++ b/poetry.lock @@ -939,6 +939,23 @@ files = [ {file = "packaging-23.2.tar.gz", hash = "sha256:048fb0e9405036518eaaf48a55953c750c11e1a1b68e0dd1a9d62ed0c092cfc5"}, ] +[[package]] +name = "passlib" +version = "1.7.4" +description = "comprehensive password hashing framework supporting over 30 schemes" +optional = false +python-versions = "*" +files = [ + {file = "passlib-1.7.4-py2.py3-none-any.whl", hash = "sha256:aa6bca462b8d8bda89c70b382f0c298a20b5560af6cbfa2dce410c0a2fb669f1"}, + {file = "passlib-1.7.4.tar.gz", hash = "sha256:defd50f72b65c5402ab2c573830a6978e5f202ad0d984793c8dde2c4152ebe04"}, +] + +[package.extras] +argon2 = ["argon2-cffi (>=18.2.0)"] +bcrypt = ["bcrypt (>=3.1.0)"] +build-docs = ["cloud-sptheme (>=1.10.1)", "sphinx (>=1.6)", "sphinxcontrib-fulltoc (>=1.2.0)"] +totp = ["cryptography"] + [[package]] name = "pathspec" version = "0.11.2" @@ -1401,4 +1418,4 @@ multidict = ">=4.0" [metadata] lock-version = "2.0" python-versions = "^3.11" -content-hash = "45ad8bdea6aea6fc1ede32fa3c44d9cdb36e48f02333009833a28d1d71efac68" +content-hash = "ee606dd31a24ced4709bd82d354df9958af1e2f56d8c18ba96af2534907e0d25" diff --git a/pyproject.toml b/pyproject.toml index dcc0832..57cab21 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,6 +21,7 @@ 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" [tool.poetry.group.dev.dependencies] mypy = "^1.6.1" diff --git a/src/pwncore/config.py b/src/pwncore/config.py index 80ab2bd..1bade51 100644 --- a/src/pwncore/config.py +++ b/src/pwncore/config.py @@ -24,6 +24,11 @@ "container_not_found": 6, "container_already_running": 7, "container_limit_reached": 8, + "signup_success": 9, + "wrong_password": 10, + "login_success": 11, + "team_not_found": 12, + "team_exists": 13 } @@ -35,7 +40,8 @@ class Config: docker_url: str | None flag: str max_containers_per_team: int - jwt_secret : str + jwt_secret: str + jwt_valid_duration: int config = Config( @@ -44,16 +50,7 @@ class Config: docker_url=None, # None for default system docker flag="C0D", max_containers_per_team=3, - jwt_secret='mysecret', - 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, - }, + jwt_secret="mysecret", + jwt_valid_duration=12, # In hours + msg_codes=msg_codes, ) diff --git a/src/pwncore/models/user.py b/src/pwncore/models/user.py index 6ddd319..85afad0 100644 --- a/src/pwncore/models/user.py +++ b/src/pwncore/models/user.py @@ -7,9 +7,8 @@ from tortoise.models import Model from tortoise.expressions import Q from passlib.hash import bcrypt +from pwncore.models.container import Container -if TYPE_CHECKING: - from pwncore.models.container import Container __all__ = ("User", "Team") @@ -18,10 +17,10 @@ class User(Model): # Registration numbers and other identity tags # abstractly just represents one person, expand this # field for Identity providers - tag : fields.CharField = fields.CharField(128, unique=True) - name : str = fields.TextField() - email : str = fields.TextField() - phone_num : fields.CharField = fields.CharField(15) + tag = fields.CharField(128, unique=True) + name = fields.TextField() + email = fields.TextField() + phone_num = fields.CharField(15) team: fields.ForeignKeyNullableRelation[Team] = fields.ForeignKeyField( "models.Team", "members", null=True, on_delete=fields.OnDelete.SET_NULL @@ -38,14 +37,8 @@ async def save(self, *args, **kwargs): class Team(Model): - name : fields.CharField = fields.CharField(255, unique=True) - password : fields.CharField = fields.CharField(255, unique=True) - current_points : int = fields.IntField(null=True) - current_stage : int = fields.IntField(null=True) - last_timestamp : fields.DatetimeField = fields.DatetimeField(null=True) + name = fields.CharField(255, unique=True) + secret_hash = fields.TextField() members: fields.ReverseRelation[User] containers: fields.ReverseRelation[Container] - - async def check_password(self, password : str) -> bool: - return bcrypt.verify(self.password, password) diff --git a/src/pwncore/routes/__init__.py b/src/pwncore/routes/__init__.py index 2c6095a..b231b01 100644 --- a/src/pwncore/routes/__init__.py +++ b/src/pwncore/routes/__init__.py @@ -1,16 +1,11 @@ from fastapi import APIRouter -from pwncore.routes import ctf, team +from pwncore.routes import ctf, team, auth # 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) - - -# Miscellaneous routes -@router.get("/asd") -async def a(): - return {"ASD": "asd"} diff --git a/src/pwncore/routes/auth.py b/src/pwncore/routes/auth.py new file mode 100644 index 0000000..a45b5ab --- /dev/null +++ b/src/pwncore/routes/auth.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +import jwt +from fastapi import APIRouter, HTTPException, Depends, Header, Response +from pwncore.models import User, Team +from pwncore.config import config +import datetime +from passlib.hash import bcrypt +from pydantic import BaseModel +from tortoise.transactions import atomic + +# 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): + reponse.status_code = 406 + return {"msg_code": config.msg_codes["team_exists"]} + 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): + + # 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 " + 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) +async def get_jwt(*, authorization: str = Header()): + token = authorization.split(" ")[1] # Remove Bearer + if token is None: + raise HTTPException(status_code=401) + try: + decoded_token = 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 diff --git a/src/pwncore/routes/ctf/start.py b/src/pwncore/routes/ctf/start.py index 6395d72..bc1c92a 100644 --- a/src/pwncore/routes/ctf/start.py +++ b/src/pwncore/routes/ctf/start.py @@ -1,26 +1,20 @@ from __future__ import annotations -from fastapi import APIRouter, Response +from fastapi import APIRouter, Response, Depends import uuid from tortoise.transactions import atomic from pwncore.models import Problem, Container, Ports, Team from pwncore.container import docker_client from pwncore.config import config - -# temporary helper functions -if config.development: - - def get_team_id(): - return 1 - +from pwncore.routes.auth import get_jwt router = APIRouter(tags=["ctf"]) @atomic() @router.post("/start/{ctf_id}") -async def start_docker_container(ctf_id: int, response: Response): +async def start_docker_container(ctf_id: int, response: Response, jwt = Depends(get_jwt)): """ image_config contains the raw POST data that gets sent to the Docker Remote API. For now it just contains the guest ports that need to be opened on the host. @@ -49,7 +43,7 @@ async def start_docker_container(ctf_id: int, response: Response): response.status_code = 404 return {"msg_code": config.msg_codes["ctf_not_found"]} - team_id = get_team_id() # From JWT + team_id = jwt["team_id"] # From JWT team_container = await Container.get_or_none(team=team_id, problem=ctf_id) if team_container: db_ports = await team_container.ports.all().values("port") # Get ports from DB @@ -121,8 +115,8 @@ async def start_docker_container(ctf_id: int, response: Response): @atomic() @router.post("/stopall") -async def stopall_docker_container(response: Response): - team_id = get_team_id() # From JWT +async def stopall_docker_container(response: Response, jwt = Depends(get_jwt)): + team_id = jwt["team_id"] # From JWT containers = await Container.filter(team_id=team_id).values() @@ -144,13 +138,13 @@ async def stopall_docker_container(response: Response): @atomic() @router.post("/stop/{ctf_id}") -async def stop_docker_container(ctf_id: int, response: Response): +async def stop_docker_container(ctf_id: int, response: Response, jwt = Depends(get_jwt)): ctf = await Problem.get_or_none(id=ctf_id) if not ctf: response.status_code = 404 return {"msg_code": config.msg_codes["ctf_not_found"]} - team_id = get_team_id() + team_id = jwt["team_id"] team_container = await Container.get_or_none(team_id=team_id, problem_id=ctf_id) if not team_container: return {"msg_code": config.msg_codes["container_not_found"]} diff --git a/src/pwncore/routes/team.py b/src/pwncore/routes/team.py index 807b904..5e50308 100644 --- a/src/pwncore/routes/team.py +++ b/src/pwncore/routes/team.py @@ -13,77 +13,3 @@ metadata = {"name": "team", "description": "Operations with teams"} router = APIRouter(prefix="/team", tags=["team"]) - -Team_pydantic = pydantic_model_creator(Team, name='Team') -TeamIn_pydantic = pydantic_model_creator(Team, name='TeamIn', exclude_readonly=True) - - -oauth2_scheme = OAuth2PasswordBearer(tokenUrl='login') - - -async def get_current_team(token : str = Depends(oauth2_scheme)): - try: - payload = jwt.decode(token, config.jwt_secret, algorithms=['HS256']) - team = await Team.get(id=payload.get('id')) - except: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail="Invalid team name or password" - ) - return await Team_pydantic.from_tortoise_orm(team) - - -async def generate_token(team_data : TeamIn_pydantic): - team = await Team.get_or_none(name=team_data.name) - if team is not None and await team.check_password(team_data.password): - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail="Invalid team name or password" - ) - - team_obj = await Team_pydantic.from_tortoise_orm(team) - expiration_time = datetime.datetime.utcnow() + datetime.timedelta(hours=2) - payload = team_obj.dict() - payload['exp'] = expiration_time - token = jwt.encode(payload, config.jwt_secret, algorithm='HS256') - - return { - "access_token" : token, - "token_type" : "bearer", - "expires_at" : expiration_time - } - - -@router.post('/signup', response_model=Team_pydantic) -async def signup_team(team : TeamIn_pydantic): - team_obj = await Team.create( - name=team.name, - password=bcrypt.hash(team.password) - ) - return await Team_pydantic.from_tortoise_orm(team_obj) - - -@router.post('/login') -async def team_login( - team_data : TeamIn_pydantic, - token : str = Depends(oauth2_scheme), - team : Team_pydantic = Depends(get_current_team) -): - - if team_data.name in (jwt.decode(token, config.jwt_secret, algorithms=['HS256'])): - issued_at = datetime.datetime.utcfromtimestamp(token['iat']) - current_time = datetime.datetime.utcnow() - idle_time = current_time - issued_at - - if idle_time < datetime.timedelta(hours=2): - expiration_time = current_time + datetime.timedelta(hours=2) - token_payload = jwt.decode(token, config.jwt_secret, algorithms=['HS256']) - token_payload['exp'] = expiration_time - token = jwt.encode(token_payload, config.jwt_secret, algorithm='HS256') - - return { - "access_token": token, - "token_type": "bearer" - } - else: - return generate_token(team_data) From a192b60d0e99d0cb051dd4942bd082aa01fb0d1e Mon Sep 17 00:00:00 2001 From: mradigen <55953083+mradigen@users.noreply.github.com> Date: Thu, 23 Nov 2023 23:28:38 +0530 Subject: [PATCH 04/13] chore: Pass Python 3.11 tests --- poetry.lock | 32 +++++++++++++++++++++++++-- pyproject.toml | 2 ++ setup.cfg | 3 +++ src/pwncore/config.py | 2 +- src/pwncore/models/user.py | 4 +--- src/pwncore/routes/auth.py | 38 +++++++++++++++++---------------- src/pwncore/routes/ctf/start.py | 11 ++++++---- src/pwncore/routes/team.py | 9 +------- 8 files changed, 65 insertions(+), 36 deletions(-) create mode 100644 setup.cfg diff --git a/poetry.lock b/poetry.lock index 8260a2a..009f5f3 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.7.0 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand. [[package]] name = "aiodocker" @@ -1156,6 +1156,23 @@ files = [ {file = "pyflakes-3.1.0.tar.gz", hash = "sha256:a0aae034c444db0071aa077972ba4768d40c830d9539fd45bf4cd3f8f6992efc"}, ] +[[package]] +name = "pyjwt" +version = "2.8.0" +description = "JSON Web Token implementation in Python" +optional = false +python-versions = ">=3.7" +files = [ + {file = "PyJWT-2.8.0-py3-none-any.whl", hash = "sha256:59127c392cc44c2da5bb3192169a91f429924e17aff6534d70fdc02ab3e04320"}, + {file = "PyJWT-2.8.0.tar.gz", hash = "sha256:57e28d156e3d5c10088e0c68abb90bfac3df82b40a71bd0daa20c65ccd5c23de"}, +] + +[package.extras] +crypto = ["cryptography (>=3.4.0)"] +dev = ["coverage[toml] (==5.0.4)", "cryptography (>=3.4.0)", "pre-commit", "pytest (>=6.0.0,<7.0.0)", "sphinx (>=4.5.0,<5.0.0)", "sphinx-rtd-theme", "zope.interface"] +docs = ["sphinx (>=4.5.0,<5.0.0)", "sphinx-rtd-theme", "zope.interface"] +tests = ["coverage[toml] (==5.0.4)", "pytest (>=6.0.0,<7.0.0)"] + [[package]] name = "pypika-tortoise" version = "0.1.6" @@ -1255,6 +1272,17 @@ asyncodbc = ["asyncodbc (>=0.1.1,<0.2.0)"] asyncpg = ["asyncpg"] psycopg = ["psycopg[binary,pool] (>=3.0.12,<4.0.0)"] +[[package]] +name = "types-passlib" +version = "1.7.7.13" +description = "Typing stubs for passlib" +optional = false +python-versions = "*" +files = [ + {file = "types-passlib-1.7.7.13.tar.gz", hash = "sha256:f152639f1f2103d7f59a56e2aec5f9398a75a80830991d0d68aac5c2b9c32a77"}, + {file = "types_passlib-1.7.7.13-py3-none-any.whl", hash = "sha256:414b5ee9c88313357c9261cfcf816509b1e8e4673f0796bd61e9ef249f6fe076"}, +] + [[package]] name = "typing-extensions" version = "4.8.0" @@ -1418,4 +1446,4 @@ multidict = ">=4.0" [metadata] lock-version = "2.0" python-versions = "^3.11" -content-hash = "ee606dd31a24ced4709bd82d354df9958af1e2f56d8c18ba96af2534907e0d25" +content-hash = "0615f2c4261b1fd9cc1ef6f701824148e16cde379c246253c315a9626c6e9cad" diff --git a/pyproject.toml b/pyproject.toml index 57cab21..df6f684 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,8 @@ 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" diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 0000000..2da2e6b --- /dev/null +++ b/setup.cfg @@ -0,0 +1,3 @@ +[flake8] +# Use Depends and Header from FastAPI without flake8 throwing an error. +extend-immutable-calls = Depends, Header diff --git a/src/pwncore/config.py b/src/pwncore/config.py index 1bade51..48e02be 100644 --- a/src/pwncore/config.py +++ b/src/pwncore/config.py @@ -28,7 +28,7 @@ "wrong_password": 10, "login_success": 11, "team_not_found": 12, - "team_exists": 13 + "team_exists": 13, } diff --git a/src/pwncore/models/user.py b/src/pwncore/models/user.py index 85afad0..41348e3 100644 --- a/src/pwncore/models/user.py +++ b/src/pwncore/models/user.py @@ -1,12 +1,9 @@ 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 passlib.hash import bcrypt from pwncore.models.container import Container @@ -37,6 +34,7 @@ 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() diff --git a/src/pwncore/routes/auth.py b/src/pwncore/routes/auth.py index a45b5ab..32bd195 100644 --- a/src/pwncore/routes/auth.py +++ b/src/pwncore/routes/auth.py @@ -1,8 +1,8 @@ from __future__ import annotations import jwt -from fastapi import APIRouter, HTTPException, Depends, Header, Response -from pwncore.models import User, Team +from fastapi import APIRouter, Header, Response, HTTPException +from pwncore.models import Team from pwncore.config import config import datetime from passlib.hash import bcrypt @@ -10,7 +10,10 @@ from tortoise.transactions import atomic # Metadata at the top for instant accessibility -metadata = {"name": "auth", "description": "Authentication using a JWT using a single access token."} +metadata = { + "name": "auth", + "description": "Authentication using a JWT using a single access token.", +} router = APIRouter(prefix="/auth", tags=["auth"]) @@ -27,45 +30,44 @@ class LoginBody(BaseModel): @atomic() -@router.post('/signup') +@router.post("/signup") async def signup_team(team: SignupBody, response: Response): try: if await Team.exists(name=team.name): - reponse.status_code = 406 + response.status_code = 406 return {"msg_code": config.msg_codes["team_exists"]} - await Team.create( - name=team.name, - secret_hash=bcrypt.hash(team.password) - ) + + # 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): - +@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): + 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') + 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 " return { "msg_code": config.msg_codes["login_success"], "access_token": token, - "token_type": "bearer" + "token_type": "bearer", } @@ -75,7 +77,7 @@ async def get_jwt(*, authorization: str = Header()): if token is None: raise HTTPException(status_code=401) try: - decoded_token = jwt.decode(token, config.jwt_secret, algorithms="HS256") + decoded_token = jwt.decode(token, config.jwt_secret, algorithm="HS256") except Exception: # Will filter for invalid signature/expired tokens raise HTTPException(status_code=401) return decoded_token diff --git a/src/pwncore/routes/ctf/start.py b/src/pwncore/routes/ctf/start.py index bc1c92a..5095d3a 100644 --- a/src/pwncore/routes/ctf/start.py +++ b/src/pwncore/routes/ctf/start.py @@ -14,7 +14,7 @@ @atomic() @router.post("/start/{ctf_id}") -async def start_docker_container(ctf_id: int, response: Response, jwt = Depends(get_jwt)): +async def start_docker_container(ctf_id: int, response: Response, jwt=Depends(get_jwt)): """ image_config contains the raw POST data that gets sent to the Docker Remote API. For now it just contains the guest ports that need to be opened on the host. @@ -54,7 +54,10 @@ async def start_docker_container(ctf_id: int, response: Response, jwt = Depends( "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 @@ -115,7 +118,7 @@ async def start_docker_container(ctf_id: int, response: Response, jwt = Depends( @atomic() @router.post("/stopall") -async def stopall_docker_container(response: Response, jwt = Depends(get_jwt)): +async def stopall_docker_container(response: Response, jwt=Depends(get_jwt)): team_id = jwt["team_id"] # From JWT containers = await Container.filter(team_id=team_id).values() @@ -138,7 +141,7 @@ async def stopall_docker_container(response: Response, jwt = Depends(get_jwt)): @atomic() @router.post("/stop/{ctf_id}") -async def stop_docker_container(ctf_id: int, response: Response, jwt = Depends(get_jwt)): +async def stop_docker_container(ctf_id: int, response: Response, jwt=Depends(get_jwt)): ctf = await Problem.get_or_none(id=ctf_id) if not ctf: response.status_code = 404 diff --git a/src/pwncore/routes/team.py b/src/pwncore/routes/team.py index 5e50308..d97af88 100644 --- a/src/pwncore/routes/team.py +++ b/src/pwncore/routes/team.py @@ -1,13 +1,6 @@ from __future__ import annotations -import jwt -from fastapi import APIRouter, HTTPException, status, Depends -from fastapi.security import OAuth2PasswordBearer -from tortoise.contrib.pydantic import pydantic_model_creator -from pwncore.models import User, Team -from pwncore.config import config -import datetime -from passlib.hash import bcrypt +from fastapi import APIRouter # Metadata at the top for instant accessibility metadata = {"name": "team", "description": "Operations with teams"} From cc0769fba90bac3c70d2a2d71408a4e32d26fd97 Mon Sep 17 00:00:00 2001 From: mradigen <55953083+mradigen@users.noreply.github.com> Date: Thu, 23 Nov 2023 23:36:29 +0530 Subject: [PATCH 05/13] chore: Pass Python 3.11 tests (flake8) --- setup.cfg | 2 ++ tests/test_login.py | 4 ++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/setup.cfg b/setup.cfg index 2da2e6b..6dc5c29 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,3 +1,5 @@ [flake8] +# Ignore 79 line character limit +extend-ignore = E501 # Use Depends and Header from FastAPI without flake8 throwing an error. extend-immutable-calls = Depends, Header diff --git a/tests/test_login.py b/tests/test_login.py index 6153f91..dfa6d9e 100644 --- a/tests/test_login.py +++ b/tests/test_login.py @@ -5,10 +5,10 @@ client = TestClient(app) +# Example test def test_login(): # Send a GET response to the specified endpoint response = client.get("/api/team/login") # Evaluate the response against expected values - assert response.status_code == 200 - assert response.json() == {"status": "logged in!"} + assert response.status_code == 404 From 327618cfd681bb88a00e53431876baa561765faa Mon Sep 17 00:00:00 2001 From: mradigen <55953083+mradigen@users.noreply.github.com> Date: Fri, 24 Nov 2023 17:02:01 +0530 Subject: [PATCH 06/13] chore: Apply requested changes --- setup.cfg | 5 ----- src/pwncore/models/user.py | 4 +++- src/pwncore/routes/auth.py | 12 +++++++----- src/pwncore/routes/ctf/start.py | 12 ++++++------ tox.ini | 2 +- 5 files changed, 17 insertions(+), 18 deletions(-) delete mode 100644 setup.cfg diff --git a/setup.cfg b/setup.cfg deleted file mode 100644 index 6dc5c29..0000000 --- a/setup.cfg +++ /dev/null @@ -1,5 +0,0 @@ -[flake8] -# Ignore 79 line character limit -extend-ignore = E501 -# Use Depends and Header from FastAPI without flake8 throwing an error. -extend-immutable-calls = Depends, Header diff --git a/src/pwncore/models/user.py b/src/pwncore/models/user.py index 41348e3..f929422 100644 --- a/src/pwncore/models/user.py +++ b/src/pwncore/models/user.py @@ -34,7 +34,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 + 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() diff --git a/src/pwncore/routes/auth.py b/src/pwncore/routes/auth.py index 32bd195..c70bb3e 100644 --- a/src/pwncore/routes/auth.py +++ b/src/pwncore/routes/auth.py @@ -1,7 +1,7 @@ from __future__ import annotations import jwt -from fastapi import APIRouter, Header, Response, HTTPException +from fastapi import APIRouter, Header, Response, HTTPException, Depends from pwncore.models import Team from pwncore.config import config import datetime @@ -72,12 +72,14 @@ async def team_login(team_data: LoginBody, response: Response): # Custom JWT processing (since FastAPI's implentation deals with refresh tokens) -async def get_jwt(*, authorization: str = Header()): - token = authorization.split(" ")[1] # Remove Bearer - if token is None: - raise HTTPException(status_code=401) +# Supressing B008 in order to be able to use Header() in arguments +async def get_jwt(*, authorization: str = Header()): # noqa: B008 try: + token = authorization.split(" ")[1] # Remove Bearer decoded_token = jwt.decode(token, config.jwt_secret, algorithm="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 +require_jwt = Depends(get_jwt) diff --git a/src/pwncore/routes/ctf/start.py b/src/pwncore/routes/ctf/start.py index 5095d3a..cd8d348 100644 --- a/src/pwncore/routes/ctf/start.py +++ b/src/pwncore/routes/ctf/start.py @@ -1,20 +1,20 @@ from __future__ import annotations -from fastapi import APIRouter, Response, Depends +from fastapi import APIRouter, Response import uuid from tortoise.transactions import atomic from pwncore.models import Problem, Container, Ports, Team from pwncore.container import docker_client from pwncore.config import config -from pwncore.routes.auth import get_jwt +from pwncore.routes.auth import require_jwt router = APIRouter(tags=["ctf"]) @atomic() @router.post("/start/{ctf_id}") -async def start_docker_container(ctf_id: int, response: Response, jwt=Depends(get_jwt)): +async def start_docker_container(ctf_id: int, response: Response, jwt=require_jwt): """ image_config contains the raw POST data that gets sent to the Docker Remote API. For now it just contains the guest ports that need to be opened on the host. @@ -57,7 +57,7 @@ async def start_docker_container(ctf_id: int, response: Response, jwt=Depends(ge 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 @@ -118,7 +118,7 @@ async def start_docker_container(ctf_id: int, response: Response, jwt=Depends(ge @atomic() @router.post("/stopall") -async def stopall_docker_container(response: Response, jwt=Depends(get_jwt)): +async def stopall_docker_container(response: Response, jwt=require_jwt): team_id = jwt["team_id"] # From JWT containers = await Container.filter(team_id=team_id).values() @@ -141,7 +141,7 @@ async def stopall_docker_container(response: Response, jwt=Depends(get_jwt)): @atomic() @router.post("/stop/{ctf_id}") -async def stop_docker_container(ctf_id: int, response: Response, jwt=Depends(get_jwt)): +async def stop_docker_container(ctf_id: int, response: Response, jwt=require_jwt): ctf = await Problem.get_or_none(id=ctf_id) if not ctf: response.status_code = 404 diff --git a/tox.ini b/tox.ini index 0352b29..89b3472 100644 --- a/tox.ini +++ b/tox.ini @@ -43,7 +43,7 @@ skip_install = true deps = flake8 flake8-bugbear -commands = flake8 src/ tests/ +commands = flake8 --extend-ignore E501 src/ tests/ [testenv:debug] description = check if the package runs after any configuration changes, not for debugging dev changes From c50789bffa6dcf232952c9394cf8a724ae98acab Mon Sep 17 00:00:00 2001 From: mradigen <55953083+mradigen@users.noreply.github.com> Date: Fri, 24 Nov 2023 17:06:13 +0530 Subject: [PATCH 07/13] chore: Fix goddamn flake8 --- src/pwncore/routes/auth.py | 1 + src/pwncore/routes/ctf/start.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/src/pwncore/routes/auth.py b/src/pwncore/routes/auth.py index c70bb3e..9fd0eef 100644 --- a/src/pwncore/routes/auth.py +++ b/src/pwncore/routes/auth.py @@ -81,5 +81,6 @@ async def get_jwt(*, authorization: str = Header()): # noqa: B008 raise HTTPException(status_code=401) return decoded_token + # Using a pre-assigned variable everywhere else to follow flake8's B008 require_jwt = Depends(get_jwt) diff --git a/src/pwncore/routes/ctf/start.py b/src/pwncore/routes/ctf/start.py index cd8d348..795d3b4 100644 --- a/src/pwncore/routes/ctf/start.py +++ b/src/pwncore/routes/ctf/start.py @@ -56,7 +56,7 @@ async def start_docker_container(ctf_id: int, response: Response, jwt=require_jw if ( await Container.filter(team_id=team_id).count() - >= config.max_containers_per_team + >= config.max_containers_per_team # noqa: W503 ): return {"msg_code": config.msg_codes["container_limit_reached"]} From 3aba3f082b6f6cc5e0a118745ec2b455c0f55509 Mon Sep 17 00:00:00 2001 From: mradigen <55953083+mradigen@users.noreply.github.com> Date: Sat, 25 Nov 2023 16:45:03 +0530 Subject: [PATCH 08/13] chore: Remove redundant flake8 flag from tox --- tox.ini | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tox.ini b/tox.ini index 89b3472..0352b29 100644 --- a/tox.ini +++ b/tox.ini @@ -43,7 +43,7 @@ skip_install = true deps = flake8 flake8-bugbear -commands = flake8 --extend-ignore E501 src/ tests/ +commands = flake8 src/ tests/ [testenv:debug] description = check if the package runs after any configuration changes, not for debugging dev changes From 255c251341e39eec2d3306bef8f26ad20c606b4a Mon Sep 17 00:00:00 2001 From: WizzyGeek Date: Sat, 25 Nov 2023 19:56:57 +0530 Subject: [PATCH 09/13] fix: jwt.decode arguments misc: use PEP 563 syntax --- src/pwncore/routes/auth.py | 18 ++++++++++++------ src/pwncore/routes/ctf/start.py | 8 ++++---- 2 files changed, 16 insertions(+), 10 deletions(-) diff --git a/src/pwncore/routes/auth.py b/src/pwncore/routes/auth.py index 9fd0eef..7e7cf6e 100644 --- a/src/pwncore/routes/auth.py +++ b/src/pwncore/routes/auth.py @@ -1,14 +1,17 @@ from __future__ import annotations +import datetime +import typing as t + import jwt from fastapi import APIRouter, Header, Response, HTTPException, Depends -from pwncore.models import Team -from pwncore.config import config -import datetime 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", @@ -73,14 +76,17 @@ async def team_login(team_data: LoginBody, response: Response): # Custom JWT processing (since FastAPI's implentation deals with refresh tokens) # Supressing B008 in order to be able to use Header() in arguments -async def get_jwt(*, authorization: str = Header()): # noqa: B008 +def get_jwt(*, authorization: t.Annotated[str, Header()]) -> JwtInfo: # noqa: B008 try: token = authorization.split(" ")[1] # Remove Bearer - decoded_token = jwt.decode(token, config.jwt_secret, algorithm="HS256") + 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 -require_jwt = Depends(get_jwt) +JwtInfo = t.TypedDict("JwtInfo", {"team_id": int, "exp": int}) +RequireJwt = t.Annotated[JwtInfo, Depends(get_jwt)] diff --git a/src/pwncore/routes/ctf/start.py b/src/pwncore/routes/ctf/start.py index 795d3b4..2dbd2ef 100644 --- a/src/pwncore/routes/ctf/start.py +++ b/src/pwncore/routes/ctf/start.py @@ -7,14 +7,14 @@ from pwncore.models import Problem, Container, Ports, Team from pwncore.container import docker_client from pwncore.config import config -from pwncore.routes.auth import require_jwt +from pwncore.routes.auth import RequireJwt router = APIRouter(tags=["ctf"]) @atomic() @router.post("/start/{ctf_id}") -async def start_docker_container(ctf_id: int, response: Response, jwt=require_jwt): +async def start_docker_container(ctf_id: int, response: Response, jwt: RequireJwt): """ image_config contains the raw POST data that gets sent to the Docker Remote API. For now it just contains the guest ports that need to be opened on the host. @@ -118,7 +118,7 @@ async def start_docker_container(ctf_id: int, response: Response, jwt=require_jw @atomic() @router.post("/stopall") -async def stopall_docker_container(response: Response, jwt=require_jwt): +async def stopall_docker_container(response: Response, jwt: RequireJwt): team_id = jwt["team_id"] # From JWT containers = await Container.filter(team_id=team_id).values() @@ -141,7 +141,7 @@ async def stopall_docker_container(response: Response, jwt=require_jwt): @atomic() @router.post("/stop/{ctf_id}") -async def stop_docker_container(ctf_id: int, response: Response, jwt=require_jwt): +async def stop_docker_container(ctf_id: int, response: Response, jwt: RequireJwt): ctf = await Problem.get_or_none(id=ctf_id) if not ctf: response.status_code = 404 From 88a2d443d4af125c056e30f04b240538c1bf28cf Mon Sep 17 00:00:00 2001 From: WizzyGeek Date: Sat, 25 Nov 2023 20:12:14 +0530 Subject: [PATCH 10/13] chore: fix linting checks --- src/pwncore/config.py | 1 - src/pwncore/routes/ctf/start.py | 6 +++++- src/pwncore/routes/team.py | 3 ++- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/src/pwncore/config.py b/src/pwncore/config.py index f31190d..af36cae 100644 --- a/src/pwncore/config.py +++ b/src/pwncore/config.py @@ -31,7 +31,6 @@ "signup_success": 13, "wrong_password": 14, "login_success": 15, - "team_not_found": 16, "team_exists": 17, } diff --git a/src/pwncore/routes/ctf/start.py b/src/pwncore/routes/ctf/start.py index 05bf1d2..984f4ca 100644 --- a/src/pwncore/routes/ctf/start.py +++ b/src/pwncore/routes/ctf/start.py @@ -54,7 +54,11 @@ async def start_docker_container(ctf_id: int, response: Response, jwt: RequireJw "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 + ): # fmt: skip 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 18df487..4fcf4f6 100644 --- a/src/pwncore/routes/team.py +++ b/src/pwncore/routes/team.py @@ -8,6 +8,7 @@ router = APIRouter(prefix="/team", tags=["team"]) + # Retrieve team_id from cookies def get_team_id(): return 1 @@ -18,10 +19,10 @@ async def team_list(): teams = await Team_Pydantic.from_queryset(Team.all()) return teams + # 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 - \ No newline at end of file From b1442f217550ced0477455e8bcb4432e00e60e8a Mon Sep 17 00:00:00 2001 From: WizzyGeek Date: Sat, 25 Nov 2023 20:17:28 +0530 Subject: [PATCH 11/13] misc: remove unused noqa --- src/pwncore/routes/auth.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pwncore/routes/auth.py b/src/pwncore/routes/auth.py index 7e7cf6e..baa5cc0 100644 --- a/src/pwncore/routes/auth.py +++ b/src/pwncore/routes/auth.py @@ -76,7 +76,7 @@ async def team_login(team_data: LoginBody, response: Response): # 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: # noqa: B008 +def get_jwt(*, authorization: t.Annotated[str, Header()]) -> JwtInfo: try: token = authorization.split(" ")[1] # Remove Bearer decoded_token: JwtInfo = jwt.decode( From 97a4dbf389a9e1fbd1f2776aa47ab142c2d4e8d7 Mon Sep 17 00:00:00 2001 From: mradigen <55953083+mradigen@users.noreply.github.com> Date: Sat, 25 Nov 2023 22:38:47 +0530 Subject: [PATCH 12/13] chore: RequireJwt, rename start.py routes, admin routes --- src/pwncore/routes/__init__.py | 5 +- src/pwncore/routes/admin.py | 79 ++++++++++++++++++++++++++++++ src/pwncore/routes/ctf/__init__.py | 39 ++++++++------- src/pwncore/routes/ctf/start.py | 4 +- src/pwncore/routes/team.py | 11 ++--- 5 files changed, 109 insertions(+), 29 deletions(-) diff --git a/src/pwncore/routes/__init__.py b/src/pwncore/routes/__init__.py index b231b01..a12543c 100644 --- a/src/pwncore/routes/__init__.py +++ b/src/pwncore/routes/__init__.py @@ -1,6 +1,7 @@ from fastapi import APIRouter -from pwncore.routes import ctf, team, auth +from pwncore.routes import ctf, team, auth, admin +from pwncore.config import config # Main router (all routes go under /api) router = APIRouter(prefix="/api") @@ -9,3 +10,5 @@ router.include_router(auth.router) router.include_router(ctf.router) router.include_router(team.router) +if config.development: + router.include_router(admin.router) diff --git a/src/pwncore/routes/admin.py b/src/pwncore/routes/admin.py index e69de29..9f4dc65 100644 --- a/src/pwncore/routes/admin.py +++ b/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="email1@xyz.org", + ) + await User.create( + tag="23BCE1000", + name="def", + team_id=2, + phone_num=2222222222, + email="email1@xyz.org", + ) + await User.create( + tag="23BAI1000", + name="ghi", + team_id=2, + phone_num=3333333333, + email="email1@xyz.org", + ) + await User.create( + tag="23BRS2000", + name="ABC", + team_id=1, + phone_num=4444444444, + email="email1@xyz.org", + ) + await User.create( + tag="23BCE2000", + name="DEF", + team_id=1, + phone_num=5555555555, + email="email1@xyz.org", + ) + await User.create( + tag="23BAI2000", + name="GHI", + team_id=1, + phone_num=6666666666, + email="email1@xyz.org", + ) + 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") diff --git a/src/pwncore/routes/ctf/__init__.py b/src/pwncore/routes/ctf/__init__.py index 8f8d342..d10fbcc 100644 --- a/src/pwncore/routes/ctf/__init__.py +++ b/src/pwncore/routes/ctf/__init__.py @@ -1,6 +1,7 @@ from __future__ import annotations from fastapi import APIRouter, Response +from pydantic import BaseModel from pwncore.models import ( Problem, @@ -12,8 +13,8 @@ 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 +from pwncore.routes.auth import RequireJwt # Metadata at the top for instant accessibility metadata = { @@ -25,12 +26,9 @@ 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}" +class Flag(BaseModel): + text: str @router.get("/list") @@ -39,37 +37,38 @@ async def ctf_list(): 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): +@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) 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) + status = await SolvedProblem.exists(team_id=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 + team_id=team_id, flag=flag.text, problem_id=ctf_id ) if check_solved: - await SolvedProblem.create(team_id=get_team_id(), problem_id=ctf_id) + await SolvedProblem.create(team_id=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): +async def hint_get(ctf_id: int, response: Response, jwt: RequireJwt): + team_id = jwt["team_id"] 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()) + await Hint.filter(problem_id=ctf_id, viewedhints__team_id=team_id) .order_by("-order") .first() ) @@ -83,22 +82,24 @@ async def hint_get(ctf_id: int, response: Response): else: hint = await Hint.get(problem_id=ctf_id, order=0) - await ViewedHint.create(hint_id=hint.id, team_id=get_team_id()) + await ViewedHint.create(hint_id=hint.id, team_id=team_id) return {"text": hint.text, "order": hint.order} @router.get("/{ctf_id}/viewed_hints") -async def viewed_problem_hints_get(ctf_id: int): +async def viewed_problem_hints_get(ctf_id: int, jwt: RequireJwt): + team_id = jwt["team_id"] viewed_hints = await Hint_Pydantic.from_queryset( - Hint.filter(problem_id=ctf_id, viewedhints__team_id=get_team_id()) + Hint.filter(problem_id=ctf_id, viewedhints__team_id=team_id) ) return viewed_hints @router.get("/completed") -async def completed_problem_get(): +async def completed_problem_get(jwt: RequireJwt): + team_id = jwt["team_id"] problems = await Problem_Pydantic.from_queryset( - Problem.filter(solvedproblems__team_id=get_team_id()) + Problem.filter(solvedproblems__team_id=team_id) ) return problems diff --git a/src/pwncore/routes/ctf/start.py b/src/pwncore/routes/ctf/start.py index 984f4ca..1d76ee6 100644 --- a/src/pwncore/routes/ctf/start.py +++ b/src/pwncore/routes/ctf/start.py @@ -13,7 +13,7 @@ @atomic() -@router.post("/start/{ctf_id}") +@router.post("/{ctf_id}/start") async def start_docker_container(ctf_id: int, response: Response, jwt: RequireJwt): """ image_config contains the raw POST data that gets sent to the Docker Remote API. @@ -141,7 +141,7 @@ async def stopall_docker_container(response: Response, jwt: RequireJwt): @atomic() -@router.post("/stop/{ctf_id}") +@router.post("/{ctf_id}/stop") async def stop_docker_container(ctf_id: int, response: Response, jwt: RequireJwt): ctf = await Problem.get_or_none(id=ctf_id) if not ctf: diff --git a/src/pwncore/routes/team.py b/src/pwncore/routes/team.py index 4fcf4f6..a6f3640 100644 --- a/src/pwncore/routes/team.py +++ b/src/pwncore/routes/team.py @@ -2,6 +2,7 @@ from fastapi import APIRouter from pwncore.models import Team, User, Team_Pydantic, User_Pydantic +from pwncore.routes.auth import RequireJwt # Metadata at the top for instant accessibility metadata = {"name": "team", "description": "Operations with teams"} @@ -9,11 +10,6 @@ router = APIRouter(prefix="/team", tags=["team"]) -# Retrieve team_id from cookies -def get_team_id(): - return 1 - - @router.get("/list") async def team_list(): teams = await Team_Pydantic.from_queryset(Team.all()) @@ -22,7 +18,8 @@ async def team_list(): # 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())) +async def team_members(jwt: RequireJwt): + team_id = jwt["team_id"] + members = await User_Pydantic.from_queryset(User.filter(team_id=team_id)) # Incase of no members, it just returns an empty list. return members From 99d204cc24aa00d4196ec4b7b0921fbf921b75ae Mon Sep 17 00:00:00 2001 From: mradigen <55953083+mradigen@users.noreply.github.com> Date: Sat, 25 Nov 2023 22:41:47 +0530 Subject: [PATCH 13/13] chore: Rename text to flag for flag post --- src/pwncore/routes/ctf/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pwncore/routes/ctf/__init__.py b/src/pwncore/routes/ctf/__init__.py index d10fbcc..6410db5 100644 --- a/src/pwncore/routes/ctf/__init__.py +++ b/src/pwncore/routes/ctf/__init__.py @@ -28,7 +28,7 @@ class Flag(BaseModel): - text: str + flag: str @router.get("/list") @@ -51,7 +51,7 @@ async def flag_post(ctf_id: int, flag: Flag, response: Response, jwt: RequireJwt return {"msg_code": config.msg_codes["ctf_solved"]} check_solved = await Container.exists( - team_id=team_id, flag=flag.text, problem_id=ctf_id + team_id=team_id, flag=flag.flag, problem_id=ctf_id ) if check_solved: await SolvedProblem.create(team_id=team_id, problem_id=ctf_id)