From bd5768527e6cec8966c451bb30159b018a5355be Mon Sep 17 00:00:00 2001
From: H4CK3R-01 <resiaknairolf@gmail.com>
Date: Tue, 22 Mar 2022 11:20:04 +0100
Subject: [PATCH] Improve api for bot

---
 api/auth.py             |  3 +++
 api/helper_functions.py | 31 +++++++++++++++++++++++++------
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/api/auth.py b/api/auth.py
index 7db89a7..8096012 100644
--- a/api/auth.py
+++ b/api/auth.py
@@ -11,6 +11,9 @@ def verify_token(token):
     if token is None:
         return False
 
+    if ':' in token:  # Bot token
+        token = token.split(":")[0]
+
     try:
         jwt.decode(token, os.getenv('SECRET_KEY'), algorithms=["HS256"])
         return True
diff --git a/api/helper_functions.py b/api/helper_functions.py
index 2791494..f5e6862 100644
--- a/api/helper_functions.py
+++ b/api/helper_functions.py
@@ -38,11 +38,30 @@ def extract_token_data(token):
         return None
 
 
-def get_username_from_token_data(token_data):
-    if token_data is not None:
-        return token_data['username']
-    else:
-        return None
+def get_username_from_token_data():
+    if 'Authorization' in request.headers:
+        token = request.headers['Authorization'].split(" ")[1]
+
+        if token is not None:
+            if ':' in token:  # Maybe bot token, check if token valid and return username after ":" then
+                username = token.split(":")[1]
+                token = token.split(":")[0]
+
+                try:
+                    if jwt.decode(token, os.getenv('SECRET_KEY'), algorithms=["HS256"])['username'] == "bot":
+                        return username
+                    else:
+                        return None
+                except jwt.exceptions.DecodeError:
+                    return None
+
+            else:  # "Normal" token, extract username from token
+                try:
+                    return jwt.decode(token, os.getenv('SECRET_KEY'), algorithms=["HS256"])['username']
+                except jwt.exceptions.DecodeError:
+                    return None
+
+    return None
 
 
 def get_user_id_from_username(username):
@@ -54,7 +73,7 @@ def get_user_id_from_username(username):
 
 def get_username_or_abort_401():
     # get username from jwt token
-    username = get_username_from_token_data(extract_token_data(get_token()))
+    username = get_username_from_token_data()
 
     if username is None:  # If token not provided or invalid -> return 401 code
         abort(401, message="Unable to login")