cleaned jwk, added class checking on key, added tests

Author: Marko Mrdjenovic

jose/constants.py

@@ -45,9 +45,15 @@ def get_key(self, algorithm):
             return RSAKey
         elif algorithm in self.EC:
             return ECKey
+        return None
 
     def register_key(self, algorithm, key_class):
-        self.KEYS[algorithm] = key_class
-        self.SUPPORTED.add(algorithm)
+        from jose.jwk import Key
+        if issubclass(key_class, Key):
+            self.KEYS[algorithm] = key_class
+            self.SUPPORTED.add(algorithm)
+            return True
+        else:
+            return False
 
 ALGORITHMS = Algorithms()

jose/jwk.py

@@ -87,11 +87,8 @@ class Key(object):
     """
     A simple interface for implementing JWK keys.
     """
-    prepared_key = None
-    hash_alg = None
-
-    def _process_jwk(self, jwk_dict):
-        raise NotImplementedError()
+    def __init__(self, key, algorithm):
+        pass
 
     def sign(self, msg):
         raise NotImplementedError()
@@ -109,9 +106,6 @@ class HMACKey(Key):
     SHA384 = hashlib.sha384
     SHA512 = hashlib.sha512
 
-    prepared_key = None
-    hash_alg = None
-
     def __init__(self, key, algorithm):
         if algorithm not in ALGORITHMS.HMAC:
             raise JWKError('hash_alg: %s is not a valid hash algorithm' % algorithm)
@@ -170,9 +164,6 @@ class RSAKey(Key):
     SHA384 = Crypto.Hash.SHA384
     SHA512 = Crypto.Hash.SHA512
 
-    prepared_key = None
-    hash_alg = None
-
     def __init__(self, key, algorithm):
 
         if algorithm not in ALGORITHMS.RSA:
@@ -236,7 +227,7 @@ def verify(self, msg, sig):
         try:
             return PKCS1_v1_5.new(self.prepared_key).verify(self.hash_alg.new(msg), sig)
         except Exception as e:
-            raise JWKError(e)
+            return False
 
 
 class ECKey(Key):
@@ -252,22 +243,18 @@ class ECKey(Key):
     SHA384 = hashlib.sha384
     SHA512 = hashlib.sha512
 
-    curve_map = {
+    CURVE_MAP = {
         SHA256: ecdsa.curves.NIST256p,
         SHA384: ecdsa.curves.NIST384p,
         SHA512: ecdsa.curves.NIST521p,
     }
 
-    prepared_key = None
-    hash_alg = None
-    curve = None
-
     def __init__(self, key, algorithm):
         if algorithm not in ALGORITHMS.EC:
             raise JWKError('hash_alg: %s is not a valid hash algorithm' % algorithm)
         self.hash_alg = get_algorithm_object(algorithm)
 
-        self.curve = self.curve_map.get(self.hash_alg)
+        self.curve = self.CURVE_MAP.get(self.hash_alg)
 
         if isinstance(key, (ecdsa.SigningKey, ecdsa.VerifyingKey)):
             self.prepared_key = key

tests/algorithms/test_base.py

@@ -1,25 +1,37 @@
 
-# from jose.jwk import Key
-# from jose.exceptions import JOSEError
+from jose.jwk import Key, HMACKey, RSAKey, ECKey
+from jose.constants import ALGORITHMS
 
-# import pytest
+import pytest
 
 
-# @pytest.fixture
-# def alg():
-#     return Key()
+@pytest.fixture
+def alg():
+    return Key("key", "ALG")
 
 
-# class TestBaseAlgorithm:
+class TestBaseAlgorithm:
 
-#     def test_prepare_key_is_interface(self, alg):
-#         with pytest.raises(JOSEError):
-#             alg.prepare_key('secret')
+    def test_sign_is_interface(self, alg):
+        with pytest.raises(NotImplementedError):
+            alg.sign('msg')
 
-#     def test_sign_is_interface(self, alg):
-#         with pytest.raises(JOSEError):
-#             alg.sign('msg', 'secret')
+    def test_verify_is_interface(self, alg):
+        with pytest.raises(NotImplementedError):
+            alg.verify('msg', 'sig')
 
-#     def test_verify_is_interface(self, alg):
-#         with pytest.raises(JOSEError):
-#             alg.verify('msg', 'secret', 'sig')
+
+class TestAlgorithms:
+
+    def test_get_key(self):
+        assert ALGORITHMS.get_key("HS256") == HMACKey
+        assert ALGORITHMS.get_key("RS256") == RSAKey
+        assert ALGORITHMS.get_key("ES256") == ECKey
+
+        assert ALGORITHMS.get_key("NONEXISTENT") == None
+    
+    def test_register_key(self):
+        assert ALGORITHMS.register_key("ALG", Key) == True
+        assert ALGORITHMS.get_key("ALG") == Key
+
+        assert ALGORITHMS.register_key("ALG", object) == False

tests/test_jwk.py

@@ -35,10 +35,7 @@ class TestJWK:
 
     def test_interface(self):
 
-        key = jwk.Key()
-
-        with pytest.raises(NotImplementedError):
-            key._process_jwk(None)
+        key = jwk.Key("key", "ALG")
 
         with pytest.raises(NotImplementedError):
             key.sign('')
@@ -115,3 +112,6 @@ def test_construct_from_jwk_missing_alg(self):
 
         with pytest.raises(JWKError):
             key = jwk.construct(hmac_key)
+
+        with pytest.raises(JWKError):
+            key = jwk.construct("key", algorithm="NONEXISTENT")