add ecdh class (#149)

Add support for ECDH operations.

Author: merlokk

README.md

@@ -466,3 +466,18 @@ comp_str = '022799c0d0ee09772fdd337d4f28dc155581951d07082fb19a38aa396b67e77759'
 vk = VerifyingKey.from_string(bytearray.fromhex(comp_str), curve=NIST256p)
 print(vk.to_string("uncompressed").hex())
 ```
+
+ECDH key exchange with remote party
+
+```
+from ecdsa import ECDH, NIST256p
+
+ecdh = ECDH(curve=NIST256p)
+ecdh.generate_private_key()
+local_public_key = ecdh.get_public_key()
+#send `local_public_key` to remote party and receive `remote_public_key` from remote party
+with open("remote_public_key.pem") as e:
+    remote_public_key = e.read()
+ecdh.load_received_public_key_pem(remote_public_key)
+secret = ecdh.generate_sharedsecret_bytes()
+```

src/ecdsa/__init__.py

@@ -3,6 +3,8 @@
 from .curves import NIST192p, NIST224p, NIST256p, NIST384p, NIST521p,\
         SECP256k1, BRAINPOOLP160r1, BRAINPOOLP192r1, BRAINPOOLP224r1,\
         BRAINPOOLP256r1, BRAINPOOLP320r1, BRAINPOOLP384r1, BRAINPOOLP512r1
+from .ecdh import ECDH, NoKeyError, NoCurveError, InvalidCurveError, \
+    InvalidSharedSecretError
 from .der import UnexpectedDER
 
 # This code comes from http://github.com/warner/python-ecdsa
@@ -14,7 +16,8 @@
            "test_pyecdsa", "util", "six"]
 
 _hush_pyflakes = [SigningKey, VerifyingKey, BadSignatureError, BadDigestError,
-                  MalformedPointError, UnexpectedDER,
+                  MalformedPointError, UnexpectedDER, InvalidCurveError,
+                  NoKeyError, InvalidSharedSecretError, ECDH, NoCurveError,
                   NIST192p, NIST224p, NIST256p, NIST384p, NIST521p, SECP256k1,
                   BRAINPOOLP160r1, BRAINPOOLP192r1, BRAINPOOLP224r1,
                   BRAINPOOLP256r1, BRAINPOOLP320r1, BRAINPOOLP384r1,

src/ecdsa/ecdh.py

@@ -0,0 +1,306 @@
+"""
+Class for performing Elliptic-curve Diffie-Hellman (ECDH) operations.
+"""
+
+from .util import number_to_string
+from .ellipticcurve import INFINITY
+from .keys import SigningKey, VerifyingKey
+
+
+__all__ = ["ECDH", "NoKeyError", "NoCurveError", "InvalidCurveError",
+           "InvalidSharedSecretError"]
+
+
+class NoKeyError(Exception):
+    """ECDH. Key not found but it is needed for operation."""
+
+    pass
+
+
+class NoCurveError(Exception):
+    """ECDH. Curve not set but it is needed for operation."""
+
+    pass
+
+
+class InvalidCurveError(Exception):
+    """ECDH. Raised in case the public and private keys use different curves."""
+
+    pass
+
+
+class InvalidSharedSecretError(Exception):
+    """ECDH. Raised in case the shared secret we obtained is an INFINITY."""
+
+    pass
+
+
+class ECDH(object):
+    """
+    Elliptic-curve Diffie-Hellman (ECDH). A key agreement protocol.
+    
+    Allows two parties, each having an elliptic-curve public-private key 
+    pair, to establish a shared secret over an insecure channel
+    """""
+
+    def __init__(self, curve=None, private_key=None, public_key=None):
+        """
+        ECDH init.
+
+        Call can be initialised without parameters, then the first operation
+        (loading either key) will set the used curve.
+        All parameters must be ultimately set before shared secret
+        calculation will be allowed.
+
+        :param curve: curve for operations
+        :type curve: Curve
+        :param private_key: `my` private key for ECDH
+        :type private_key: SigningKey
+        :param public_key:  `their` public key for ECDH
+        :type public_key: VerifyingKey
+       """
+        self.curve = curve
+        self.private_key = None
+        self.public_key = None
+        if private_key:
+            self.load_private_key(private_key)
+        if public_key:
+            self.load_received_public_key(public_key)
+
+    def _get_shared_secret(self, remote_public_key):
+        if not self.private_key:
+            raise NoKeyError(
+                "Private key needs to be set to create shared secret")
+        if not self.public_key:
+            raise NoKeyError(
+                "Public key needs to be set to create shared secret")
+        if not (self.private_key.curve == self.curve == remote_public_key.curve):
+            raise InvalidCurveError(
+                "Curves for public key and private key is not equal.")
+
+        # shared secret = PUBKEYtheirs * PRIVATEKEYours
+        result = remote_public_key.pubkey.point * self.private_key.privkey.secret_multiplier
+        if result == INFINITY:
+            raise InvalidSharedSecretError(
+                "Invalid shared secret (INFINITY).")
+
+        return result.x()
+
+    def set_curve(self, key_curve):
+        """
+        Set the working curve for ecdh operations.
+
+        :param key_curve: curve from `curves` module
+        :type key_curve: Curve
+        """
+        self.curve = key_curve
+
+    def generate_private_key(self):
+        """
+        Generate local private key for ecdh operation with curve that was set.
+
+        :raises NoCurveError: Curve must be set before key generation.
+
+        :return: public (verifying) key from this private key.
+        :rtype: VerifyingKey object
+        """
+        if not self.curve:
+            raise NoCurveError("Curve must be set prior to key generation.")
+        return self.load_private_key(SigningKey.generate(curve=self.curve))
+
+    def load_private_key(self, private_key):
+        """
+        Load private key from SigningKey (keys.py) object.
+
+        Needs to have the same curve as was set with set_curve method.
+        If curve is not set - it sets from this SigningKey
+
+        :param private_key: Initialised SigningKey class
+        :type private_key: SigningKey
+
+        :raises InvalidCurveError: private_key curve not the same as self.curve
+
+        :return: public (verifying) key from this private key.
+        :rtype: VerifyingKey object
+        """
+        if not self.curve:
+            self.curve = private_key.curve
+        if self.curve != private_key.curve:
+            raise InvalidCurveError("Curve mismatch.")
+        self.private_key = private_key
+        return self.private_key.get_verifying_key()
+
+    def load_private_key_bytes(self, private_key):
+        """
+        Load private key from byte string.
+
+        Uses current curve and checks if the provided key matches
+        the curve of ECDH key agreement.
+        Key loads via from_string method of SigningKey class
+
+        :param private_key: private key in bytes string format
+        :type private_key: :term:`bytes-like object`
+
+        :raises NoCurveError: Curve must be set before loading.
+
+        :return: public (verifying) key from this private key.
+        :rtype: VerifyingKey object
+        """
+        if not self.curve:
+            raise NoCurveError("Curve must be set prior to key load.")
+        return self.load_private_key(
+            SigningKey.from_string(private_key, curve=self.curve))
+
+    def load_private_key_der(self, private_key_der):
+        """
+        Load private key from DER byte string.
+
+        Compares the curve of the DER-encoded key with the ECDH set curve,
+        uses the former if unset.
+
+        Note, the only DER format supported is the RFC5915
+        Look at keys.py:SigningKey.from_der()
+
+        :param private_key_der: string with the DER encoding of private ECDSA key
+        :type private_key_der: string
+
+        :raises InvalidCurveError: private_key curve not the same as self.curve
+
+        :return: public (verifying) key from this private key.
+        :rtype: VerifyingKey object
+        """
+        return self.load_private_key(SigningKey.from_der(private_key_der))
+
+    def load_private_key_pem(self, private_key_pem):
+        """
+        Load private key from PEM string.
+
+        Compares the curve of the DER-encoded key with the ECDH set curve,
+        uses the former if unset.
+
+        Note, the only PEM format supported is the RFC5915
+        Look at keys.py:SigningKey.from_pem()
+        it needs to have `EC PRIVATE KEY` section
+
+        :param private_key_pem: string with PEM-encoded private ECDSA key
+        :type private_key_pem: string
+
+        :raises InvalidCurveError: private_key curve not the same as self.curve
+
+        :return: public (verifying) key from this private key.
+        :rtype: VerifyingKey object
+        """
+        return self.load_private_key(SigningKey.from_pem(private_key_pem))
+
+    def get_public_key(self):
+        """
+        Provides a public key that matches the local private key.
+
+        Needs to be sent to the remote party.
+
+        :return: public (verifying) key from local private key.
+        :rtype: VerifyingKey object
+       """
+        return self.private_key.get_verifying_key()
+
+    def load_received_public_key(self, public_key):
+        """
+        Load public key from VerifyingKey (keys.py) object.
+
+        Needs to have the same curve as set as current for ecdh operation.
+        If curve is not set - it sets it from VerifyingKey.
+
+        :param public_key: Initialised VerifyingKey class
+        :type public_key: VerifyingKey
+
+        :raises InvalidCurveError: public_key curve not the same as self.curve
+        """
+        if not self.curve:
+            self.curve = public_key.curve
+        if self.curve != public_key.curve:
+            raise InvalidCurveError("Curve mismatch.")
+        self.public_key = public_key
+
+    def load_received_public_key_bytes(self, public_key_str):
+        """
+        Load public key from byte string.
+
+        Uses current curve and checks if key length corresponds to
+        the current curve.
+        Key loads via from_string method of VerifyingKey class
+
+        :param public_key_str: public key in bytes string format
+        :type public_key_str: :term:`bytes-like object`
+        """
+        return self.load_received_public_key(
+            VerifyingKey.from_string(public_key_str, self.curve))
+
+    def load_received_public_key_der(self, public_key_der):
+        """
+        Load public key from DER byte string.
+
+        Compares the curve of the DER-encoded key with the ECDH set curve,
+        uses the former if unset.
+
+        Note, the only DER format supported is the RFC5912
+        Look at keys.py:VerifyingKey.from_der()
+
+        :param public_key_der: string with the DER encoding of public ECDSA key
+        :type public_key_der: string
+
+        :raises InvalidCurveError: public_key curve not the same as self.curve
+        """
+        return self.load_received_public_key(VerifyingKey.from_der(public_key_der))
+
+    def load_received_public_key_pem(self, public_key_pem):
+        """
+        Load public key from PEM string.
+
+        Compares the curve of the PEM-encoded key with the ECDH set curve,
+        uses the former if unset.
+
+        Note, the only PEM format supported is the RFC5912
+        Look at keys.py:VerifyingKey.from_pem()
+
+        :param public_key_pem: string with PEM-encoded public ECDSA key
+        :type public_key_pem: string
+
+        :raises InvalidCurveError: public_key curve not the same as self.curve
+        """
+        return self.load_received_public_key(VerifyingKey.from_pem(public_key_pem))
+
+    def generate_sharedsecret_bytes(self):
+        """
+        Generate shared secret from local private key and remote public key.
+
+        The objects needs to have both private key and received public key
+        before generation is allowed.
+
+        :raises InvalidCurveError: public_key curve not the same as self.curve
+        :raises NoKeyError: public_key or private_key is not set
+
+        :return: shared secret
+        :rtype: byte string
+        """
+        return number_to_string(
+                    self.generate_sharedsecret(),
+                    self.private_key.curve.order)
+
+    def generate_sharedsecret(self):
+        """
+        Generate shared secret from local private key and remote public key.
+
+        The objects needs to have both private key and received public key
+        before generation is allowed.
+
+        It's the same for local and remote party.
+        shared secret(local private key, remote public key ) ==
+                shared secret (local public key, remote private key)
+
+        :raises InvalidCurveError: public_key curve not the same as self.curve
+        :raises NoKeyError: public_key or private_key is not set
+
+        :return: shared secret
+        :rtype: int
+        """
+        return self._get_shared_secret(self.public_key)