@decode_key_loader callback added for dynamic decode keys (#191)

Author: steinitzu

docs/api.rst

@@ -24,6 +24,7 @@ Configuring Flask-JWT-Extended
   .. automethod:: user_loader_callback_loader
   .. automethod:: user_loader_error_loader
   .. automethod:: unauthorized_loader
+  .. automethod:: decode_key_loader
 
 
 Protected endpoint decorators

docs/changing_default_behavior.rst

@@ -43,6 +43,8 @@ and what the return values of your callback functions need to be.
       - Function that is called to verify the user_claims data. Must return True or False
     * - :meth:`~flask_jwt_extended.JWTManager.claims_verification_failed_loader`
       - Function that is called when the user claims verification callback returns False
+    * - :meth:`~flask_jwt_extended.JWTManager.decode_key_loader`
+      - Function that is called to load the decode/secret key before verifying a token
 
 Dynamic token expires time
 ~~~~~~~~~~~~~~~~~~~~~~~~~~

flask_jwt_extended/default_callbacks.py

@@ -101,3 +101,7 @@ def default_verify_claims_failed_callback():
     error message with a 400 status code
     """
     return jsonify({config.error_msg_key: 'User claims verification failed'}), 400
+
+
+def default_decode_key_callback(claims):
+    return config.decode_key

flask_jwt_extended/jwt_manager.py

@@ -13,7 +13,8 @@
     default_user_identity_callback, default_invalid_token_callback,
     default_unauthorized_callback, default_needs_fresh_token_callback,
     default_revoked_token_callback, default_user_loader_error_callback,
-    default_claims_verification_callback, default_verify_claims_failed_callback
+    default_claims_verification_callback, default_verify_claims_failed_callback,
+    default_decode_key_callback
 )
 from flask_jwt_extended.tokens import (
     encode_refresh_token, encode_access_token
@@ -53,6 +54,7 @@ def __init__(self, app=None):
         self._token_in_blacklist_callback = None
         self._claims_verification_callback = default_claims_verification_callback
         self._verify_claims_failed_callback = default_verify_claims_failed_callback
+        self._decode_key_callback = default_decode_key_callback
 
         # Register this extension with the flask app now (if it is provided)
         if app is not None:
@@ -378,6 +380,21 @@ def claims_verification_failed_loader(self, callback):
         self._verify_claims_failed_callback = callback
         return callback
 
+    def decode_key_loader(self, callback):
+        """
+        This decorator sets the callback function for getting the JWT decode key and
+        can be used to dynamically choose the appropriate decode key based on token
+        contents.
+        The default implementation returns the decode key from config (either
+        `JWT_SECRET_KEY` or `JWT_PUBLIC_KEY` depending on signing algorithm).
+
+        *HINT*: The callback function must be a function that takes only **one** argument,
+        which is a dictionary of the claims encoded in the JWT and must return a *string*
+        which is the decode key to verify the token.
+        """
+        self._decode_key_callback = callback
+        return callback
+
     def _create_refresh_token(self, identity, expires_delta=None):
         if expires_delta is None:
             expires_delta = config.refresh_expires

flask_jwt_extended/utils.py

@@ -11,6 +11,7 @@
     RevokedTokenError, UserClaimsVerificationError, WrongTokenError
 )
 from flask_jwt_extended.tokens import decode_jwt
+import jwt
 
 
 # Proxy to access the current user
@@ -71,9 +72,14 @@ def decode_token(encoded_token, csrf_value=None):
     :param encoded_token: The encoded JWT to decode into a python dict.
     :param csrf_value: Expected CSRF double submit value (optional)
     """
+    jwt_manager = _get_jwt_manager()
+    unverified_claims = jwt.decode(
+        encoded_token, verify=False, algorithms=config.algorithm
+    )
+    secret = jwt_manager._decode_key_callback(unverified_claims)
     return decode_jwt(
         encoded_token=encoded_token,
-        secret=config.decode_key,
+        secret=secret,
         algorithm=config.algorithm,
         identity_claim_key=config.identity_claim_key,
         user_claims_key=config.user_claims_key,

tests/test_decode_tokens.py

@@ -3,7 +3,7 @@
 from datetime import timedelta
 
 from flask import Flask
-from jwt import ExpiredSignatureError
+from jwt import ExpiredSignatureError, InvalidSignatureError
 
 from flask_jwt_extended import (
     JWTManager, create_access_token, decode_token, create_refresh_token,
@@ -48,7 +48,9 @@ def empty_user_loader_return(identity):
     # returned via the decode_token call
     with app.test_request_context():
         token = create_access_token('username')
-        pure_decoded = jwt.decode(token, config.decode_key, algorithms=[config.algorithm])
+        unverfied_claims = jwt.decode(token, verify=False, algorithms=[config.algorithm])
+        decode_key = jwtM._decode_key_callback(unverfied_claims)
+        pure_decoded = jwt.decode(token, decode_key, algorithms=[config.algorithm])
         assert config.user_claims_key not in pure_decoded
         extension_decoded = decode_token(token)
         assert config.user_claims_key in extension_decoded
@@ -117,3 +119,27 @@ def test_get_jti(app, default_access_token):
 
     with app.test_request_context():
         assert default_access_token['jti'] == get_jti(token)
+
+
+def test_decode_key_callback(app, default_access_token):
+    jwtM = get_jwt_manager(app)
+    app.config['JWT_SECRET_KEY'] = 'correct secret'
+
+    @jwtM.decode_key_loader
+    def get_decode_key_1(claims):
+        return 'different secret'
+
+    assert jwtM._decode_key_callback({}) == 'different secret'
+
+    with pytest.raises(InvalidSignatureError):
+        with app.test_request_context():
+            token = encode_token(app, default_access_token)
+            decode_token(token)
+
+    @jwtM.decode_key_loader
+    def get_decode_key_2(claims):
+        return 'correct secret'
+
+    with app.test_request_context():
+        token = encode_token(app, default_access_token)
+        decode_token(token)