Implement a basic test suite for jwtutil and add extra checks to the decode method

This commit is contained in:
Joseph Schorr 2018-10-23 17:48:20 -04:00
parent 436e8cb760
commit 4868f17832
2 changed files with 149 additions and 9 deletions

View file

@ -5,7 +5,8 @@ from datetime import datetime, timedelta
from jwt import PyJWT
from jwt.exceptions import (
InvalidTokenError, DecodeError, InvalidAudienceError, ExpiredSignatureError,
ImmatureSignatureError, InvalidIssuedAtError, InvalidIssuerError, MissingRequiredClaimError
ImmatureSignatureError, InvalidIssuedAtError, InvalidIssuerError, MissingRequiredClaimError,
InvalidAlgorithmError
)
from cryptography.hazmat.backends import default_backend
@ -17,14 +18,19 @@ from jwkest.jwk import keyrep, RSAKey, ECKey
# TOKEN_REGEX defines a regular expression for matching JWT bearer tokens.
TOKEN_REGEX = re.compile(r'\ABearer (([a-zA-Z0-9+\-_/]+\.)+[a-zA-Z0-9+\-_/]+)\Z')
# ALGORITHM_WHITELIST defines a whitelist of allowed algorithms to be used in JWTs. DO NOT ADD
# `none` here!
ALGORITHM_WHITELIST = [
'rs256'
]
class StrictJWT(PyJWT):
""" StrictJWT defines a JWT decoder with extra checks. """
class _StrictJWT(PyJWT):
""" _StrictJWT defines a JWT decoder with extra checks. """
@staticmethod
def _get_default_options():
# Weird syntax to call super on a staticmethod
defaults = super(StrictJWT, StrictJWT)._get_default_options()
defaults = super(_StrictJWT, _StrictJWT)._get_default_options()
defaults.update({
'require_exp': True,
'require_iat': True,
@ -41,14 +47,14 @@ class StrictJWT(PyJWT):
options['verify_exp'] = True
# Do all of the other checks
super(StrictJWT, self)._validate_claims(payload, options, audience, issuer, leeway, **kwargs)
super(_StrictJWT, self)._validate_claims(payload, options, audience, issuer, leeway, **kwargs)
now = timegm(datetime.utcnow().utctimetuple())
self._reject_future_iat(payload, now, leeway)
if 'exp' in payload and options.get('exp_max_s') is not None:
# Validate that the expiration was not more than exp_max_s seconds after the issue time
# or in the absense of an issue time, more than exp_max_s in the future from now
# or in the absence of an issue time, more than exp_max_s in the future from now
# This will work because the parent method already checked the type of exp
expiration = datetime.utcfromtimestamp(int(payload['exp']))
@ -73,15 +79,30 @@ class StrictJWT(PyJWT):
' the future.')
def decode(jwt, key='', verify=True, algorithms=None, options=None,
**kwargs):
""" Decodes a JWT. """
if not algorithms:
raise InvalidAlgorithmError('algorithms must be specified')
normalized = set([a.lower() for a in algorithms])
if 'none' in normalized:
raise InvalidAlgorithmError('`none` algorithm is not allowed')
if set(normalized).intersection(set(ALGORITHM_WHITELIST)) != set(normalized):
raise InvalidAlgorithmError('Algorithms `%s` are not whitelisted. Allowed: %s' %
(algorithms, ALGORITHM_WHITELIST))
return _StrictJWT().decode(jwt, key, verify, algorithms, options, **kwargs)
def exp_max_s_option(max_exp_s):
""" Returns an options dictionary that sets the maximum expiration seconds for a JWT. """
return {
'exp_max_s': max_exp_s,
}
decode = StrictJWT().decode
def jwk_dict_to_public_key(jwk):
""" Converts the specified JWK into a public key. """
jwkest_key = keyrep(jwk)

View file

@ -0,0 +1,119 @@
import time
import pytest
import jwt
from Crypto.PublicKey import RSA
from jwkest.jwk import RSAKey
from util.security.jwtutil import (decode, exp_max_s_option, jwk_dict_to_public_key,
InvalidTokenError, InvalidAlgorithmError)
@pytest.fixture(scope="session")
def private_key():
return RSA.generate(2048)
@pytest.fixture(scope="session")
def private_key_pem(private_key):
return private_key.exportKey('PEM')
@pytest.fixture(scope="session")
def public_key(private_key):
return private_key.publickey().exportKey('PEM')
def _token_data(audience, subject, iss, iat=None, exp=None, nbf=None):
return {
'iss': iss,
'aud': audience,
'nbf': nbf() if nbf is not None else int(time.time()),
'iat': iat() if iat is not None else int(time.time()),
'exp': exp() if exp is not None else int(time.time() + 3600),
'sub': subject,
}
@pytest.mark.parametrize('aud, iss, nbf, iat, exp, expected_exception', [
pytest.param('invalidaudience', 'someissuer', None, None, None, 'Invalid audience',
id='invalid audience'),
pytest.param('someaudience', 'invalidissuer', None, None, None, 'Invalid issuer',
id='invalid issuer'),
pytest.param('someaudience', 'someissuer', lambda: time.time() + 120, None, None,
'The token is not yet valid',
id='invalid not before'),
pytest.param('someaudience', 'someissuer', None, lambda: time.time() + 120, None,
'Issued At claim',
id='issued at in future'),
pytest.param('someaudience', 'someissuer', None, None, lambda: time.time() - 100,
'Signature has expired',
id='already expired'),
pytest.param('someaudience', 'someissuer', None, None, lambda: time.time() + 10000,
'Token was signed for more than',
id='expiration too far in future'),
pytest.param('someaudience', 'someissuer', lambda: time.time() + 10, None, None,
None,
id='not before in future by within leeway'),
pytest.param('someaudience', 'someissuer', None, lambda: time.time() + 10, None,
None,
id='issued at in future but within leeway'),
pytest.param('someaudience', 'someissuer', None, None, lambda: time.time() - 10,
None,
id='expiration in past but within leeway'),
])
def test_decode_jwt_validation(aud, iss, nbf, iat, exp, expected_exception, private_key_pem,
public_key):
token = jwt.encode(_token_data(aud, 'subject', iss, iat, exp, nbf), private_key_pem, 'RS256')
if expected_exception is not None:
with pytest.raises(InvalidTokenError) as ite:
max_exp = exp_max_s_option(3600)
decode(token, public_key, algorithms=['RS256'], audience='someaudience',
issuer='someissuer', options=max_exp, leeway=60)
assert ite.match(expected_exception)
else:
max_exp = exp_max_s_option(3600)
decode(token, public_key, algorithms=['RS256'], audience='someaudience',
issuer='someissuer', options=max_exp, leeway=60)
def test_decode_jwt_invalid_key(private_key_pem):
# Encode with the test private key.
token = jwt.encode(_token_data('aud', 'subject', 'someissuer'), private_key_pem, 'RS256')
# Try to decode with a different public key.
another_public_key = RSA.generate(2048).publickey().exportKey('PEM')
with pytest.raises(InvalidTokenError) as ite:
max_exp = exp_max_s_option(3600)
decode(token, another_public_key, algorithms=['RS256'], audience='aud',
issuer='someissuer', options=max_exp, leeway=60)
assert ite.match('Signature verification failed')
def test_decode_jwt_invalid_algorithm(private_key_pem, public_key):
# Encode with the test private key.
token = jwt.encode(_token_data('aud', 'subject', 'someissuer'), private_key_pem, 'RS256')
# Attempt to decode but only with a different algorithm than that used.
with pytest.raises(InvalidAlgorithmError) as ite:
max_exp = exp_max_s_option(3600)
decode(token, public_key, algorithms=['ES256'], audience='aud',
issuer='someissuer', options=max_exp, leeway=60)
assert ite.match('are not whitelisted')
def test_jwk_dict_to_public_key(private_key, private_key_pem):
public_key = private_key.publickey()
jwk = RSAKey(key=private_key.publickey()).serialize()
converted = jwk_dict_to_public_key(jwk)
# Encode with the test private key.
token = jwt.encode(_token_data('aud', 'subject', 'someissuer'), private_key_pem, 'RS256')
# Decode with the converted key.
max_exp = exp_max_s_option(3600)
decode(token, converted, algorithms=['RS256'], audience='aud',
issuer='someissuer', options=max_exp, leeway=60)