Add callback function for customizing encode key (#193)

Refs #91

Author: vimalloc

docs/api.rst

@@ -14,17 +14,18 @@ Configuring Flask-JWT-Extended
   .. automethod:: init_app
   .. automethod:: claims_verification_loader
   .. automethod:: claims_verification_failed_loader
+  .. automethod:: decode_key_loader
+  .. automethod:: encode_key_loader
   .. automethod:: expired_token_loader
   .. automethod:: invalid_token_loader
   .. automethod:: needs_fresh_token_loader
   .. automethod:: revoked_token_loader
   .. automethod:: token_in_blacklist_loader
+  .. automethod:: unauthorized_loader
   .. automethod:: user_claims_loader
   .. automethod:: user_identity_loader
   .. 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

@@ -23,28 +23,30 @@ and what the return values of your callback functions need to be.
 
     * - Loader Decorator
       - Description
+    * - :meth:`~flask_jwt_extended.JWTManager.claims_verification_loader`
+      - 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 get the decode key before verifying a token
+    * - :meth:`~flask_jwt_extended.JWTManager.encode_key_loader`
+      - Function that is called to get the encode key before creating a token
     * - :meth:`~flask_jwt_extended.JWTManager.expired_token_loader`
       - Function to call when an expired token accesses a protected endpoint
     * - :meth:`~flask_jwt_extended.JWTManager.invalid_token_loader`
       - Function to call when an invalid token accesses a protected endpoint
-    * - :meth:`~flask_jwt_extended.JWTManager.unauthorized_loader`
-      - Function to call when a request with no JWT accesses a protected endpoint
     * - :meth:`~flask_jwt_extended.JWTManager.needs_fresh_token_loader`
       - Function to call when a non-fresh token accesses a :func:`~flask_jwt_extended.fresh_jwt_required` endpoint
     * - :meth:`~flask_jwt_extended.JWTManager.revoked_token_loader`
       - Function to call when a revoked token accesses a protected endpoint
+    * - :meth:`~flask_jwt_extended.JWTManager.token_in_blacklist_loader`
+      - Function that is called to check if a token has been revoked
+    * - :meth:`~flask_jwt_extended.JWTManager.unauthorized_loader`
+      - Function to call when a request with no JWT accesses a protected endpoint
     * - :meth:`~flask_jwt_extended.JWTManager.user_loader_callback_loader`
       - Function to call to load a user object when token accesses a protected endpoint
     * - :meth:`~flask_jwt_extended.JWTManager.user_loader_error_loader`
       - Function that is called when the user_loader callback function returns `None`
-    * - :meth:`~flask_jwt_extended.JWTManager.token_in_blacklist_loader`
-      - Function that is called to check if a token has been revoked
-    * - :meth:`~flask_jwt_extended.JWTManager.claims_verification_loader`
-      - 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

@@ -104,4 +104,16 @@ def default_verify_claims_failed_callback():
 
 
 def default_decode_key_callback(claims):
+    """
+    By default, the decode key specified via the JWT_SECRET_KEY or
+    JWT_PUBLIC_KEY settings will be used to decode all tokens
+    """
     return config.decode_key
+
+
+def default_encode_key_callback(identity):
+    """
+    By default, the encode key specified via the JWT_SECRET_KEY or
+    JWT_PRIVATE_KEY settings will be used to encode all tokens
+    """
+    return config.encode_key

flask_jwt_extended/jwt_manager.py

@@ -14,7 +14,7 @@
     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_decode_key_callback
+    default_decode_key_callback, default_encode_key_callback
 )
 from flask_jwt_extended.tokens import (
     encode_refresh_token, encode_access_token
@@ -55,6 +55,7 @@ def __init__(self, app=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
+        self._encode_key_callback = default_encode_key_callback
 
         # Register this extension with the flask app now (if it is provided)
         if app is not None:
@@ -385,16 +386,34 @@ 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).
+
+        The default implementation returns the decode key specified by
+        `JWT_SECRET_KEY` or `JWT_PUBLIC_KEY`, depending on the 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 unverified claims of the jwt (dictionary) and must return a *string*
         which is the decode key to verify the token.
         """
         self._decode_key_callback = callback
         return callback
 
+    def encode_key_loader(self, callback):
+        """
+        This decorator sets the callback function for getting the JWT encode key and
+        can be used to dynamically choose the appropriate encode key based on the
+        token identity.
+
+        The default implementation returns the encode key specified by
+        `JWT_SECRET_KEY` or `JWT_PRIVATE_KEY`, depending on the signing algorithm.
+
+        *HINT*: The callback function must be a function that takes only **one**
+        argument, which is the identity as passed into the create_access_token
+        or create_refresh_token functions, and must return a *string* which is
+        the decode key to verify the token.
+        """
+        self._encode_key_callback = callback
+        return callback
+
     def _create_refresh_token(self, identity, expires_delta=None):
         if expires_delta is None:
             expires_delta = config.refresh_expires
@@ -406,7 +425,7 @@ def _create_refresh_token(self, identity, expires_delta=None):
 
         refresh_token = encode_refresh_token(
             identity=self._user_identity_callback(identity),
-            secret=config.encode_key,
+            secret=self._encode_key_callback(identity),
             algorithm=config.algorithm,
             expires_delta=expires_delta,
             user_claims=user_claims,
@@ -423,7 +442,7 @@ def _create_access_token(self, identity, fresh=False, expires_delta=None):
 
         access_token = encode_access_token(
             identity=self._user_identity_callback(identity),
-            secret=config.encode_key,
+            secret=self._encode_key_callback(identity),
             algorithm=config.algorithm,
             expires_delta=expires_delta,
             fresh=fresh,

tests/test_decode_tokens.py

@@ -48,9 +48,7 @@ def empty_user_loader_return(identity):
     # returned via the decode_token call
     with app.test_request_context():
         token = create_access_token('username')
-        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])
+        pure_decoded = jwt.decode(token, config.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
@@ -121,25 +119,44 @@ def test_get_jti(app, default_access_token):
         assert default_access_token['jti'] == get_jti(token)
 
 
-def test_decode_key_callback(app, default_access_token):
+def test_encode_decode_callback_values(app, default_access_token):
     jwtM = get_jwt_manager(app)
-    app.config['JWT_SECRET_KEY'] = 'correct secret'
+    app.config['JWT_SECRET_KEY'] = 'foobarbaz'
+    with app.test_request_context():
+        assert jwtM._decode_key_callback({}) == 'foobarbaz'
+        assert jwtM._encode_key_callback({}) == 'foobarbaz'
 
     @jwtM.decode_key_loader
     def get_decode_key_1(claims):
         return 'different secret'
 
+    @jwtM.encode_key_loader
+    def get_decode_key_2(identity):
+        return 'different secret'
+
     assert jwtM._decode_key_callback({}) == 'different secret'
+    assert jwtM._encode_key_callback('') == 'different secret'
+
+
+def test_custom_encode_decode_key_callbacks(app, default_access_token):
+    jwtM = get_jwt_manager(app)
+    app.config['JWT_SECRET_KEY'] = 'foobarbaz'
+
+    @jwtM.encode_key_loader
+    def get_encode_key_1(identity):
+        assert identity == 'username'
+        return 'different secret'
 
     with pytest.raises(InvalidSignatureError):
         with app.test_request_context():
-            token = encode_token(app, default_access_token)
+            token = create_access_token('username')
             decode_token(token)
 
     @jwtM.decode_key_loader
-    def get_decode_key_2(claims):
-        return 'correct secret'
+    def get_decode_key_1(claims):
+        assert claims['identity'] == 'username'
+        return 'different secret'
 
     with app.test_request_context():
-        token = encode_token(app, default_access_token)
+        token = create_access_token('username')
         decode_token(token)