Moved callbacks in own function

Author: Seluj78

backend/PyMatcha/__init__.py

@@ -35,7 +35,6 @@
 from PyMatcha.utils.tables import create_tables
 from pymysql.cursors import DictCursor
 from redis import StrictRedis
-from sentry_sdk import configure_scope
 from sentry_sdk.integrations.flask import FlaskIntegration
 
 
@@ -127,24 +126,6 @@
 
 jwt = JWTManager(application)
 
-logging.debug("Configuring JWT expired token handler callback")
-
-
-@jwt.expired_token_loader
-def expired_token_callback(expired_token):
-    logging.debug("Token {} expired".format(expired_token))
-    resp = {
-        "code": 401,
-        "error": {
-            "message": f"The {expired_token['type']} token has expired",
-            "name": "Unauthorized Error",
-            "solution": "Try again when you have renewed your token",
-            "type": "UnauthorizedError",
-        },
-        "success": False,
-    }
-    return jsonify(resp), 401
-
 
 logging.debug("Configuring CORS")
 CORS(application, expose_headers="Authorization", supports_credentials=True)
@@ -183,39 +164,6 @@ def expired_token_callback(expired_token):
 
 redis.flushdb()
 
-from PyMatcha.models.user import get_user
-
-
-logging.debug("Configuring JWT user callback loader")
-
-
-from PyMatcha.utils.errors import NotFoundError
-
-
-@jwt.user_loader_callback_loader
-def jwt_user_callback(identity):
-    try:
-        user = get_user(identity["id"])
-    except NotFoundError:
-        # The user who the server issues the token for was deleted in the db.
-        return None
-
-    with configure_scope() as scope:
-        scope.user = {"email": user.email, "id": user.id, "username": user.username}
-    user.is_online = True
-    user.date_lastseen = datetime.datetime.utcnow()
-    user.save()
-    return user
-
-
-@jwt.token_in_blacklist_loader
-def check_if_token_is_revoked(decrypted_token):
-    jti = decrypted_token["jti"]
-    entry = redis.get("is_revoked_jti:" + jti)
-    if entry is None:
-        return True
-    return entry == "true"
-
 
 from PyMatcha.routes.api.user import user_bp
 from PyMatcha.routes.api.auth.email import auth_email_bp
@@ -257,66 +205,10 @@ def check_if_token_is_revoked(decrypted_token):
     application.register_blueprint(debug_bp)
 
 
-@jwt.revoked_token_loader
-def jwt_revoked_token_callback():
-    return (
-        jsonify(
-            {
-                "code": 401,
-                "error": {
-                    "message": "Token has been revoked.",
-                    "name": "Unauthorized Error",
-                    "solution": "Please login again",
-                    "type": "UnauthorizedError",
-                },
-                "success": False,
-            }
-        ),
-        401,
-    )
-
-
-@jwt.unauthorized_loader
-def no_jwt_callback(error_message):
-    return (
-        jsonify(
-            {
-                "code": 401,
-                "error": {
-                    "message": error_message,
-                    "name": "Unauthorized Error",
-                    "solution": "Try again",
-                    "type": "UnauthorizedError",
-                },
-                "success": False,
-            }
-        ),
-        401,
-    )
-
-
-@jwt.invalid_token_loader
-def jwt_invalid_token_callback(error_message):
-    return (
-        jsonify(
-            {
-                "code": 400,
-                "error": {
-                    "message": error_message,
-                    "name": "Bad Request Error",
-                    "solution": "Try again (The token is invalid)",
-                    "type": "BadRequestError",
-                },
-                "success": False,
-            }
-        ),
-        400,
-    )
-
-
 # import tasks here to be registered by celery
 
 import PyMatcha.utils.tasks  # noqa
+import PyMatcha.utils.jwt_callbacks  # noqa
 
 
 @application.route("/")

backend/PyMatcha/utils/jwt_callbacks.py

@@ -0,0 +1,108 @@
+import datetime
+import logging
+
+from flask import jsonify
+from PyMatcha import jwt
+from PyMatcha import redis
+from PyMatcha.models.user import get_user
+from PyMatcha.utils.errors import NotFoundError
+from sentry_sdk import configure_scope
+
+logging.debug("Configuring JWT callbacks")
+
+
+@jwt.expired_token_loader
+def expired_token_callback(expired_token):
+    resp = {
+        "code": 401,
+        "error": {
+            "message": f"The {expired_token['type']} token has expired",
+            "name": "Unauthorized Error",
+            "solution": "Try again when you have renewed your token",
+            "type": "UnauthorizedError",
+        },
+        "success": False,
+    }
+    return jsonify(resp), 401
+
+
+@jwt.user_loader_callback_loader
+def jwt_user_callback(identity):
+    try:
+        user = get_user(identity["id"])
+    except NotFoundError:
+        # The user who the server issues the token for was deleted in the db.
+        return None
+
+    with configure_scope() as scope:
+        scope.user = {"email": user.email, "id": user.id, "username": user.username}
+    user.is_online = True
+    user.date_lastseen = datetime.datetime.utcnow()
+    user.save()
+    return user
+
+
+@jwt.token_in_blacklist_loader
+def check_if_token_is_revoked(decrypted_token):
+    jti = decrypted_token["jti"]
+    entry = redis.get("is_revoked_jti:" + jti)
+    if entry is None:
+        return True
+    return entry == "true"
+
+
+@jwt.revoked_token_loader
+def jwt_revoked_token_callback():
+    return (
+        jsonify(
+            {
+                "code": 401,
+                "error": {
+                    "message": "Token has been revoked.",
+                    "name": "Unauthorized Error",
+                    "solution": "Please login again",
+                    "type": "UnauthorizedError",
+                },
+                "success": False,
+            }
+        ),
+        401,
+    )
+
+
+@jwt.unauthorized_loader
+def no_jwt_callback(error_message):
+    return (
+        jsonify(
+            {
+                "code": 401,
+                "error": {
+                    "message": error_message,
+                    "name": "Unauthorized Error",
+                    "solution": "Try again",
+                    "type": "UnauthorizedError",
+                },
+                "success": False,
+            }
+        ),
+        401,
+    )
+
+
+@jwt.invalid_token_loader
+def jwt_invalid_token_callback(error_message):
+    return (
+        jsonify(
+            {
+                "code": 400,
+                "error": {
+                    "message": error_message,
+                    "name": "Bad Request Error",
+                    "solution": "Try again (The token is invalid)",
+                    "type": "BadRequestError",
+                },
+                "success": False,
+            }
+        ),
+        400,
+    )