From a8ca6282a7b1a3eefe20585383e0f352d0f35691 Mon Sep 17 00:00:00 2001 From: parrothacker1 Date: Thu, 23 Nov 2023 19:11:43 +0530 Subject: [PATCH] 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)