improve missing header error

noharm-ai · Nov 28, 2024 · a737258 · a737258
1 parent 1214246
commit a737258
Showing 1 changed file with 14 additions and 4 deletions.
diff --git a/app/resources/api_decorator.py b/app/resources/api_decorator.py
@@ -3,7 +3,7 @@
 from functools import wraps
 from flask import request
 from flask_api import status
-from jwt.exceptions import PyJWTError, ExpiredSignatureError
+from jwt.exceptions import PyJWTError, ExpiredSignatureError, InvalidTokenError
 
 from resources.connections import JWT_SECRET
 
@@ -15,9 +15,13 @@ def wrapper(f):
         def decorator_f(*args, **kwargs):
             try:
                 if JWT_SECRET:
-                    auth_type, auth_token = request.headers.get(
-                        "Authorization", ""
-                    ).split()
+                    authorization = request.headers.get("Authorization", "").split()
+
+                    if len(authorization) != 2:
+                        raise InvalidTokenError()
+
+                    auth_token = authorization[1]
+
                     jwt.decode(
                         jwt=auth_token,
                         key=JWT_SECRET,
@@ -33,6 +37,12 @@ def decorator_f(*args, **kwargs):
                     "message": "Token expirado",
                 }, status.HTTP_401_UNAUTHORIZED
 
+            except InvalidTokenError:
+                return {
+                    "status": "error",
+                    "message": "Token inválido",
+                }, status.HTTP_401_UNAUTHORIZED
+
             except PyJWTError as e:
                 logging.basicConfig()
                 logger = logging.getLogger("noharm.getname")