Merge pull request #51 from faisalabujabal/master

Getting Token From Encoded Token

Author: vimalloc

flask_jwt_extended/blacklist.py

@@ -5,6 +5,7 @@
 
 from flask_jwt_extended.config import config
 from flask_jwt_extended.exceptions import RevokedTokenError
+from flask_jwt_extended.utils import get_jti
 
 # TODO make simplekv an optional dependency if blacklist is disabled
 
@@ -80,13 +81,24 @@ def unrevoke_token(jti):
     """
     Revoke a token
 
-    :param jti: The jti of the token to revoke
+    :param jti: The jti of the token to unrevoke
     """
     _update_token(jti, revoked=False)
 
 
 @_verify_blacklist_enabled
-def get_stored_token(jti):
+def get_stored_token(jti=None, encoded_token=None):
+    """
+    Get the stored token for the passed in jti or encoded_token
+
+    :param jti: The jti of the token
+    :param encoded_token: The encoded JWT string
+    :return: Python dictionary with the token information
+    """
+    if jti is None and encoded_token is not None:
+        jti = get_jti(encoded_token)
+    elif jti is None and encoded_token is None:
+        raise ValueError('Either jti or encoded_token is required')
     return _get_token_from_store(jti)
 
 

flask_jwt_extended/utils.py

@@ -32,6 +32,16 @@ def get_jwt_claims():
     return get_raw_jwt().get('user_claims', {})
 
 
+def get_jti(encoded_token):
+    """
+    Returns the JTI given the JWT encoded token
+
+    :param encoded_token: The encoded JWT string
+    :return: The JTI of the token
+    """
+    return decode_jwt(encoded_token, config.secret_key, config.algorithm, config.csrf_protect).get('jti')
+
+
 def _get_jwt_manager():
     try:
         return current_app.jwt_manager

tests/test_blacklist.py

@@ -36,10 +36,15 @@ def login():
             }
             return jsonify(ret), 200
 
-        @self.app.route('/auth/token/<identity>', methods=['GET'])
-        def get_single_token(identity):
+        @self.app.route('/auth/token/jti/<jti>', methods=['GET'])
+        @self.app.route('/auth/token/encoded_token/<encoded_token>', methods=['GET'])
+        @self.app.route('/auth/token/encoded_token/', methods=['GET'])
+        def get_single_token(jti=None, encoded_token=None):
             try:
-                return jsonify(get_stored_token(identity)), 200
+                if jti is not None:
+                    return jsonify(get_stored_token(jti=jti)), 200
+                else:
+                    return jsonify(get_stored_token(encoded_token=encoded_token)), 200
             except KeyError:
                 return jsonify({"msg": "token not found"}), 404
 
@@ -399,20 +404,34 @@ def test_get_specific_identity(self):
         self.assertEqual(len(data), 0)
 
     def test_get_stored_token(self):
-        self._login('test1')
+        access_token, refresh_token = self._login('test1')
         response = self.client.get('/auth/tokens')
         data = json.loads(response.get_data(as_text=True))
         refresh_jti = data[0]['token']['jti']
 
-        response = self.client.get('/auth/token/{}'.format(refresh_jti))
+        # Test getting the token by passing in JTI
+        response = self.client.get('/auth/token/jti/{}'.format(refresh_jti))
+        status_code = response.status_code
+        data = json.loads(response.get_data(as_text=True))
+        self.assertEqual(status_code, 200)
+        self.assertIn('token', data)
+        self.assertIn('revoked', data)
+        self.assertEqual(len(data), 2)
+
+        # Test getting the token by passing in the encoded token
+        response = self.client.get('/auth/token/encoded_token/{}'.format(refresh_token))
         status_code = response.status_code
         data = json.loads(response.get_data(as_text=True))
         self.assertEqual(status_code, 200)
         self.assertIn('token', data)
         self.assertIn('revoked', data)
         self.assertEqual(len(data), 2)
 
-        response = self.client.get('/auth/token/404notokenfound')
+        # Test passing neither throws an exception
+        with self.assertRaises(ValueError):
+            self.client.get('/auth/token/encoded_token/')
+
+        response = self.client.get('/auth/token/jti/404notokenfound')
         status_code = response.status_code
         data = json.loads(response.get_data(as_text=True))
         self.assertEqual(status_code, 404)