Avoid using hardcoded 'identity'

vimalloc · Jul 13, 2017 · f150fe0 · f150fe0
1 parent 3efd522
commit f150fe0
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 9 deletions.
diff --git a/examples/database_blacklist/app.py b/examples/database_blacklist/app.py
@@ -56,8 +56,8 @@ def login():
         refresh_token = create_refresh_token(identity=username)
 
         # Store the tokens in our store with a status of not currently revoked.
-        add_token_to_database(access_token)
-        add_token_to_database(refresh_token)
+        add_token_to_database(access_token, app.config['JWT_IDENTITY_CLAIM'])
+        add_token_to_database(refresh_token, app.config['JWT_IDENTITY_CLAIM'])
 
         ret = {
             'access_token': access_token,
@@ -72,7 +72,7 @@ def refresh():
         # Do the same thing that we did in the login endpoint here
         current_user = get_jwt_identity()
         access_token = create_access_token(identity=current_user)
-        add_token_to_database(access_token)
+        add_token_to_database(access_token, app.config['JWT_IDENTITY_CLAIM'])
         return jsonify({'access_token': access_token}), 201
 
     # Provide a way for a user to look at their tokens

diff --git a/examples/database_blacklist/blacklist_helpers.py b/examples/database_blacklist/blacklist_helpers.py
@@ -16,14 +16,15 @@ def _epoch_utc_to_datetime(epoch_utc):
     return datetime.fromtimestamp(epoch_utc)
 
 
-def add_token_to_database(encoded_token):
+def add_token_to_database(encoded_token, identity_claim):
     """
     Adds a new token to the database. It is not revoked when it is added.
+    :param identity_claim:
     """
     decoded_token = decode_token(encoded_token)
     jti = decoded_token['jti']
     token_type = decoded_token['type']
-    user_identity = decoded_token['identity']
+    user_identity = decoded_token[identity_claim]
     expires = _epoch_utc_to_datetime(decoded_token['exp'])
     revoked = False
 

diff --git a/flask_jwt_extended/view_decorators.py b/flask_jwt_extended/view_decorators.py
@@ -33,7 +33,7 @@ def jwt_required(fn):
     def wrapper(*args, **kwargs):
         jwt_data = _decode_jwt_from_request(request_type='access')
         ctx_stack.top.jwt = jwt_data
-        _load_user(jwt_data['identity'])
+        _load_user(jwt_data[config.identity_claim])
         return fn(*args, **kwargs)
     return wrapper
 
@@ -53,7 +53,7 @@ def wrapper(*args, **kwargs):
         try:
             jwt_data = _decode_jwt_from_request(request_type='access')
             ctx_stack.top.jwt = jwt_data
-            _load_user(jwt_data['identity'])
+            _load_user(jwt_data[config.identity_claim])
         except NoAuthorizationError:
             pass
         return fn(*args, **kwargs)
@@ -77,7 +77,7 @@ def wrapper(*args, **kwargs):
             raise FreshTokenRequired('Fresh token required')
 
         ctx_stack.top.jwt = jwt_data
-        _load_user(jwt_data['identity'])
+        _load_user(jwt_data[config.identity_claim])
         return fn(*args, **kwargs)
     return wrapper
 
@@ -92,7 +92,7 @@ def jwt_refresh_token_required(fn):
     def wrapper(*args, **kwargs):
         jwt_data = _decode_jwt_from_request(request_type='refresh')
         ctx_stack.top.jwt = jwt_data
-        _load_user(jwt_data['identity'])
+        _load_user(jwt_data[config.identity_claim])
         return fn(*args, **kwargs)
     return wrapper