From a73725863e15a919b76beffcd6c4db5bc9724bcb Mon Sep 17 00:00:00 2001
From: Marcelo Arocha <marceloarocha@gmail.com>
Date: Thu, 28 Nov 2024 20:24:55 -0300
Subject: [PATCH] improve missing header error

---
 app/resources/api_decorator.py | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/app/resources/api_decorator.py b/app/resources/api_decorator.py
index c47f503..45dedd8 100644
--- 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")