Fix jwt_required with flask_restless (refs #10)

Author: vimalloc

flask_jwt_extended/jwt_manager.py

@@ -1,36 +1,40 @@
 from flask import jsonify
 
+from flask_jwt_extended.exceptions import JWTDecodeError, NoAuthorizationError, \
+    InvalidHeaderError, WrongTokenError, RevokedTokenError, FreshTokenRequired
+from jwt import ExpiredSignatureError, InvalidTokenError
+
 
 class JWTManager:
     def __init__(self, app=None):
         # Function that will be called to add custom user claims to a JWT.
         self.user_claims_callback = lambda _: {}
 
         # Function that will be called when an expired token is received
-        self.expired_token_callback = lambda: (
+        self._expired_token_callback = lambda: (
             jsonify({'msg': 'Token has expired'}), 401
         )
 
         # Function that will be called when an invalid token is received
-        self.invalid_token_callback = lambda err: (
+        self._invalid_token_callback = lambda err: (
             jsonify({'msg': err}), 422
         )
 
         # Function that will be called when attempting to access a protected
         # endpoint without a valid token
-        self.unauthorized_callback = lambda: (
+        self._unauthorized_callback = lambda: (
             jsonify({'msg': 'Missing Authorization Header'}), 401
         )
 
         # Function that will be called when attempting to access a fresh_jwt_required
         # endpoint with a valid token that is not fresh
-        self.needs_fresh_token_callback = lambda: (
+        self._needs_fresh_token_callback = lambda: (
             jsonify({'msg': 'Fresh token required'}), 401
         )
 
         # Function that will be called when a revoked token attempts to access
         # a protected endpoint
-        self.revoked_token_callback = lambda: (
+        self._revoked_token_callback = lambda: (
             jsonify({'msg': 'Token has been revoked'}), 401
         )
 
@@ -45,6 +49,38 @@ def init_app(self, app):
         """
         app.jwt_manager = self
 
+        @app.errorhandler(NoAuthorizationError)
+        def handle_auth_error(e):
+            return self._unauthorized_callback()
+
+        @app.errorhandler(ExpiredSignatureError)
+        def handle_expired_error(e):
+            return self._expired_token_callback()
+
+        @app.errorhandler(InvalidHeaderError)
+        def handle_invalid_header_error(e):
+            return self._invalid_token_callback(str(e))
+
+        @app.errorhandler(InvalidTokenError)
+        def handle_invalid_token_error(e):
+            return self._invalid_token_callback(str(e))
+
+        @app.errorhandler(JWTDecodeError)
+        def handle_jwt_decode_error(e):
+            return self._invalid_token_callback(str(e))
+
+        @app.errorhandler(WrongTokenError)
+        def handle_wrong_token_error(e):
+            return self._invalid_token_callback(str(e))
+
+        @app.errorhandler(RevokedTokenError)
+        def hanlde_revoked_token_error(e):
+            return self._revoked_token_callback()
+
+        @app.errorhandler(FreshTokenRequired)
+        def handle_fresh_token_required(e):
+            return self._needs_fresh_token_callback()
+
     def user_claims_loader(self, callback):
         """
         This sets the callback method for adding custom user claims to a JWT.
@@ -66,7 +102,7 @@ def expired_token_loader(self, callback):
 
         Callback must be a function that takes zero arguments.
         """
-        self.expired_token_callback = callback
+        self._expired_token_callback = callback
         return callback
 
     def invalid_token_loader(self, callback):
@@ -79,7 +115,7 @@ def invalid_token_loader(self, callback):
         Callback must be a function that takes only one argument, which is the
         error message of why the token is invalid.
         """
-        self.invalid_token_callback = callback
+        self._invalid_token_callback = callback
         return callback
 
     def unauthorized_loader(self, callback):
@@ -92,7 +128,7 @@ def unauthorized_loader(self, callback):
         Callback must be a function that takes only one argument, which is the
         error message of why the token is invalid.
         """
-        self.unauthorized_callback = callback
+        self._unauthorized_callback = callback
         return callback
 
     def needs_fresh_token_loader(self, callback):
@@ -105,7 +141,7 @@ def needs_fresh_token_loader(self, callback):
 
         Callback must be a function that takes no arguments.
         """
-        self.needs_fresh_token_callback = callback
+        self._needs_fresh_token_callback = callback
         return callback
 
     def revoked_token_loader(self, callback):
@@ -118,5 +154,5 @@ def revoked_token_loader(self, callback):
 
         Callback must be a function that takes no arguments.
         """
-        self.revoked_token_callback = callback
+        self._revoked_token_callback = callback
         return callback

flask_jwt_extended/utils.py

@@ -20,7 +20,7 @@
     get_refresh_csrf_cookie_name, get_token_location, \
     get_csrf_header_name
 from flask_jwt_extended.exceptions import JWTEncodeError, JWTDecodeError, \
-    InvalidHeaderError, NoAuthorizationError, WrongTokenError, RevokedTokenError, \
+    InvalidHeaderError, NoAuthorizationError, WrongTokenError, \
     FreshTokenRequired
 from flask_jwt_extended.blacklist import check_if_token_revoked, store_token
 
@@ -206,33 +206,6 @@ def _decode_jwt_from_request(type):
         return _decode_jwt_from_cookies(type)
 
 
-def _handle_callbacks_on_error(fn):
-    """
-    Helper decorator that will catch any exceptions we expect to encounter
-    when dealing with a JWT, and call the appropriate callback function for
-    handling that error. Callback functions can be set in using the *_loader
-    methods in jwt_manager.
-    """
-    @wraps(fn)
-    def wrapper(*args, **kwargs):
-        m = current_app.jwt_manager
-
-        try:
-            return fn(*args, **kwargs)
-        except NoAuthorizationError:
-            return m.unauthorized_callback()
-        except jwt.ExpiredSignatureError:
-            return m.expired_token_callback()
-        except (InvalidHeaderError, jwt.InvalidTokenError, JWTDecodeError,
-                WrongTokenError) as e:
-            return m.invalid_token_callback(str(e))
-        except RevokedTokenError:
-            return m.revoked_token_callback()
-        except FreshTokenRequired:
-            return m.needs_fresh_token_callback()
-    return wrapper
-
-
 def jwt_required(fn):
     """
     If you decorate a vew with this, it will ensure that the requester has a valid
@@ -243,7 +216,6 @@ def jwt_required(fn):
 
     :param fn: The view function to decorate
     """
-    @_handle_callbacks_on_error
     @wraps(fn)
     def wrapper(*args, **kwargs):
         # Attempt to decode the token
@@ -275,7 +247,6 @@ def fresh_jwt_required(fn):
 
     :param fn: The view function to decorate
     """
-    @_handle_callbacks_on_error
     @wraps(fn)
     def wrapper(*args, **kwargs):
         # Attempt to decode the token
@@ -308,7 +279,6 @@ def jwt_refresh_token_required(fn):
     valid JWT refresh token before calling the actual view. If the token is
     invalid, expired, not present, etc, the appropiate callback will be called
     """
-    @_handle_callbacks_on_error
     @wraps(fn)
     def wrapper(*args, **kwargs):
         # Get the JWT

tests/test_jwt_manager.py

@@ -37,7 +37,7 @@ def test_default_user_claims_callback(self):
     def test_default_expired_token_callback(self):
         with self.app.test_request_context():
             m = JWTManager(self.app)
-            result = m.expired_token_callback()
+            result = m._expired_token_callback()
             status_code, data = self._parse_callback_result(result)
 
             self.assertEqual(status_code, 401)
@@ -47,7 +47,7 @@ def test_default_invalid_token_callback(self):
         with self.app.test_request_context():
             m = JWTManager(self.app)
             err = "Test error"
-            result = m.invalid_token_callback(err)
+            result = m._invalid_token_callback(err)
             status_code, data = self._parse_callback_result(result)
 
             self.assertEqual(status_code, 422)
@@ -56,7 +56,7 @@ def test_default_invalid_token_callback(self):
     def test_default_unauthorized_callback(self):
         with self.app.test_request_context():
             m = JWTManager(self.app)
-            result = m.unauthorized_callback()
+            result = m._unauthorized_callback()
             status_code, data = self._parse_callback_result(result)
 
             self.assertEqual(status_code, 401)
@@ -65,7 +65,7 @@ def test_default_unauthorized_callback(self):
     def test_default_needs_fresh_token_callback(self):
         with self.app.test_request_context():
             m = JWTManager(self.app)
-            result = m.needs_fresh_token_callback()
+            result = m._needs_fresh_token_callback()
             status_code, data = self._parse_callback_result(result)
 
             self.assertEqual(status_code, 401)
@@ -74,7 +74,7 @@ def test_default_needs_fresh_token_callback(self):
     def test_default_revoked_token_callback(self):
         with self.app.test_request_context():
             m = JWTManager(self.app)
-            result = m.revoked_token_callback()
+            result = m._revoked_token_callback()
             status_code, data = self._parse_callback_result(result)
 
             self.assertEqual(status_code, 401)
@@ -98,7 +98,7 @@ def test_custom_expired_token_callback(self):
             def custom_expired_token():
                 return jsonify({"res": "TOKEN IS EXPIRED FOOL"}), 422
 
-            result = m.expired_token_callback()
+            result = m._expired_token_callback()
             status_code, data = self._parse_callback_result(result)
 
             self.assertEqual(status_code, 422)
@@ -113,7 +113,7 @@ def test_custom_invalid_token_callback(self):
             def custom_invalid_token(err):
                 return jsonify({"err": err}), 200
 
-            result = m.invalid_token_callback(err)
+            result = m._invalid_token_callback(err)
             status_code, data = self._parse_callback_result(result)
 
             self.assertEqual(status_code, 200)
@@ -127,7 +127,7 @@ def test_custom_unauthorized_callback(self):
             def custom_unauthorized():
                 return jsonify({"err": "GOTTA LOGIN FOOL"}), 200
 
-            result = m.unauthorized_callback()
+            result = m._unauthorized_callback()
             status_code, data = self._parse_callback_result(result)
 
             self.assertEqual(status_code, 200)
@@ -141,7 +141,7 @@ def test_custom_needs_fresh_token_callback(self):
             def custom_token_needs_refresh():
                 return jsonify({'sub_status': 101}), 200
 
-            result = m.needs_fresh_token_callback()
+            result = m._needs_fresh_token_callback()
             status_code, data = self._parse_callback_result(result)
 
             self.assertEqual(status_code, 200)
@@ -154,7 +154,7 @@ def test_custom_revoked_token_callback(self):
             @m.revoked_token_loader
             def custom_revoken_token():
                 return jsonify({"err": "Nice knowing you!"}), 422
-            result = m.revoked_token_callback()
+            result = m._revoked_token_callback()
             status_code, data = self._parse_callback_result(result)
 
             self.assertEqual(status_code, 422)