From 441a112fb2c32cb6e35a790e239e1d4eb3fdb5b8 Mon Sep 17 00:00:00 2001
From: Ivan <iyuvlasov_1@edu.hse.ru>
Date: Mon, 12 Dec 2022 20:33:30 +0300
Subject: [PATCH] Added availability to create BotTokens for automated access
to API
---
auth/middleware.py | 12 +++++++++---
endpoints/__init__.py | 2 +-
endpoints/routes.py | 3 +++
endpoints/user.py | 29 ++++++++++++++++++++++++++++-
models/__init__.py | 2 +-
models/users.py | 40 +++++++++++++++++++++++++++++++++++-----
responses/errors.py | 7 +++++++
7 files changed, 84 insertions(+), 11 deletions(-)
diff --git a/auth/middleware.py b/auth/middleware.py
index aa0b5a4..b4463f0 100644
--- a/auth/middleware.py
+++ b/auth/middleware.py
@@ -1,10 +1,12 @@
+from datetime import datetime
+
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import Response
from auth.key import decrypt_token
-from models import User
-from responses.errors import TokenUndefinedError, UserDontExist
+from models import User, BotToken
+from responses.errors import TokenUndefinedError, UserDontExist, BotTokenDontExist
def parse_token(request: Request) -> str:
@@ -24,13 +26,17 @@ def parse_token(request: Request) -> str:
class JWTAuthenticationMiddleware(BaseHTTPMiddleware):
""" JWT Authentication middleware """
+
async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
token = parse_token(request)
data = decrypt_token(token)
user = await User.get_or_none(id=data["user_id"])
if user is None:
raise UserDontExist
+ if "bot_token_id" in data:
+ bot_token = await BotToken.get_or_none(id=data["bot_token_id"])
+ if bot_token is None or bot_token.expires_at.timestamp() < datetime.now().timestamp():
+ raise BotTokenDontExist
request.state.token_data = data
request.state.user = user
return await call_next(request)
-
diff --git a/endpoints/__init__.py b/endpoints/__init__.py
index 71ddfbe..26bd82d 100644
--- a/endpoints/__init__.py
+++ b/endpoints/__init__.py
@@ -1,5 +1,5 @@
from .oauth import github_oauth_callback, github_oauth_redirect
from .service import ping
-from .user import get_me
+from .user import get_me, create_bot_token, get_bot_tokens, delete_bot_token
from .parser import load_report
from .indicators import get_indicators_from_group
diff --git a/endpoints/routes.py b/endpoints/routes.py
index 0c31a9f..abb0dcc 100644
--- a/endpoints/routes.py
+++ b/endpoints/routes.py
@@ -14,6 +14,9 @@ api_routes = [
Route("/getMe", get_me, methods=["GET"]),
Route("/loadReport", load_report, methods=["POST"]),
Route("/getIndicatorsFromGroup", get_indicators_from_group, methods=["GET"]),
+ Route("/createBotToken", create_bot_token, methods=["GET"]),
+ Route("/getBotTokens", get_bot_tokens, methods=["GET"]),
+ Route("/deleteBotToken", delete_bot_token, methods=["GET"])
]
admin_routes = []
diff --git a/endpoints/user.py b/endpoints/user.py
index 1dd5b26..eb22d7e 100644
--- a/endpoints/user.py
+++ b/endpoints/user.py
@@ -1,7 +1,8 @@
from starlette.requests import Request
from models import User
-from models.users import UserPD
+from models.users import UserPD, BotToken
+from responses.errors import BotTokenDontExist
from responses.responses import OkResponse
@@ -10,3 +11,29 @@ async def get_me(request: Request):
user: User = request.state.user
user_model = UserPD.from_orm(user)
return OkResponse(user_model)
+
+
+async def create_bot_token(request: Request):
+ """ Creates and returns bot token """
+ user: User = request.state.user
+ expires = int(request.query_params.get("expires", 1))
+ token = user.create_token(for_bot=True)
+ await BotToken.create_from_token(token, expires)
+ return OkResponse({"token": token})
+
+
+async def get_bot_tokens(request: Request):
+ """ Returns all bot tokens of user """
+ user: User = request.state.user
+ tokens = await BotToken.filter(user=user).all()
+ return OkResponse({"tokens": [{"id": str(token.id)[:4], "expires_in": token.expires_at} for token in tokens]})
+
+
+async def delete_bot_token(request: Request):
+ """ Deletes bot token """
+ user: User = request.state.user
+ token_id = request.query_params["token_id"]
+ token = await BotToken.get_or_none(id=token_id)
+ if token is None:
+ raise BotTokenDontExist
+ return OkResponse({"deleted": await token.delete()})
diff --git a/models/__init__.py b/models/__init__.py
index 58086f9..852d9ef 100644
--- a/models/__init__.py
+++ b/models/__init__.py
@@ -1,3 +1,3 @@
-from .users import User
+from .users import User, BotToken
from .indicator import Indicator, IndicatorGroup
from .report import Report
diff --git a/models/users.py b/models/users.py
index 3afa643..4ff7544 100644
--- a/models/users.py
+++ b/models/users.py
@@ -1,9 +1,11 @@
from datetime import datetime, timedelta
+from uuid import uuid4
from jwcrypto.jwt import JWT
from pydantic import BaseModel
from tortoise import Model, fields
-from auth.key import key
+from auth.key import key, decrypt_token
+from responses.errors import UserDontExist
class User(Model):
@@ -14,14 +16,16 @@ class User(Model):
is_admin = fields.BooleanField(default=False)
is_private = fields.BooleanField(default=False)
- def create_token(self) -> str:
+ def create_token(self, for_bot=False) -> str:
""" Creates JWT token for user """
- jwt = JWT(header={"alg": "RS256"}, claims={
+ claims = {
"user_id": self.id,
- "is_api_token": False,
"exp": int((datetime.now() + timedelta(days=30)).timestamp()),
"iss": "ioc"
- })
+ }
+ if for_bot:
+ claims["bot_token_id"] = str(uuid4())
+ jwt = JWT(header={"alg": "RS256"}, claims=claims)
jwt.make_signed_token(key.jwk)
return jwt.serialize()
@@ -57,3 +61,29 @@ class UserPD(BaseModel):
'is_admin': {'exclude': True},
'is_private': {'exclude': True}
}
+
+
+class BotToken(Model):
+ """ ORM model of Bot Token """
+ id = fields.UUIDField(pk=True)
+ user = fields.ForeignKeyField('models.User', on_delete=fields.CASCADE)
+ expires_at = fields.DatetimeField()
+ created_at = fields.DatetimeField(auto_now=True)
+
+ class Meta:
+ table = "bot_tokens"
+
+ @staticmethod
+ async def create_from_token(token: str, expires_in: int = 30) -> 'BotToken':
+ """ Creates new bot token from token
+
+ Args:
+ token: Token, which will be used to create bot token
+ expires_in: Time in days, when token will be expired
+ """
+ data = decrypt_token(token)
+ user = await User.get_or_none(id=data["user_id"])
+ if user is None:
+ raise UserDontExist
+ return await BotToken.create(user=user, id=data["bot_token_id"],
+ expires_at=datetime.now() + timedelta(days=expires_in))
diff --git a/responses/errors.py b/responses/errors.py
index 8f26a64..a999a58 100644
--- a/responses/errors.py
+++ b/responses/errors.py
@@ -73,3 +73,10 @@ class IndicatorGroupDoesNotExist(ApiError):
code = 402
description = "Indicator group does not exist"
http_code = 404
+
+
+class BotTokenDontExist(ApiError):
+ """ Bot token not found in database """
+ code = 503
+ description = "Bot token does not exist"
+ http_code = 404
--
GitLab