Merge pull request jpadilla#110 from mark-adams/algo-whitelist

Added support for whitelist validation of the `alg` header

Author: jpadilla

README.md

@@ -26,19 +26,19 @@ $ pip install cryptography
 
 ```python
 import jwt
-jwt.encode({'some': 'payload'}, 'secret')
+jwt.encode({'some': 'payload'}, 'secret', algorithm='HS256')
 ```
 
 Additional headers may also be specified.
 
 ```python
-jwt.encode({'some': 'payload'}, 'secret', headers={'kid': '230498151c214b788dd97f22b85410a5'})
+jwt.encode({'some': 'payload'}, 'secret', algorithm='HS256', headers={'kid': '230498151c214b788dd97f22b85410a5'})
 ```
 
 Note the resulting JWT will not be encrypted, but verifiable with a secret key.
 
 ```python
-jwt.decode('someJWTstring', 'secret')
+jwt.decode('someJWTstring', 'secret', algorithms=['HS256'])
 ```
 
 If the secret is wrong, it will raise a `jwt.DecodeError` telling you as such.
@@ -83,12 +83,27 @@ currently supports:
 * RS384 - RSASSA-PKCS1-v1_5 signature algorithm using SHA-384 hash algorithm
 * RS512 - RSASSA-PKCS1-v1_5 signature algorithm using SHA-512 hash algorithm
 
-Change the algorithm with by setting it in encode:
+### Encoding
+You can specify which algorithm you would like to use to sign the JWT
+by using the `algorithm` parameter:
 
 ```python
-jwt.encode({'some': 'payload'}, 'secret', 'HS512')
+jwt.encode({'some': 'payload'}, 'secret', algorithm='HS512')
 ```
 
+### Decoding
+When decoding, you can specify which algorithms you would like to permit
+when validating the JWT by using the `algorithms` parameter which takes a list
+of allowed algorithms:
+
+```python
+jwt.decode(some_jwt, 'secret', algorithms=['HS512', 'HS256'])
+```
+
+In the above case, if the JWT has any value for its alg header other than
+HS512 or HS256, the claim will be rejected with an `InvalidAlgorithmError`.
+
+### Asymmetric (Public-key) Algorithms
 Usage of RSA (RS\*) and EC (EC\*) algorithms require a basic understanding
 of how public-key cryptography is used with regards to digital signatures.
 If you are unfamiliar, you may want to read
@@ -103,6 +118,7 @@ When using the ECDSA algorithms, the `key` argument is expected to
 be an Elliptic Curve public or private key in PEM format. The type of key
 (private or public) depends on whether you are signing or verifying.
 
+
 ## Support of registered claim names
 
 JSON Web Token defines some registered claim names and defines how they should

jwt/__init__.py

@@ -16,7 +16,9 @@
 __copyright__ = 'Copyright 2015 José Padilla'
 
 
-from .api import encode, decode, register_algorithm, PyJWT
+from .api import (
+    encode, decode, register_algorithm, unregister_algorithm, PyJWT
+)
 from .exceptions import (
     InvalidTokenError, DecodeError, ExpiredSignatureError,
     InvalidAudienceError, InvalidIssuerError,

jwt/algorithms.py

@@ -18,23 +18,28 @@
     has_crypto = False
 
 
-def _register_default_algorithms(pyjwt_obj):
+def get_default_algorithms():
     """
-    Registers the algorithms that are implemented by the library.
+    Returns the algorithms that are implemented by the library.
     """
-    pyjwt_obj.register_algorithm('none', NoneAlgorithm())
-    pyjwt_obj.register_algorithm('HS256', HMACAlgorithm(HMACAlgorithm.SHA256))
-    pyjwt_obj.register_algorithm('HS384', HMACAlgorithm(HMACAlgorithm.SHA384))
-    pyjwt_obj.register_algorithm('HS512', HMACAlgorithm(HMACAlgorithm.SHA512))
+    default_algorithms = {
+        'none': NoneAlgorithm(),
+        'HS256': HMACAlgorithm(HMACAlgorithm.SHA256),
+        'HS384': HMACAlgorithm(HMACAlgorithm.SHA384),
+        'HS512': HMACAlgorithm(HMACAlgorithm.SHA512)
+    }
 
     if has_crypto:
-        pyjwt_obj.register_algorithm('RS256', RSAAlgorithm(RSAAlgorithm.SHA256))
-        pyjwt_obj.register_algorithm('RS384', RSAAlgorithm(RSAAlgorithm.SHA384))
-        pyjwt_obj.register_algorithm('RS512', RSAAlgorithm(RSAAlgorithm.SHA512))
-
-        pyjwt_obj.register_algorithm('ES256', ECAlgorithm(ECAlgorithm.SHA256))
-        pyjwt_obj.register_algorithm('ES384', ECAlgorithm(ECAlgorithm.SHA384))
-        pyjwt_obj.register_algorithm('ES512', ECAlgorithm(ECAlgorithm.SHA512))
+        default_algorithms.update({
+            'RS256': RSAAlgorithm(RSAAlgorithm.SHA256),
+            'RS384': RSAAlgorithm(RSAAlgorithm.SHA384),
+            'RS512': RSAAlgorithm(RSAAlgorithm.SHA512),
+            'ES256': ECAlgorithm(ECAlgorithm.SHA256),
+            'ES384': ECAlgorithm(ECAlgorithm.SHA384),
+            'ES512': ECAlgorithm(ECAlgorithm.SHA512)
+        })
+
+    return default_algorithms
 
 
 class Algorithm(object):

jwt/api.py

@@ -5,24 +5,24 @@
 from collections import Mapping
 from datetime import datetime, timedelta
 
-from .algorithms import Algorithm, _register_default_algorithms  # NOQA
+from .algorithms import Algorithm, get_default_algorithms  # NOQA
 from .compat import string_types, text_type, timedelta_total_seconds
 from .exceptions import (
-    DecodeError, ExpiredSignatureError,
+    DecodeError, ExpiredSignatureError, InvalidAlgorithmError,
     InvalidAudienceError, InvalidIssuerError
 )
 from .utils import base64url_decode, base64url_encode
 
 
 class PyJWT(object):
     def __init__(self, algorithms=None):
-        self._algorithms = {}
+        self._algorithms = get_default_algorithms()
+        self._valid_algs = set(algorithms) if algorithms is not None else set(self._algorithms)
 
-        if algorithms is None:
-            _register_default_algorithms(self)
-        else:
-            for key, algo in algorithms.items():
-                self.register_algorithm(key, algo)
+        # Remove algorithms that aren't on the whitelist
+        for key in list(self._algorithms.keys()):
+            if key not in self._valid_algs:
+                del self._algorithms[key]
 
     def register_algorithm(self, alg_id, alg_obj):
         """
@@ -35,19 +35,34 @@ def register_algorithm(self, alg_id, alg_obj):
             raise TypeError('Object is not of type `Algorithm`')
 
         self._algorithms[alg_id] = alg_obj
+        self._valid_algs.add(alg_id)
 
-    def get_supported_algorithms(self):
+    def unregister_algorithm(self, alg_id):
+        """
+        Unregisters an Algorithm for use when creating and verifying tokens
+        Throws KeyError if algorithm is not registered.
+        """
+        if alg_id not in self._algorithms:
+            raise KeyError('The specified algorithm could not be removed because it is not registered.')
+
+        del self._algorithms[alg_id]
+        self._valid_algs.remove(alg_id)
+
+    def get_algorithms(self):
         """
         Returns a list of supported values for the 'alg' parameter.
         """
-        return self._algorithms.keys()
+        return list(self._valid_algs)
 
     def encode(self, payload, key, algorithm='HS256', headers=None, json_encoder=None):
         segments = []
 
         if algorithm is None:
             algorithm = 'none'
 
+        if algorithm not in self._valid_algs:
+            pass
+
         # Check that we get a mapping
         if not isinstance(payload, Mapping):
             raise TypeError('Expecting a mapping object, as json web token only'
@@ -94,12 +109,12 @@ def encode(self, payload, key, algorithm='HS256', headers=None, json_encoder=Non
 
         return b'.'.join(segments)
 
-    def decode(self, jwt, key='', verify=True, **kwargs):
+    def decode(self, jwt, key='', verify=True, algorithms=None, **kwargs):
         payload, signing_input, header, signature = self._load(jwt)
 
         if verify:
             self._verify_signature(payload, signing_input, header, signature,
-                                   key, **kwargs)
+                                   key, algorithms, **kwargs)
 
         return payload
 
@@ -142,24 +157,29 @@ def _load(self, jwt):
         return (payload, signing_input, header, signature)
 
     def _verify_signature(self, payload, signing_input, header, signature,
-                          key='', verify_expiration=True, leeway=0,
+                          key='', algorithms=None, verify_expiration=True, leeway=0,
                           audience=None, issuer=None):
 
+        alg = header['alg']
+
+        if algorithms is not None and alg not in algorithms:
+            raise InvalidAlgorithmError('The specified alg value is not allowed')
+
         if isinstance(leeway, timedelta):
             leeway = timedelta_total_seconds(leeway)
 
         if not isinstance(audience, (string_types, type(None))):
             raise TypeError('audience must be a string or None')
 
         try:
-            alg_obj = self._algorithms[header['alg']]
+            alg_obj = self._algorithms[alg]
             key = alg_obj.prepare_key(key)
 
             if not alg_obj.verify(signing_input, key, signature):
                 raise DecodeError('Signature verification failed')
 
         except KeyError:
-            raise DecodeError('Algorithm not supported')
+            raise InvalidAlgorithmError('Algorithm not supported')
 
         if 'nbf' in payload and verify_expiration:
             utc_timestamp = timegm(datetime.utcnow().utctimetuple())
@@ -196,3 +216,4 @@ def _verify_signature(self, payload, signing_input, header, signature,
 encode = _jwt_global_obj.encode
 decode = _jwt_global_obj.decode
 register_algorithm = _jwt_global_obj.register_algorithm
+unregister_algorithm = _jwt_global_obj.unregister_algorithm

jwt/exceptions.py

@@ -22,6 +22,10 @@ class InvalidKeyError(Exception):
     pass
 
 
+class InvalidAlgorithmError(InvalidTokenError):
+    pass
+
+
 # Compatibility aliases (deprecated)
 ExpiredSignature = ExpiredSignatureError
 InvalidAudience = InvalidAudienceError

tests/test_api.py

@@ -8,7 +8,10 @@
 
 from jwt.algorithms import Algorithm
 from jwt.api import PyJWT
-from jwt.exceptions import DecodeError, ExpiredSignatureError, InvalidAudienceError, InvalidIssuerError
+from jwt.exceptions import (
+    DecodeError, ExpiredSignatureError, InvalidAlgorithmError,
+    InvalidAudienceError, InvalidIssuerError
+)
 
 from .compat import text_type, unittest
 from .utils import ensure_bytes
@@ -35,20 +38,6 @@ def setUp(self):  # noqa
                         'claim': 'insanity'}
         self.jwt = PyJWT()
 
-    def test_if_algorithms_param_is_empty_dict_then_no_algorithms(self):
-        jwt_obj = PyJWT(algorithms={})
-        jwt_algorithms = jwt_obj.get_supported_algorithms()
-
-        self.assertEqual(len(jwt_algorithms), 0)
-
-    def test_algorithms_param_sets_algorithms(self):
-        algorithms = {'TESTALG': Algorithm()}
-        jwt_obj = PyJWT(algorithms=algorithms)
-
-        supported_algs = jwt_obj.get_supported_algorithms()
-        self.assertEqual(len(supported_algs), 1)
-        self.assertIn('TESTALG', supported_algs)
-
     def test_register_algorithm_does_not_allow_duplicate_registration(self):
         self.jwt.register_algorithm('AAA', Algorithm())
 
@@ -59,13 +48,43 @@ def test_register_algorithm_rejects_non_algorithm_obj(self):
         with self.assertRaises(TypeError):
             self.jwt.register_algorithm('AAA123', {})
 
+    def test_unregister_algorithm_removes_algorithm(self):
+        supported = self.jwt.get_algorithms()
+        self.assertIn('none', supported)
+        self.assertIn('HS256', supported)
+
+        self.jwt.unregister_algorithm('HS256')
+
+        supported = self.jwt.get_algorithms()
+        self.assertNotIn('HS256', supported)
+
+    def test_unregister_algorithm_throws_error_if_not_registered(self):
+        with self.assertRaises(KeyError):
+            self.jwt.unregister_algorithm('AAA')
+
+    def test_algorithms_parameter_removes_alg_from_algorithms_list(self):
+        self.assertIn('none', self.jwt.get_algorithms())
+        self.assertIn('HS256', self.jwt.get_algorithms())
+
+        self.jwt = PyJWT(algorithms=['HS256'])
+        self.assertNotIn('none', self.jwt.get_algorithms())
+        self.assertIn('HS256', self.jwt.get_algorithms())
+
     def test_encode_decode(self):
         secret = 'secret'
         jwt_message = self.jwt.encode(self.payload, secret)
         decoded_payload = self.jwt.decode(jwt_message, secret)
 
         self.assertEqual(decoded_payload, self.payload)
 
+    def test_decode_fails_when_alg_is_not_on_method_algorithms_param(self):
+        secret = 'secret'
+        jwt_token = self.jwt.encode(self.payload, secret, algorithm='HS256')
+        self.jwt.decode(jwt_token, secret)
+
+        with self.assertRaises(InvalidAlgorithmError):
+            self.jwt.decode(jwt_token, secret, algorithms=['HS384'])
+
     def test_decode_works_with_unicode_token(self):
         secret = 'secret'
         unicode_jwt = text_type(
@@ -170,7 +189,7 @@ def test_decode_algorithm_param_should_be_case_sensitive(self):
                        '.eyJoZWxsbyI6IndvcmxkIn0'
                        '.5R_FEPE7SW2dT9GgIxPgZATjFGXfUDOSwo7TtO_Kd_g')
 
-        with self.assertRaises(DecodeError) as context:
+        with self.assertRaises(InvalidAlgorithmError) as context:
             self.jwt.decode(example_jwt, 'secret')
 
         exception = context.exception
@@ -670,7 +689,7 @@ def test_encode_decode_with_rsa_sha512(self):
 
     def test_rsa_related_algorithms(self):
         self.jwt = PyJWT()
-        jwt_algorithms = self.jwt.get_supported_algorithms()
+        jwt_algorithms = self.jwt.get_algorithms()
 
         if has_crypto:
             self.assertTrue('RS256' in jwt_algorithms)
@@ -773,7 +792,7 @@ def test_encode_decode_with_ecdsa_sha512(self):
 
     def test_ecdsa_related_algorithms(self):
         self.jwt = PyJWT()
-        jwt_algorithms = self.jwt.get_supported_algorithms()
+        jwt_algorithms = self.jwt.get_algorithms()
 
         if has_crypto:
             self.assertTrue('ES256' in jwt_algorithms)