Adds creating token from complex object (refs #11)

Author: vimalloc

docs/index.rst

@@ -15,6 +15,7 @@ documentation is coming soon!
    installation
    basic_usage
    add_custom_data_claims
+   tokens_from_complex_object
    refresh_tokens
    token_freshness
    changing_default_behavior

docs/tokens_from_complex_object.rst

@@ -0,0 +1,27 @@
+Tokens from Complex Objects
+===========================
+
+A common pattern will be to have your users information (such as username and
+password) stored on disk somewhere. Lets say for example that we have a database
+object stores a username, hashed password, and what roles this user is. In the
+access token, we want the identity to be the username or user_id. We also want
+to store the roles this user has access to in the access_token so that we don't
+have to look that information up from the database on every request. We could
+do this simple enough with the **user_claims_loader** mentioned in the last section.
+However, is we pass just the identity (username or userid) to the **user_claims_loader**,
+we would have to look up this user from the database twice. First time, when
+they access the login endpoint and we need to verify their username and password,
+and second time in the **user_claims_loader** function, so that we can fine what roles
+this user has. This isn't a huge deal, but obviously it could be more efficient.
+
+This extension provides the ability to pass any object to the **create_access_token**
+method, which will then be passed to the **user_claims_loader** method. This lets
+us access the database only once. However, as we still need the identity to be
+a JSON serializable object unique to this user, we need
+to take an addition step and use the optional **identity_lookup** kwarg in the
+**create_access_token** method. This lets us tell the system how to get the identity from
+an object.
+
+Here is an example of this in action
+
+.. literalinclude:: ../examples/tokens_from_complex_objects.py

examples/tokens_from_complex_objects.py

@@ -0,0 +1,63 @@
+from flask import Flask, jsonify, request
+from flask_jwt_extended import JWTManager, jwt_required, \
+    create_access_token, get_jwt_identity
+
+app = Flask(__name__)
+app.secret_key = 'super-secret'  # Change this!
+
+# Setup the Flask-JWT-Extended extension
+jwt = JWTManager(app)
+
+
+# This is a simple example of a complex object that we could build
+# a JWT from. In practice, this will normally be something that
+# requires a lookup from disk (such as SQLAlchemy)
+class UserObject:
+    def __init__(self, username, roles):
+        self.username = username
+        self.roles = roles
+
+
+# This method will get whatever object is passed into the
+# create_access_token method.
+@jwt.user_claims_loader
+def add_claims_to_access_token(user):
+    return {'roles': user.roles}
+
+
+@app.route('/login', methods=['POST'])
+def login():
+    username = request.json.get('username', None)
+    password = request.json.get('password', None)
+    if username != 'test' and password != 'test':
+        return jsonify({"msg": "Bad username or password"}), 401
+
+    # Create an example UserObject
+    user = UserObject(username='test', roles=['foo', 'bar'])
+
+    # We can now pass this complex object directly to the
+    # create_access_token method. This will allow us to access
+    # the properties of this object in the user_claims_loader
+    # function. Because this object is not json serializable itself,
+    # we also need to provide a way to get some which is json
+    # serializable and represents the identity of this token from
+    # the complex object. We pass a function to  the optional
+    # identity_lookup kwarg, which tells the create_access_token
+    # function how to get the identity from this object
+    access_token = create_access_token(
+        identity=user,
+        identity_lookup=lambda u: u.username
+    )
+
+    ret = {'access_token': access_token}
+    return jsonify(ret), 200
+
+
+@app.route('/protected', methods=['GET'])
+@jwt_required
+def protected():
+    current_user = get_jwt_identity()
+    return jsonify({'hello_from': current_user}), 200
+
+if __name__ == '__main__':
+    app.run()

flask_jwt_extended/utils.py

@@ -312,12 +312,30 @@ def create_refresh_token(identity):
     return refresh_token
 
 
-def create_access_token(identity, fresh=True):
+def create_access_token(identity, fresh=False, identity_lookup=None):
+    """
+    Creates a new access token
+
+    :param identity: The identity of this token. This can be any data that is
+                     json serializable. It can also be an object, in which case
+                     you can pass a function to identity_lookup which tells us
+                     how to get the identity out of this object. This is useful
+                     so you don't need to query disk twice, once for initially
+                     finding the identity in your login endpoint, and once for
+                     setting addition data in the JWT via the user_claims_loader
+    :param fresh: If this token should me markded as fresh, and can thus access
+                  fresh_jwt_required protected endpoints. Defaults to False
+    :param identity_lookup: Function to generate a json serilizable identity
+                            from the identity object
+    :return: A newly encoded JWT access token
+    """
     # Token options
     secret = _get_secret_key()
     access_expire_delta = get_access_expires()
     algorithm = get_algorithm()
     user_claims = current_app.jwt_manager.user_claims_callback(identity)
+    if identity_lookup:
+        identity = identity_lookup(identity)
 
     access_token = _encode_access_token(identity, secret, algorithm, access_expire_delta,
                                         fresh=fresh, user_claims=user_claims)

tests/test_jwt_encode_decode.py

@@ -6,7 +6,8 @@
 from flask import Flask
 from flask_jwt_extended.exceptions import JWTEncodeError, JWTDecodeError
 from flask_jwt_extended.utils import _encode_access_token, _encode_refresh_token, \
-    _decode_jwt
+    _decode_jwt, create_access_token
+from flask_jwt_extended.jwt_manager import JWTManager
 
 
 class JWTEncodeDecodeTests(unittest.TestCase):
@@ -330,3 +331,35 @@ def test_decode_invalid_jwt(self):
                 }
                 encoded_token = jwt.encode(token_data, 'secret', 'HS256').decode('utf-8')
                 _decode_jwt(encoded_token, 'secret', 'HS256')
+
+    def test_create_access_token_with_object(self):
+        # Complex object to test building a JWT from. Normally if you are using
+        # this functionality, this is something that would be retrieved from
+        # disk somewhere (think sqlalchemy)
+        class TestObject:
+            def __init__(self, username, roles):
+                self.username = username
+                self.roles = roles
+
+        # Setup the flask stuff
+        app = Flask(__name__)
+        app.secret_key = 'super=secret'
+        app.config['JWT_ALGORITHM'] = 'HS256'
+        jwt = JWTManager(app)
+
+        @jwt.user_claims_loader
+        def custom_claims(object):
+            return {
+                'roles': object.roles
+            }
+
+        # Create the token using the complex object
+        with app.test_request_context():
+            user = TestObject(username='foo', roles=['bar', 'baz'])
+            token = create_access_token(identity=user,
+                                        identity_lookup=lambda obj: obj.username)
+
+            # Decode the token and make sure the values are set properly
+            token_data = _decode_jwt(token, app.secret_key, app.config['JWT_ALGORITHM'])
+            self.assertEqual(token_data['identity'], 'foo')
+            self.assertEqual(token_data['user_claims']['roles'], ['bar', 'baz'])