Move ecdsa python implementation into separate backend. Key.to_pem returns six.binary_type - str on PY2, bytes on PY3.

Gasper Zejn · Gasper Zejn · commit 4b59bb7048fa · 2017-04-13T18:23:14.000+02:00
diff --git a/jose/backends/__init__.py b/jose/backends/__init__.py
@@ -1,5 +1,10 @@
 
 try:
-    from .pycrypto_backend import RSAKey
+    from jose.backends.pycrypto_backend import RSAKey
 except ImportError:
-    from .cryptography_backend import CryptographyRSAKey as RSAKey
+    from jose.backends.cryptography_backend import CryptographyRSAKey as RSAKey
+
+try:
+    from jose.backends.ecdsa_backend import ECKey
+except ImportError:
+    from jose.backends.cryptography_backend import CryptographyECKey as ECKey
diff --git a/jose/backends/cryptography_backend.py b/jose/backends/cryptography_backend.py
@@ -52,9 +52,9 @@ def __init__(self, key, algorithm, cryptography_backend=default_backend):
             return
 
         if isinstance(key, six.string_types):
-            if isinstance(key, six.text_type):
-                key = key.encode('utf-8')
+            key = key.encode('utf-8')
 
+        if isinstance(key, six.binary_type):
             # Attempt to load key. We don't know if it's
             # a Public Key or a Private Key, so we try
             # the Public Key first.
@@ -109,13 +109,13 @@ def to_pem(self):
                 encoding=serialization.Encoding.PEM,
                 format=serialization.PublicFormat.SubjectPublicKeyInfo
             )
-            return pem.decode('utf-8')
+            return pem
         pem = self.prepared_key.private_bytes(
             encoding=serialization.Encoding.PEM,
             format=serialization.PrivateFormat.TraditionalOpenSSL,
             encryption_algorithm=serialization.NoEncryption()
         )
-        return pem.decode('utf-8')
+        return pem
 
 
 class CryptographyRSAKey(Key):
@@ -146,9 +146,9 @@ def __init__(self, key, algorithm, cryptography_backend=default_backend):
             return
 
         if isinstance(key, six.string_types):
-            if isinstance(key, six.text_type):
-                key = key.encode('utf-8')
+            key = key.encode('utf-8')
 
+        if isinstance(key, six.binary_type):
             try:
                 try:
                     key = load_pem_public_key(key, self.cryptography_backend())
@@ -203,10 +203,10 @@ def to_pem(self):
             return self.prepared_key.public_bytes(
                 encoding=serialization.Encoding.PEM,
                 format=serialization.PublicFormat.SubjectPublicKeyInfo
-            ).decode('utf-8')
+            )
 
         return self.prepared_key.private_bytes(
             encoding=serialization.Encoding.PEM,
             format=serialization.PrivateFormat.TraditionalOpenSSL,
             encryption_algorithm=serialization.NoEncryption()
-        ).decode('utf-8')
+        )
diff --git a/jose/backends/ecdsa_backend.py b/jose/backends/ecdsa_backend.py
@@ -0,0 +1,101 @@
+import hashlib
+import six
+
+from jose.backends.base import Key
+import ecdsa
+
+from jose.constants import ALGORITHMS
+from jose.exceptions import JWKError
+from jose.utils import base64_to_long
+
+
+class ECKey(Key):
+    """
+    Performs signing and verification operations using
+    ECDSA and the specified hash function
+
+    This class requires the ecdsa package to be installed.
+
+    This is based off of the implementation in PyJWT 0.3.2
+    """
+    SHA256 = hashlib.sha256
+    SHA384 = hashlib.sha384
+    SHA512 = hashlib.sha512
+
+    CURVE_MAP = {
+        SHA256: ecdsa.curves.NIST256p,
+        SHA384: ecdsa.curves.NIST384p,
+        SHA512: ecdsa.curves.NIST521p,
+    }
+
+    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 = {
+            ALGORITHMS.ES256: self.SHA256,
+            ALGORITHMS.ES384: self.SHA384,
+            ALGORITHMS.ES512: self.SHA512
+        }.get(algorithm)
+        self._algorithm = algorithm
+
+        self.curve = self.CURVE_MAP.get(self.hash_alg)
+
+        if isinstance(key, (ecdsa.SigningKey, ecdsa.VerifyingKey)):
+            self.prepared_key = key
+            return
+
+        if isinstance(key, dict):
+            self.prepared_key = self._process_jwk(key)
+            return
+
+        if isinstance(key, six.string_types):
+            key = key.encode('utf-8')
+
+        if isinstance(key, six.binary_type):
+            # Attempt to load key. We don't know if it's
+            # a Signing Key or a Verifying Key, so we try
+            # the Verifying Key first.
+            try:
+                key = ecdsa.VerifyingKey.from_pem(key)
+            except ecdsa.der.UnexpectedDER:
+                key = ecdsa.SigningKey.from_pem(key)
+            except Exception as e:
+                raise JWKError(e)
+
+            self.prepared_key = key
+            return
+
+        raise JWKError('Unable to parse an ECKey from key: %s' % key)
+
+    def _process_jwk(self, jwk_dict):
+        if not jwk_dict.get('kty') == 'EC':
+            raise JWKError("Incorrect key type.  Expected: 'EC', Recieved: %s" % jwk_dict.get('kty'))
+
+        x = base64_to_long(jwk_dict.get('x'))
+        y = base64_to_long(jwk_dict.get('y'))
+
+        if not ecdsa.ecdsa.point_is_valid(self.curve.generator, x, y):
+            raise JWKError("Point: %s, %s is not a valid point" % (x, y))
+
+        point = ecdsa.ellipticcurve.Point(self.curve.curve, x, y, self.curve.order)
+        verifying_key = ecdsa.keys.VerifyingKey.from_public_point(point, self.curve)
+
+        return verifying_key
+
+    def sign(self, msg):
+        return self.prepared_key.sign(msg, hashfunc=self.hash_alg, sigencode=ecdsa.util.sigencode_string)
+
+    def verify(self, msg, sig):
+        try:
+            return self.prepared_key.verify(sig, msg, hashfunc=self.hash_alg, sigdecode=ecdsa.util.sigdecode_string)
+        except:
+            return False
+
+    def public_key(self):
+        if isinstance(self.prepared_key, ecdsa.VerifyingKey):
+            return self
+        return self.__class__(self.prepared_key.get_verifying_key(), self._algorithm)
+
+    def to_pem(self):
+        return self.prepared_key.to_pem()
diff --git a/jose/backends/pycrypto_backend.py b/jose/backends/pycrypto_backend.py
@@ -9,7 +9,7 @@
 from Crypto.Util.asn1 import DerSequence
 
 from jose.backends.base import Key
-from jose.jwk import base64_to_long
+from jose.utils import base64_to_long
 from jose.constants import ALGORITHMS
 from jose.exceptions import JWKError
 from jose.utils import base64url_decode
@@ -55,9 +55,9 @@ def __init__(self, key, algorithm):
             return
 
         if isinstance(key, six.string_types):
-            if isinstance(key, six.text_type):
-                key = key.encode('utf-8')
+            key = key.encode('utf-8')
 
+        if isinstance(key, six.binary_type):
             if key.startswith(b'-----BEGIN CERTIFICATE-----'):
                 try:
                     self._process_cert(key)
diff --git a/jose/jwk.py b/jose/jwk.py
@@ -3,14 +3,12 @@
 import hmac
 import six
 
-import ecdsa
-
 from jose.constants import ALGORITHMS
 from jose.exceptions import JWKError
-from jose.utils import base64url_decode, base64_to_long
+from jose.utils import base64url_decode
 from jose.utils import constant_time_string_compare
 from jose.backends.base import Key
-from jose.backends import RSAKey
+from jose.backends import RSAKey, ECKey
 
 
 def get_key(algorithm):
@@ -122,90 +120,3 @@ def sign(self, msg):
 
     def verify(self, msg, sig):
         return constant_time_string_compare(sig, self.sign(msg))
-
-
-class ECKey(Key):
-    """
-    Performs signing and verification operations using
-    ECDSA and the specified hash function
-
-    This class requires the ecdsa package to be installed.
-
-    This is based off of the implementation in PyJWT 0.3.2
-    """
-    SHA256 = hashlib.sha256
-    SHA384 = hashlib.sha384
-    SHA512 = hashlib.sha512
-
-    CURVE_MAP = {
-        SHA256: ecdsa.curves.NIST256p,
-        SHA384: ecdsa.curves.NIST384p,
-        SHA512: ecdsa.curves.NIST521p,
-    }
-
-    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._algorithm = algorithm
-
-        self.curve = self.CURVE_MAP.get(self.hash_alg)
-
-        if isinstance(key, (ecdsa.SigningKey, ecdsa.VerifyingKey)):
-            self.prepared_key = key
-            return
-
-        if isinstance(key, dict):
-            self.prepared_key = self._process_jwk(key)
-            return
-
-        if isinstance(key, six.string_types):
-            if isinstance(key, six.text_type):
-                key = key.encode('utf-8')
-
-            # Attempt to load key. We don't know if it's
-            # a Signing Key or a Verifying Key, so we try
-            # the Verifying Key first.
-            try:
-                key = ecdsa.VerifyingKey.from_pem(key)
-            except ecdsa.der.UnexpectedDER:
-                key = ecdsa.SigningKey.from_pem(key)
-            except Exception as e:
-                raise JWKError(e)
-
-            self.prepared_key = key
-            return
-
-        raise JWKError('Unable to parse an ECKey from key: %s' % key)
-
-    def _process_jwk(self, jwk_dict):
-        if not jwk_dict.get('kty') == 'EC':
-            raise JWKError("Incorrect key type.  Expected: 'EC', Recieved: %s" % jwk_dict.get('kty'))
-
-        x = base64_to_long(jwk_dict.get('x'))
-        y = base64_to_long(jwk_dict.get('y'))
-
-        if not ecdsa.ecdsa.point_is_valid(self.curve.generator, x, y):
-            raise JWKError("Point: %s, %s is not a valid point" % (x, y))
-
-        point = ecdsa.ellipticcurve.Point(self.curve.curve, x, y, self.curve.order)
-        verifying_key = ecdsa.keys.VerifyingKey.from_public_point(point, self.curve)
-
-        return verifying_key
-
-    def sign(self, msg):
-        return self.prepared_key.sign(msg, hashfunc=self.hash_alg, sigencode=ecdsa.util.sigencode_string)
-
-    def verify(self, msg, sig):
-        try:
-            return self.prepared_key.verify(sig, msg, hashfunc=self.hash_alg, sigdecode=ecdsa.util.sigdecode_string)
-        except:
-            return False
-
-    def public_key(self):
-        if isinstance(self.prepared_key, ecdsa.VerifyingKey):
-            return self
-        return self.__class__(self.prepared_key.get_verifying_key(), self._algorithm)
-
-    def to_pem(self):
-        return self.prepared_key.to_pem()
diff --git a/tests/algorithms/test_RSA.py b/tests/algorithms/test_RSA.py
@@ -108,7 +108,7 @@ def test_RSA_key_instance(self):
 
         pubkey = CryptographyRSAKey(key, ALGORITHMS.RS256)
         pem = pubkey.to_pem()
-        assert pem.startswith('-----BEGIN PUBLIC KEY-----')
+        assert pem.startswith(b'-----BEGIN PUBLIC KEY-----')
 
     def test_RSA_jwk(self):
         d = {
diff --git a/tests/algorithms/test_cryptography_EC.py b/tests/algorithms/test_cryptography_EC.py
@@ -7,7 +7,7 @@
 import ecdsa
 import pytest
 
-private_key = """-----BEGIN EC PRIVATE KEY-----
+private_key = b"""-----BEGIN EC PRIVATE KEY-----
 MHQCAQEEIIAK499svJugZZfsTsgL2tc7kH/CpzQbkr4g55CEWQyPoAcGBSuBBAAK
 oUQDQgAEsOnVqWVPfjte2nI0Ay3oTZVehCUtH66nJM8z6flUluHxhLG8ZTTCkJAZ
 W6xQdXHfqGUy3Dx40NDhgTaM8xAdSw==
@@ -20,8 +20,6 @@ def test_EC_key(self):
         key = ecdsa.SigningKey.from_pem(private_key)
         k = CryptographyECKey(key, ALGORITHMS.ES256)
 
-        print(repr(k.to_pem().strip()))
-        print(repr(private_key.strip()))
         assert k.to_pem().strip() == private_key.strip()
         public_pem = k.public_key().to_pem()
         public_key = CryptographyECKey(public_pem, ALGORITHMS.ES256)
@@ -42,7 +40,7 @@ def test_cryptography_EC_key(self):
 
     def test_signing_parity(self):
         key1 = ECKey(private_key, ALGORITHMS.ES256)
-        public_key = key1.prepared_key.get_verifying_key().to_pem().decode('utf-8')
+        public_key = key1.public_key().to_pem()
         vkey1 = ECKey(public_key, ALGORITHMS.ES256)
         key2 = CryptographyECKey(private_key, ALGORITHMS.ES256)
         vkey2 = CryptographyECKey(public_key, ALGORITHMS.ES256)
