Merge pull request #304 from Seluj78/bugfix/jwt

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)
@@ -181,41 +162,6 @@ def expired_token_callback(expired_token):
 
 redis = StrictRedis(host=REDIS_HOST, port=REDIS_PORT, decode_responses=True, db=2)
 
-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,28 +203,10 @@ def check_if_token_is_revoked(decrypted_token):
     application.register_blueprint(debug_bp)
 
 
-@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,
-    )
-
-
 # import tasks here to be registered by celery
 
 import PyMatcha.utils.tasks  # noqa
+import PyMatcha.utils.jwt_callbacks  # noqa
 
 
 @application.route("/")

backend/PyMatcha/routes/api/auth/login.py

@@ -25,9 +25,7 @@
 from flask_jwt_extended import create_refresh_token
 from flask_jwt_extended import get_jti
 from flask_jwt_extended import get_jwt_identity
-from flask_jwt_extended import get_raw_jwt
 from flask_jwt_extended import jwt_refresh_token_required
-from flask_jwt_extended import jwt_required
 from PyMatcha import ACCESS_TOKEN_EXPIRES
 from PyMatcha import redis
 from PyMatcha import REFRESH_TOKEN_EXPIRES
@@ -89,17 +87,14 @@ def refresh():
     return SuccessOutput("access_token", access_token)
 
 
-@auth_login_bp.route("/auth/access_revoke", methods=["DELETE"])
-@jwt_required
+@auth_login_bp.route("/auth/logout", methods=["POST"])
+@validate_params({"access_token": str, "refresh_token": str})
 def logout():
-    jti = get_raw_jwt()["jti"]
-    redis.set("is_revoked_jti:" + jti, "true", ACCESS_TOKEN_EXPIRES * 1.2)
-    return Success("Access token revoked")
-
-
-@auth_login_bp.route("/auth/refresh_revoke", methods=["DELETE"])
-@jwt_refresh_token_required
-def logout2():
-    jti = get_raw_jwt()["jti"]
-    redis.set("is_revoked_jti:" + jti, "true", REFRESH_TOKEN_EXPIRES * 1.2)
-    return Success("Refresh token revoked")
+    data = request.get_json()
+    access_token = data["access_token"]
+    refresh_token = data["refresh_token"]
+    access_jti = get_jti(access_token)
+    refresh_jti = get_jti(refresh_token)
+    redis.set("is_revoked_jti:" + access_jti, "true", ACCESS_TOKEN_EXPIRES * 1.2)
+    redis.set("is_revoked_jti:" + refresh_jti, "true", REFRESH_TOKEN_EXPIRES * 1.2)
+    return Success("Logout successful.")

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,
+    )

backend/schemas/swagger.yaml

@@ -160,52 +160,13 @@ paths:
                     type: string
                     description: The new and fresh access token
                     example: eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpYXQiOjE2MDAzMjYzMjksIm5iZiI6MTYwMDMyNjMyOSwianRpIjoiZTA0ZWQ3NmMtMTcxMC00ZjQ4LWFkNzMtZDBmOTMxZTkxNzM0IiwiZXhwIjoxNjAwMzI3MjI5LCJpZGVudGl0eSI6eyJpZCI6MSwiZW1haWwiOiJmb29AZXhhbXBsZS5vcmciLCJ1c2VybmFtZSI6ImZvbyIsImlzX29ubGluZSI6MCwiZGF0ZV9sYXN0c2VlbiI6IldlZCwgMTYgU2VwIDIwMjAgMTU6Mzk6MDQgR01UIn0sImZyZXNoIjpmYWxzZSwidHlwZSI6ImFjY2VzcyJ9.dFJMy-04thKA9Q268bgQhiqLDrkqnCOCaiwHJ1XXook
-  /auth/access_revoke:
+  /auth/logout:
     post:
       tags:
         - Authentication
-      summary: Revokes the access token
+      summary: Logout user
       description: ""
-      operationId: revokeAccessToken
-      security:
-        - accessToken: [ ]
-      responses:
-        "409":
-          description: Unauthorized error, token already revoked.
-          content:
-            application/json:
-              schema:
-                type: object
-                properties:
-                  msg:
-                    type: string
-                    example: Token has been revoked
-                    default: Token has been revoked
-        "200":
-          description: Token successfully revoked
-          content:
-            application/json:
-              schema:
-                type: object
-                properties:
-                  success:
-                    type: boolean
-                    description: If the request is a success
-                    example: true
-                  code:
-                    type: integer
-                    description: The status code
-                    example: 200
-                  message:
-                    type: string
-                    example: Access token revoked
-  /auth/refresh_revoke:
-    post:
-      tags:
-        - Authentication
-      summary: Revokes the refresh token
-      description: ""
-      operationId: revokeRefreshToken
+      operationId: logoutUser
       security:
         - refreshToken: [ ]
       responses:
@@ -221,7 +182,7 @@ paths:
                     example: Token has been revoked
                     default: Token has been revoked
         "200":
-          description: Token successfully revoked
+          description: Logout successful.
           content:
             application/json:
               schema:
@@ -237,7 +198,7 @@ paths:
                     example: 200
                   message:
                     type: string
-                    example: Refresh token revoked
+                    example: Logout successful.
   /auth/confirm/{token}:
     post:
       summary: Confirm a user with a token