Skip to content

Commit ca8efcf

Browse files
author
Michael Davis
committed
Refactor JWK algorithms
1 parent 67e3e5e commit ca8efcf

File tree

6 files changed

+65
-53
lines changed

6 files changed

+65
-53
lines changed

‎jose/jwk.py‎

Lines changed: 34 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -58,32 +58,31 @@ def construct(key_data, algorithm=None):
5858
if not algorithm:
5959
raise JWKError('Unable to find a algorithm for key: %s' % key_data)
6060

61-
if algorithm == ALGORITHMS.HS256:
62-
return HMACKey(key_data, HMACKey.SHA256)
61+
if algorithm in ALGORITHMS.HMAC:
62+
return HMACKey(key_data, algorithm)
6363

64-
if algorithm == ALGORITHMS.HS384:
65-
return HMACKey(key_data, HMACKey.SHA384)
64+
if algorithm in ALGORITHMS.RSA:
65+
return RSAKey(key_data, algorithm)
6666

67-
if algorithm == ALGORITHMS.HS512:
68-
return HMACKey(key_data, HMACKey.SHA512)
67+
if algorithm in ALGORITHMS.EC:
68+
return ECKey(key_data, algorithm)
6969

70-
if algorithm == ALGORITHMS.RS256:
71-
return RSAKey(key_data, RSAKey.SHA256)
7270

73-
if algorithm == ALGORITHMS.RS384:
74-
return RSAKey(key_data, RSAKey.SHA384)
71+
def get_algorithm_object(algorithm):
7572

76-
if algorithm == ALGORITHMS.RS512:
77-
return RSAKey(key_data, RSAKey.SHA512)
78-
79-
if algorithm == ALGORITHMS.ES256:
80-
return ECKey(key_data, ECKey.SHA256)
81-
82-
if algorithm == ALGORITHMS.ES384:
83-
return ECKey(key_data, ECKey.SHA384)
73+
algorithms = {
74+
ALGORITHMS.HS256: HMACKey.SHA256,
75+
ALGORITHMS.HS384: HMACKey.SHA384,
76+
ALGORITHMS.HS512: HMACKey.SHA512,
77+
ALGORITHMS.RS256: RSAKey.SHA256,
78+
ALGORITHMS.RS384: RSAKey.SHA384,
79+
ALGORITHMS.RS512: RSAKey.SHA512,
80+
ALGORITHMS.ES256: ECKey.SHA256,
81+
ALGORITHMS.ES384: ECKey.SHA384,
82+
ALGORITHMS.ES512: ECKey.SHA512,
83+
}
8484

85-
if algorithm == ALGORITHMS.ES512:
86-
return ECKey(key_data, ECKey.SHA512)
85+
return algorithms.get(algorithm, None)
8786

8887

8988
class Key(object):
@@ -111,15 +110,15 @@ class HMACKey(Key):
111110
SHA256 = hashlib.sha256
112111
SHA384 = hashlib.sha384
113112
SHA512 = hashlib.sha512
114-
valid_hash_algs = (SHA256, SHA384, SHA512)
113+
valid_hash_algs = ALGORITHMS.HMAC
115114

116115
prepared_key = None
117116
hash_alg = None
118117

119-
def __init__(self, key, hash_alg):
120-
if hash_alg not in self.valid_hash_algs:
121-
raise JWKError('hash_alg: %s is not a valid hash algorithm' % hash_alg)
122-
self.hash_alg = hash_alg
118+
def __init__(self, key, algorithm):
119+
if algorithm not in self.valid_hash_algs:
120+
raise JWKError('hash_alg: %s is not a valid hash algorithm' % algorithm)
121+
self.hash_alg = get_algorithm_object(algorithm)
123122

124123
if isinstance(key, dict):
125124
self.prepared_key = self._process_jwk(key)
@@ -173,16 +172,16 @@ class RSAKey(Key):
173172
SHA256 = Crypto.Hash.SHA256
174173
SHA384 = Crypto.Hash.SHA384
175174
SHA512 = Crypto.Hash.SHA512
176-
valid_hash_algs = (SHA256, SHA384, SHA512)
175+
valid_hash_algs = ALGORITHMS.RSA
177176

178177
prepared_key = None
179178
hash_alg = None
180179

181-
def __init__(self, key, hash_alg):
180+
def __init__(self, key, algorithm):
182181

183-
if hash_alg not in self.valid_hash_algs:
184-
raise JWKError('hash_alg: %s is not a valid hash algorithm' % hash_alg)
185-
self.hash_alg = hash_alg
182+
if algorithm not in self.valid_hash_algs:
183+
raise JWKError('hash_alg: %s is not a valid hash algorithm' % algorithm)
184+
self.hash_alg = get_algorithm_object(algorithm)
186185

187186
if isinstance(key, _RSAKey):
188187
self.prepared_key = key
@@ -240,7 +239,7 @@ class ECKey(Key):
240239
SHA256 = hashlib.sha256
241240
SHA384 = hashlib.sha384
242241
SHA512 = hashlib.sha512
243-
valid_hash_algs = (SHA256, SHA384, SHA512)
242+
valid_hash_algs = ALGORITHMS.EC
244243

245244
curve_map = {
246245
SHA256: ecdsa.curves.NIST256p,
@@ -252,10 +251,10 @@ class ECKey(Key):
252251
hash_alg = None
253252
curve = None
254253

255-
def __init__(self, key, hash_alg):
256-
if hash_alg not in self.valid_hash_algs:
257-
raise JWKError('hash_alg: %s is not a valid hash algorithm' % hash_alg)
258-
self.hash_alg = hash_alg
254+
def __init__(self, key, algorithm):
255+
if algorithm not in self.valid_hash_algs:
256+
raise JWKError('hash_alg: %s is not a valid hash algorithm' % algorithm)
257+
self.hash_alg = get_algorithm_object(algorithm)
259258

260259
self.curve = self.curve_map.get(self.hash_alg)
261260

‎requirements-dev.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,4 +11,4 @@ ecdsa==0.13
1111
wsgiref==0.1.2
1212

1313
-r requirements.txt
14-
-r requirements-rtd.txt
14+
-r requirements-rtd.txt

‎tests/algorithms/test_EC.py‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11

2-
from jose.jwk import ECKey
2+
from jose.constants import ALGORITHMS
33
from jose.exceptions import JOSEError
4+
from jose.jwk import ECKey
45

56
import ecdsa
67
import pytest
@@ -16,14 +17,14 @@ class TestECAlgorithm:
1617

1718
def test_EC_key(self):
1819
key = ecdsa.SigningKey.from_pem(private_key)
19-
ECKey(key, ECKey.SHA256)
20+
ECKey(key, ALGORITHMS.ES256)
2021

2122
def test_string_secret(self):
2223
key = 'secret'
2324
with pytest.raises(JOSEError):
24-
ECKey(key, ECKey.SHA256)
25+
ECKey(key, ALGORITHMS.ES256)
2526

2627
def test_object(self):
2728
key = object()
2829
with pytest.raises(JOSEError):
29-
ECKey(key, ECKey.SHA256)
30+
ECKey(key, ALGORITHMS.ES256)

‎tests/algorithms/test_HMAC.py‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11

2-
from jose.jwk import HMACKey
2+
from jose.constants import ALGORITHMS
33
from jose.exceptions import JOSEError
4+
from jose.jwk import HMACKey
45

56
import pytest
67

@@ -9,17 +10,17 @@ class TestHMACAlgorithm:
910

1011
def test_non_string_key(self):
1112
with pytest.raises(JOSEError):
12-
HMACKey(object(), HMACKey.SHA256)
13+
HMACKey(object(), ALGORITHMS.HS256)
1314

1415
def test_RSA_key(self):
1516
key = "-----BEGIN PUBLIC KEY-----"
1617
with pytest.raises(JOSEError):
17-
HMACKey(key, HMACKey.SHA256)
18+
HMACKey(key, ALGORITHMS.HS256)
1819

1920
key = "-----BEGIN CERTIFICATE-----"
2021
with pytest.raises(JOSEError):
21-
HMACKey(key, HMACKey.SHA256)
22+
HMACKey(key, ALGORITHMS.HS256)
2223

2324
key = "ssh-rsa"
2425
with pytest.raises(JOSEError):
25-
HMACKey(key, HMACKey.SHA256)
26+
HMACKey(key, ALGORITHMS.HS256)

‎tests/algorithms/test_RSA.py‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,19 @@
11

2+
import sys
3+
24
from jose.jwk import RSAKey
5+
from jose.constants import ALGORITHMS
36
from jose.exceptions import JOSEError
47

58
from Crypto.PublicKey import RSA
69

710
import pytest
811

12+
# Deal with integer compatibilities between Python 2 and 3.
13+
# Using `from builtins import int` is not supported on AppEngine.
14+
if sys.version_info > (3,):
15+
long = int
16+
917
private_key = """-----BEGIN RSA PRIVATE KEY-----
1018
MIIJKwIBAAKCAgEAtSKfSeI0fukRIX38AHlKB1YPpX8PUYN2JdvfM+XjNmLfU1M7
1119
4N0VmdzIX95sneQGO9kC2xMIE+AIlt52Yf/KgBZggAlS9Y0Vx8DsSL2HvOjguAdX
@@ -62,14 +70,18 @@
6270
class TestRSAAlgorithm:
6371

6472
def test_RSA_key(self):
65-
RSAKey(private_key, RSAKey.SHA256)
73+
RSAKey(private_key, ALGORITHMS.RS256)
74+
75+
def test_RSA_key_instance(self):
76+
key = RSA.construct((long(26057131595212989515105618545799160306093557851986992545257129318694524535510983041068168825614868056510242030438003863929818932202262132630250203397069801217463517914103389095129323580576852108653940669240896817348477800490303630912852266209307160550655497615975529276169196271699168537716821419779900117025818140018436554173242441334827711966499484119233207097432165756707507563413323850255548329534279691658369466534587631102538061857114141268972476680597988266772849780811214198186940677291891818952682545840788356616771009013059992237747149380197028452160324144544057074406611859615973035412993832273216732343819), long(65537)))
77+
RSAKey(key, ALGORITHMS.RS256)
6678

6779
def test_string_secret(self):
6880
key = 'secret'
6981
with pytest.raises(JOSEError):
70-
RSAKey(key, RSAKey.SHA256)
82+
RSAKey(key, ALGORITHMS.RS256)
7183

7284
def test_object(self):
7385
key = object()
7486
with pytest.raises(JOSEError):
75-
RSAKey(key, RSAKey.SHA256)
87+
RSAKey(key, ALGORITHMS.RS256)

‎tests/test_jwk.py‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
import pytest
66

7+
78
hmac_key = {
89
"kty": "oct",
910
"kid": "018c0ae5-4d9b-471b-bfd6-eef314bc7037",
@@ -58,13 +59,13 @@ def test_invalid_hash_alg(self):
5859
def test_invalid_jwk(self):
5960

6061
with pytest.raises(JWKError):
61-
key = jwk.HMACKey(rsa_key, 'HS512')
62+
key = jwk.HMACKey(rsa_key, 'HS256')
6263

6364
with pytest.raises(JWKError):
64-
key = jwk.HMACKey(hmac_key, 'RS512')
65+
key = jwk.RSAKey(hmac_key, 'RS256')
6566

6667
with pytest.raises(JWKError):
67-
key = jwk.HMACKey(rsa_key, 'EC512')
68+
key = jwk.ECKey(rsa_key, 'ES256')
6869

6970
def test_RSAKey_errors(self):
7071

@@ -90,8 +91,6 @@ def test_RSAKey_errors(self):
9091
with pytest.raises(JWKError):
9192
key = jwk.RSAKey(rsa_key, 'RS256')
9293

93-
94-
9594
def test_construct_from_jwk(self):
9695

9796
hmac_key = {

0 commit comments

Comments
 (0)