Add role check

Author: utix

flask_jwt/__init__.py

@@ -33,6 +33,7 @@
     'JWT_AUTH_USERNAME_KEY': 'username',
     'JWT_AUTH_PASSWORD_KEY': 'password',
     'JWT_ALGORITHM': 'HS256',
+    'JWT_ROLE': 'role',
     'JWT_LEEWAY': timedelta(seconds=10),
     'JWT_AUTH_HEADER_PREFIX': 'JWT',
     'JWT_EXPIRATION_DELTA': timedelta(seconds=300),
@@ -141,7 +142,7 @@ def _default_jwt_error_handler(error):
     ])), error.status_code, error.headers
 
 
-def _jwt_required(realm):
+def _jwt_required(realm, roles):
     """Does the actual work of verifying the JWT data in the current request.
     This is done automatically for you by `jwt_required()` but you could call it manually.
     Doing so would be useful in the context of optional JWT access in your APIs.
@@ -163,17 +164,30 @@ def _jwt_required(realm):
 
     if identity is None:
         raise JWTError('Invalid JWT', 'User does not exist')
+    if roles:
+        identity_role = identity.get(current_app.config['JWT_ROLE'])
+        if not identity_role:
+            raise JWTError('Bad Request', 'Invalid credentials')
+        if not hasattr(identity_role, "__iter__"):
+            identity_role = [identity_role]
+        if not hasattr(roles, "__iter__"):
+            roles = [roles]
+        if not identity_role or not set(roles).intersection(identity_role):
+            raise JWTError('Bad Request', 'Invalid credentials')
 
 
-def jwt_required(realm=None):
+
+def jwt_required(realm=None, roles=None):
     """View decorator that requires a valid JWT token to be present in the request
 
     :param realm: an optional realm
+    :param roles: an optional list of roles allowed,
+                  the role is pick in JWT_ROLE field of identity
     """
     def wrapper(fn):
         @wraps(fn)
         def decorator(*args, **kwargs):
-            _jwt_required(realm or current_app.config['JWT_DEFAULT_REALM'])
+            _jwt_required(realm or current_app.config['JWT_DEFAULT_REALM'], roles)
             return fn(*args, **kwargs)
         return decorator
     return wrapper