store cookies in jwt and option to csrf protect them

vimalloc · vimalloc · commit 01603bce2676 · 2016-10-10T14:56:06.000-06:00
closes #5
diff --git a/examples/token_in_cookie.py b/examples/token_in_cookie.py
@@ -5,47 +5,54 @@
     set_access_cookies, set_refresh_cookie
 
 
-# NOTE: This is being actively worked on, and is not complete yet. At present,
-#       this code will not work! It should be rolled out next week sometime
-
-
 app = Flask(__name__)
 app.secret_key = 'super-secret'  # Change this!
 jwt = JWTManager(app)
 
 
-# Configure application to store jwts in cookies with double submit csrf protection
-app.config['JWT_TOKEN_LOCATION'] = 'cookie'
-app.config['JWT_COOKIE_HTTPONLY'] = True
-app.config['JWT_COOKIE_SECURE'] = True
+# Configure application to store JWTs in cookies
+app.config['JWT_TOKEN_LOCATION'] = 'cookies'
+app.config['JWT_COOKIE_SECURE'] = False  # In prod this should likely be True
 
-app.config['JWT_ACCESS_COOKIE_NAME'] = 'access_token_cookie'
+# Set the cookie paths, so that you are only sending your access token cookie
+# to the access endpoints, and only sending your refresh token to the refresh
+# endpoint.
 app.config['JWT_ACCESS_COOKIE_PATH'] = '/api/'
-
-app.config['JWT_REFRESH_COOKIE_NAME'] = 'refresh_token_cookie'
 app.config['JWT_REFRESH_COOKIE_PATH'] = '/token/refresh'
 
+# Enable csrf double submit protection. Check out this for a simple overview
+# of what this is: http://stackoverflow.com/a/37396572/272689.
 app.config['JWT_COOKIE_CSRF_PROTECT'] = True
-app.config['JWT_ACCESS_CSRF_COOKIE_NAME'] = 'x_xsrf_access_token'
-app.config['JWT_REFRESH_CSRF_COOKIE_NAME'] = 'x_xsrf_refresh_token'
+
+
+# Now, whenever you make a request to a protected endpoint, you will need to
+# send in the access or refresh JWT via a cookie, as well as a custom header
+# which has the same csrf token that is in the cookie. You cannot access the
+# csrf token from the JWT, as httponly is set to true (and javascript thus
+# cannot see it), but you can get the JWT from a secondary cookie (that only
+# javascript on your site can access), and thus verify a csrf attack isn't
+# happening.
+#
+# You can modify the cookie name, csrf cookie name, and csrf header name via
+# various app.config options. Check the options page for details.
 
 
 @app.route('/token/auth', 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
+        return jsonify({'login': False}), 401
 
     # Create the tokens we will be sending back to the user
     access_token = create_access_token(identity=username)
     refresh_token = create_refresh_token(identity=username)
 
     # Set the JWTs and the CSRF double submit protection cookies in this response
-    resp = jsonify({'login': True}), 200
+    resp = jsonify({'login': True})
     set_access_cookies(resp, access_token)
     set_refresh_cookie(resp, refresh_token)
-    return resp
+    return resp, 200
 
 
 @app.route('/token/refresh', methods=['POST'])
@@ -56,9 +63,9 @@ def refresh():
     access_token = create_access_token(identity=current_user)
 
     # Set the access JWT and CSRF double submit protection cookies in this response
-    resp = jsonify({'refresh': True}), 200
+    resp = jsonify({'refresh': True})
     set_access_cookies(resp, access_token)
-    return resp
+    return resp, 200
 
 
 # We do not need to make any changes here, all of the protected endpoints will
diff --git a/flask_jwt_extended/exceptions.py b/flask_jwt_extended/exceptions.py
@@ -26,7 +26,7 @@ class InvalidHeaderError(JWTExtendedException):
     pass
 
 
-class NoAuthHeaderError(JWTExtendedException):
+class NoAuthorizationError(JWTExtendedException):
     """
     An error getting header information from a request
     """
diff --git a/flask_jwt_extended/utils.py b/flask_jwt_extended/utils.py
@@ -18,9 +18,10 @@
     get_access_cookie_name, get_cookie_secure, get_access_cookie_path, \
     get_cookie_csrf_protect, get_access_csrf_cookie_name, \
     get_refresh_cookie_name, get_refresh_cookie_path, \
-    get_refresh_csrf_cookie_name
+    get_refresh_csrf_cookie_name, get_token_location, \
+    get_access_csrf_header_name, get_refresh_csrf_header_name
 from flask_jwt_extended.exceptions import JWTEncodeError, JWTDecodeError, \
-    InvalidHeaderError, NoAuthHeaderError, WrongTokenError, RevokedTokenError, \
+    InvalidHeaderError, NoAuthorizationError, WrongTokenError, RevokedTokenError, \
     FreshTokenRequired
 from flask_jwt_extended.blacklist import check_if_token_revoked, store_token
 
@@ -42,8 +43,8 @@ def get_jwt_claims():
 
 
 # TODO set csrf token in jwt when creating tokens (if enabled)
-def _create_xsrf_token():
-    return binascii.hexlify(os.urandom(60))
+def _create_csrf_token():
+    return str(uuid.uuid4())
 
 
 def _encode_access_token(identity, secret, algorithm, token_expire_delta,
@@ -80,6 +81,8 @@ def _encode_access_token(identity, secret, algorithm, token_expire_delta,
         'type': 'access',
         'user_claims': user_claims,
     }
+    if get_token_location() == 'cookies' and get_cookie_csrf_protect():
+        token_data['csrf'] = _create_csrf_token()
     encoded_token = jwt.encode(token_data, secret, algorithm).decode('utf-8')
 
     # If blacklisting is enabled and configured to store access and refresh tokens,
@@ -110,6 +113,8 @@ def _encode_refresh_token(identity, secret, algorithm, token_expire_delta):
         'identity': identity,
         'type': 'refresh',
     }
+    if get_token_location() == 'cookies' and get_cookie_csrf_protect():
+        token_data['csrf'] = _create_csrf_token()
     encoded_token = jwt.encode(token_data, secret, algorithm).decode('utf-8')
 
     # If blacklisting is enabled, store this token in our key-value store
@@ -142,14 +147,17 @@ def _decode_jwt(token, secret, algorithm):
             raise JWTDecodeError("Missing or invalid claim: fresh")
         if 'user_claims' not in data or not isinstance(data['user_claims'], dict):
             raise JWTDecodeError("Missing or invalid claim: user_claims")
+    if get_token_location() == 'cookies' and get_cookie_csrf_protect():
+        if 'csrf' not in data or not isinstance(data['csrf'], six.string_types):
+            raise JWTDecodeError("Missing or invalid claim: csrf")
     return data
 
 
-def _decode_jwt_from_request():
+def _decode_jwt_from_headers():
     # Verify we have the auth header
     auth_header = request.headers.get('Authorization', None)
     if not auth_header:
-        raise NoAuthHeaderError("Missing Authorization Header")
+        raise NoAuthorizationError("Missing Authorization Header")
 
     # Make sure the header is valid
     expected_header = get_jwt_header_type()
@@ -170,6 +178,37 @@ def _decode_jwt_from_request():
     return _decode_jwt(token, secret, algorithm)
 
 
+def _decode_jwt_from_cookies(type):
+    if type == 'access':
+        cookie_key = get_access_cookie_name()
+        csrf_header_key = get_access_csrf_header_name()
+    else:
+        cookie_key = get_refresh_cookie_name()
+        csrf_header_key = get_refresh_csrf_header_name()
+
+    token = request.cookies.get(cookie_key)
+    if not token:
+        raise NoAuthorizationError('Missing cookie "{}"'.format(cookie_key))
+    secret = _get_secret_key()
+    algorithm = get_algorithm()
+    token = _decode_jwt(token, secret, algorithm)
+
+    if get_cookie_csrf_protect():
+        csrf = request.headers.get(csrf_header_key, None)
+        if not csrf or csrf != token['csrf']:
+            raise NoAuthorizationError("Missing or invalid csrf double submit header")
+
+    return token
+
+
+def _decode_jwt_from_request(type):
+    token_location = get_token_location()
+    if token_location == 'headers':
+        return _decode_jwt_from_headers()
+    else:
+        return _decode_jwt_from_cookies(type)
+
+
 def _handle_callbacks_on_error(fn):
     """
     Helper decorator that will catch any exceptions we expect to encounter
@@ -183,7 +222,7 @@ def wrapper(*args, **kwargs):
 
         try:
             return fn(*args, **kwargs)
-        except NoAuthHeaderError:
+        except NoAuthorizationError:
             return m.unauthorized_callback()
         except jwt.ExpiredSignatureError:
             return m.expired_token_callback()
@@ -211,7 +250,7 @@ def jwt_required(fn):
     @wraps(fn)
     def wrapper(*args, **kwargs):
         # Attempt to decode the token
-        jwt_data = _decode_jwt_from_request()
+        jwt_data = _decode_jwt_from_request(type='access')
 
         # Verify this is an access token
         if jwt_data['type'] != 'access':
@@ -243,7 +282,7 @@ def fresh_jwt_required(fn):
     @wraps(fn)
     def wrapper(*args, **kwargs):
         # Attempt to decode the token
-        jwt_data = _decode_jwt_from_request()
+        jwt_data = _decode_jwt_from_request(type='access')
 
         # Verify this is an access token
         if jwt_data['type'] != 'access':
@@ -276,7 +315,7 @@ def jwt_refresh_token_required(fn):
     @wraps(fn)
     def wrapper(*args, **kwargs):
         # Get the JWT
-        jwt_data = _decode_jwt_from_request()
+        jwt_data = _decode_jwt_from_request(type='refresh')
 
         # verify this is a refresh token
         if jwt_data['type'] != 'refresh':
@@ -329,11 +368,7 @@ def _get_csrf_token(encoded_token):
     secret = _get_secret_key()
     algorithm = get_algorithm()
     token = _decode_jwt(encoded_token, secret, algorithm)
-    try:
-        return token['csrf']
-    except KeyError:
-        raise RuntimeError('JWT does not have csrf token set. Is '
-                           'JWT_COOKIE_CSRF_PROTECT set to True?')
+    return token['csrf']
 
 
 def set_access_cookies(response, encoded_access_token):
diff --git a/tests/test_config.py b/tests/test_config.py
@@ -6,10 +6,16 @@
 
 from flask_jwt_extended.config import get_access_expires, get_refresh_expires, \
     get_algorithm, get_blacklist_enabled, get_blacklist_store, \
-    get_blacklist_checks, get_jwt_header_type, get_jwt_header_name
+    get_blacklist_checks, get_jwt_header_type, get_jwt_header_name, \
+    get_token_location, get_cookie_secure, get_access_cookie_name, \
+    get_refresh_cookie_name, get_access_cookie_path, get_refresh_cookie_path, \
+    get_cookie_csrf_protect, get_access_csrf_cookie_name, \
+    get_refresh_csrf_cookie_name, get_access_csrf_header_name, \
+    get_refresh_csrf_header_name
 from flask_jwt_extended import JWTManager
 
 
+
 class TestEndpoints(unittest.TestCase):
 
     def setUp(self):
@@ -18,37 +24,78 @@ def setUp(self):
         JWTManager(self.app)
         self.client = self.app.test_client()
 
+    #57, 69, 73, 77, 81, 85, 89, 93, 97, 101, 105
     def test_default_configs(self):
         with self.app.test_request_context():
+            self.assertEqual(get_token_location(), 'headers')
+            self.assertEqual(get_jwt_header_name(), 'Authorization')
+            self.assertEqual(get_jwt_header_type(), 'Bearer')
+
+            self.assertEqual(get_cookie_secure(), False)
+            self.assertEqual(get_access_cookie_name(), 'access_token_cookie')
+            self.assertEqual(get_refresh_cookie_name(), 'refresh_token_cookie')
+            self.assertEqual(get_access_cookie_path(), None)
+            self.assertEqual(get_refresh_cookie_path(), None)
+            self.assertEqual(get_cookie_csrf_protect(), True)
+            self.assertEqual(get_access_csrf_cookie_name(), 'csrf_access_token')
+            self.assertEqual(get_refresh_csrf_cookie_name(), 'csrf_refresh_token')
+            self.assertEqual(get_access_csrf_header_name(), 'X-CSRF-ACCESS-TOKEN')
+            self.assertEqual(get_refresh_csrf_header_name(), 'X-CSRF-REFRESH-TOKEN')
+
             self.assertEqual(get_access_expires(), timedelta(minutes=15))
             self.assertEqual(get_refresh_expires(), timedelta(days=30))
             self.assertEqual(get_algorithm(), 'HS256')
             self.assertEqual(get_blacklist_enabled(), False)
             self.assertEqual(get_blacklist_store(), None)
             self.assertEqual(get_blacklist_checks(), 'refresh')
-            self.assertEqual(get_jwt_header_name(), 'Authorization')
-            self.assertEqual(get_jwt_header_type(), 'Bearer')
 
     def test_override_configs(self):
+        self.app.config['JWT_TOKEN_LOCATION'] = 'cookies'
+        self.app.config['JWT_HEADER_NAME'] = 'Auth'
+        self.app.config['JWT_HEADER_TYPE'] = 'JWT'
+
+        self.app.config['JWT_COOKIE_SECURE'] = True
+        self.app.config['JWT_ACCESS_COOKIE_NAME'] = 'banana1'
+        self.app.config['JWT_REFRESH_COOKIE_NAME'] = 'banana2'
+        self.app.config['JWT_ACCESS_COOKIE_PATH'] = '/banana/'
+        self.app.config['JWT_REFRESH_COOKIE_PATH'] = '/banana2/'
+        self.app.config['JWT_COOKIE_CSRF_PROTECT'] = False
+        self.app.config['JWT_ACCESS_CSRF_COOKIE_NAME'] = 'banana1a'
+        self.app.config['JWT_REFRESH_CSRF_COOKIE_NAME'] = 'banana2a'
+        self.app.config['JWT_ACCESS_CSRF_HEADER_NAME'] = 'banana1b'
+        self.app.config['JWT_REFRESH_CSRF_HEADER_NAME'] = 'banana2b'
+
         self.app.config['JWT_ACCESS_TOKEN_EXPIRES'] = timedelta(minutes=5)
         self.app.config['JWT_REFRESH_TOKEN_EXPIRES'] = timedelta(days=7)
         self.app.config['JWT_ALGORITHM'] = 'HS512'
         self.app.config['JWT_BLACKLIST_ENABLED'] = True
         self.app.config['JWT_BLACKLIST_STORE'] = simplekv.memory.DictStore()
         self.app.config['JWT_BLACKLIST_TOKEN_CHECKS'] = 'all'
-        self.app.config['JWT_HEADER_NAME'] = 'Auth'
-        self.app.config['JWT_HEADER_TYPE'] = 'JWT'
 
         with self.app.test_request_context():
+            self.assertEqual(get_token_location(), 'cookies')
+            self.assertEqual(get_jwt_header_name(), 'Auth')
+            self.assertEqual(get_jwt_header_type(), 'JWT')
+
+            self.assertEqual(get_cookie_secure(), True)
+            self.assertEqual(get_access_cookie_name(), 'banana1')
+            self.assertEqual(get_refresh_cookie_name(), 'banana2')
+            self.assertEqual(get_access_cookie_path(), '/banana/')
+            self.assertEqual(get_refresh_cookie_path(), '/banana2/')
+            self.assertEqual(get_cookie_csrf_protect(), False)
+            self.assertEqual(get_access_csrf_cookie_name(), 'banana1a')
+            self.assertEqual(get_refresh_csrf_cookie_name(), 'banana2a')
+            self.assertEqual(get_access_csrf_header_name(), 'banana1b')
+            self.assertEqual(get_refresh_csrf_header_name(), 'banana2b')
+
             self.assertEqual(get_access_expires(), timedelta(minutes=5))
             self.assertEqual(get_refresh_expires(), timedelta(days=7))
             self.assertEqual(get_algorithm(), 'HS512')
             self.assertEqual(get_blacklist_enabled(), True)
             self.assertIsInstance(get_blacklist_store(), simplekv.memory.DictStore)
             self.assertEqual(get_blacklist_checks(), 'all')
-            self.assertEqual(get_jwt_header_name(), 'Auth')
-            self.assertEqual(get_jwt_header_type(), 'JWT')
 
+        self.app.config['JWT_TOKEN_LOCATION'] = 'banana'
         self.app.config['JWT_HEADER_NAME'] = ''
         self.app.config['JWT_ACCESS_TOKEN_EXPIRES'] = 'banana'
         self.app.config['JWT_REFRESH_TOKEN_EXPIRES'] = 'banana'
@@ -61,3 +108,5 @@ def test_override_configs(self):
                 get_access_expires()
             with self.assertRaises(RuntimeError):
                 get_refresh_expires()
+            with self.assertRaises(RuntimeError):
+                get_token_location()
diff --git a/tests/test_jwt_encode_decode.py b/tests/test_jwt_encode_decode.py
@@ -300,3 +300,33 @@ def test_decode_invalid_jwt(self):
                 }
                 encoded_token = jwt.encode(token_data, 'secret', 'HS256').decode('utf-8')
                 _decode_jwt(encoded_token, 'secret', 'HS256')
+
+        # Missing and bad csrf tokens
+        self.app.config['JWT_TOKEN_LOCATION'] = 'cookies'
+        self.app.config['JWT_COOKIE_CSRF_PROTECTION'] = True
+        with self.app.test_request_context():
+            now = datetime.utcnow()
+            with self.assertRaises(JWTDecodeError):
+                token_data = {
+                    'exp': now + timedelta(minutes=5),
+                    'iat': now,
+                    'nbf': now,
+                    'jti': 'banana',
+                    'identity': 'banana',
+                    'type': 'refresh',
+                }
+                encoded_token = jwt.encode(token_data, 'secret', 'HS256').decode('utf-8')
+                _decode_jwt(encoded_token, 'secret', 'HS256')
+
+            with self.assertRaises(JWTDecodeError):
+                token_data = {
+                    'exp': now + timedelta(minutes=5),
+                    'iat': now,
+                    'nbf': now,
+                    'jti': 'banana',
+                    'identity': 'banana',
+                    'type': 'refresh',
+                    'csrf': True
+                }
+                encoded_token = jwt.encode(token_data, 'secret', 'HS256').decode('utf-8')
+                _decode_jwt(encoded_token, 'secret', 'HS256')
diff --git a/tests/test_protected_endpoints.py b/tests/test_protected_endpoints.py
