add EC type hinting (#5729)

4 files changed, 160 insertions, 70 deletions

diff --git a/src/cryptography/hazmat/backends/openssl/ec.py b/src/cryptography/hazmat/backends/openssl/ec.py
index 8977f489f..3f54d7281 100644
--- a/src/cryptography/hazmat/backends/openssl/ec.py
+++ b/src/cryptography/hazmat/backends/openssl/ec.py
@@ -22,7 +22,9 @@ from cryptography.hazmat.primitives.asymmetric import (
 )
 
 
-def _check_signature_algorithm(signature_algorithm):
+def _check_signature_algorithm(
+    signature_algorithm: ec.EllipticCurveSignatureAlgorithm,
+):
     if not isinstance(signature_algorithm, ec.ECDSA):
         raise UnsupportedAlgorithm(
             "Unsupported elliptic curve signature algorithm.",
@@ -104,42 +106,50 @@ def _ecdsa_sig_verify(backend, public_key, signature, data):
         raise InvalidSignature
 
 
-@utils.register_interface(AsymmetricSignatureContext)
-class _ECDSASignatureContext(object):
-    def __init__(self, backend, private_key, algorithm):
+class _ECDSASignatureContext(AsymmetricSignatureContext):
+    def __init__(
+        self,
+        backend,
+        private_key: ec.EllipticCurvePrivateKey,
+        algorithm: hashes.HashAlgorithm,
+    ):
         self._backend = backend
         self._private_key = private_key
         self._digest = hashes.Hash(algorithm, backend)
 
-    def update(self, data):
+    def update(self, data: bytes) -> None:
         self._digest.update(data)
 
-    def finalize(self):
+    def finalize(self) -> bytes:
         digest = self._digest.finalize()
 
         return _ecdsa_sig_sign(self._backend, self._private_key, digest)
 
 
-@utils.register_interface(AsymmetricVerificationContext)
-class _ECDSAVerificationContext(object):
-    def __init__(self, backend, public_key, signature, algorithm):
+class _ECDSAVerificationContext(AsymmetricVerificationContext):
+    def __init__(
+        self,
+        backend,
+        public_key: ec.EllipticCurvePublicKey,
+        signature: bytes,
+        algorithm: hashes.HashAlgorithm,
+    ):
         self._backend = backend
         self._public_key = public_key
         self._signature = signature
         self._digest = hashes.Hash(algorithm, backend)
 
-    def update(self, data):
+    def update(self, data: bytes) -> None:
         self._digest.update(data)
 
-    def verify(self):
+    def verify(self) -> None:
         digest = self._digest.finalize()
         _ecdsa_sig_verify(
             self._backend, self._public_key, self._signature, digest
         )
 
 
-@utils.register_interface(ec.EllipticCurvePrivateKey)
-class _EllipticCurvePrivateKey(object):
+class _EllipticCurvePrivateKey(ec.EllipticCurvePrivateKey):
     def __init__(self, backend, ec_key_cdata, evp_pkey):
         self._backend = backend
         self._ec_key = ec_key_cdata
@@ -152,18 +162,24 @@ class _EllipticCurvePrivateKey(object):
     curve = utils.read_only_property("_curve")
 
     @property
-    def key_size(self):
+    def key_size(self) -> int:
         return self.curve.key_size
 
-    def signer(self, signature_algorithm):
+    def signer(
+        self, signature_algorithm: ec.EllipticCurveSignatureAlgorithm
+    ) -> AsymmetricSignatureContext:
         _warn_sign_verify_deprecated()
         _check_signature_algorithm(signature_algorithm)
         _check_not_prehashed(signature_algorithm.algorithm)
+        # This assert is to help mypy realize what type this object holds
+        assert isinstance(signature_algorithm.algorithm, hashes.HashAlgorithm)
         return _ECDSASignatureContext(
             self._backend, self, signature_algorithm.algorithm
         )
 
-    def exchange(self, algorithm, peer_public_key):
+    def exchange(
+        self, algorithm: ec.ECDH, peer_public_key: ec.EllipticCurvePublicKey
+    ) -> bytes:
         if not (
             self._backend.elliptic_curve_exchange_algorithm_supported(
                 algorithm, self.curve
@@ -184,7 +200,7 @@ class _EllipticCurvePrivateKey(object):
         self._backend.openssl_assert(z_len > 0)
         z_buf = self._backend._ffi.new("uint8_t[]", z_len)
         peer_key = self._backend._lib.EC_KEY_get0_public_key(
-            peer_public_key._ec_key
+            peer_public_key._ec_key  # type: ignore[attr-defined]
         )
 
         r = self._backend._lib.ECDH_compute_key(
@@ -193,7 +209,7 @@ class _EllipticCurvePrivateKey(object):
         self._backend.openssl_assert(r > 0)
         return self._backend._ffi.buffer(z_buf)[:z_len]
 
-    def public_key(self):
+    def public_key(self) -> ec.EllipticCurvePublicKey:
         group = self._backend._lib.EC_KEY_get0_group(self._ec_key)
         self._backend.openssl_assert(group != self._backend._ffi.NULL)
 
@@ -210,7 +226,7 @@ class _EllipticCurvePrivateKey(object):
 
         return _EllipticCurvePublicKey(self._backend, public_ec_key, evp_pkey)
 
-    def private_numbers(self):
+    def private_numbers(self) -> ec.EllipticCurvePrivateNumbers:
         bn = self._backend._lib.EC_KEY_get0_private_key(self._ec_key)
         private_value = self._backend._bn_to_int(bn)
         return ec.EllipticCurvePrivateNumbers(
@@ -218,7 +234,12 @@ class _EllipticCurvePrivateKey(object):
             public_numbers=self.public_key().public_numbers(),
         )
 
-    def private_bytes(self, encoding, format, encryption_algorithm):
+    def private_bytes(
+        self,
+        encoding: serialization.Encoding,
+        format: serialization.PrivateFormat,
+        encryption_algorithm: serialization.KeySerializationEncryption,
+    ) -> bytes:
         return self._backend._private_key_bytes(
             encoding,
             format,
@@ -228,16 +249,21 @@ class _EllipticCurvePrivateKey(object):
             self._ec_key,
         )
 
-    def sign(self, data, signature_algorithm):
+    def sign(
+        self,
+        data: bytes,
+        signature_algorithm: ec.EllipticCurveSignatureAlgorithm,
+    ) -> bytes:
         _check_signature_algorithm(signature_algorithm)
         data, algorithm = _calculate_digest_and_algorithm(
-            self._backend, data, signature_algorithm._algorithm
+            self._backend,
+            data,
+            signature_algorithm._algorithm,  # type: ignore[attr-defined]
         )
         return _ecdsa_sig_sign(self._backend, self, data)
 
 
-@utils.register_interface(ec.EllipticCurvePublicKey)
-class _EllipticCurvePublicKey(object):
+class _EllipticCurvePublicKey(ec.EllipticCurvePublicKey):
     def __init__(self, backend, ec_key_cdata, evp_pkey):
         self._backend = backend
         self._ec_key = ec_key_cdata
@@ -250,20 +276,26 @@ class _EllipticCurvePublicKey(object):
     curve = utils.read_only_property("_curve")
 
     @property
-    def key_size(self):
+    def key_size(self) -> int:
         return self.curve.key_size
 
-    def verifier(self, signature, signature_algorithm):
+    def verifier(
+        self,
+        signature: bytes,
+        signature_algorithm: ec.EllipticCurveSignatureAlgorithm,
+    ) -> AsymmetricVerificationContext:
         _warn_sign_verify_deprecated()
         utils._check_bytes("signature", signature)
 
         _check_signature_algorithm(signature_algorithm)
         _check_not_prehashed(signature_algorithm.algorithm)
+        # This assert is to help mypy realize what type this object holds
+        assert isinstance(signature_algorithm.algorithm, hashes.HashAlgorithm)
         return _ECDSAVerificationContext(
             self._backend, self, signature, signature_algorithm.algorithm
         )
 
-    def public_numbers(self):
+    def public_numbers(self) -> ec.EllipticCurvePublicNumbers:
         get_func, group = self._backend._ec_key_determine_group_get_func(
             self._ec_key
         )
@@ -282,7 +314,7 @@ class _EllipticCurvePublicKey(object):
 
         return ec.EllipticCurvePublicNumbers(x=x, y=y, curve=self._curve)
 
-    def _encode_point(self, format):
+    def _encode_point(self, format: serialization.PublicFormat) -> bytes:
         if format is serialization.PublicFormat.CompressedPoint:
             conversion = self._backend._lib.POINT_CONVERSION_COMPRESSED
         else:
@@ -306,8 +338,11 @@ class _EllipticCurvePublicKey(object):
 
         return self._backend._ffi.buffer(buf)[:]
 
-    def public_bytes(self, encoding, format):
-
+    def public_bytes(
+        self,
+        encoding: serialization.Encoding,
+        format: serialization.PublicFormat,
+    ) -> bytes:
         if (
             encoding is serialization.Encoding.X962
             or format is serialization.PublicFormat.CompressedPoint
@@ -328,9 +363,16 @@ class _EllipticCurvePublicKey(object):
                 encoding, format, self, self._evp_pkey, None
             )
 
-    def verify(self, signature, data, signature_algorithm):
+    def verify(
+        self,
+        signature: bytes,
+        data: bytes,
+        signature_algorithm: ec.EllipticCurveSignatureAlgorithm,
+    ) -> None:
         _check_signature_algorithm(signature_algorithm)
         data, algorithm = _calculate_digest_and_algorithm(
-            self._backend, data, signature_algorithm._algorithm
+            self._backend,
+            data,
+            signature_algorithm._algorithm,  # type: ignore[attr-defined]
         )
         _ecdsa_sig_verify(self._backend, self, signature, data)
diff --git a/src/cryptography/hazmat/primitives/asymmetric/ec.py b/src/cryptography/hazmat/primitives/asymmetric/ec.py
index c1b7473fc..734226920 100644
--- a/src/cryptography/hazmat/primitives/asymmetric/ec.py
+++ b/src/cryptography/hazmat/primitives/asymmetric/ec.py
@@ -10,6 +10,12 @@ import warnings
 from cryptography import utils
 from cryptography.hazmat._oid import ObjectIdentifier
 from cryptography.hazmat.backends import _get_backend
+from cryptography.hazmat.primitives import _serialization, hashes
+from cryptography.hazmat.primitives.asymmetric import (
+    AsymmetricSignatureContext,
+    AsymmetricVerificationContext,
+    utils as asym_utils,
+)
 
 
 class EllipticCurveOID(object):
@@ -36,13 +42,13 @@ class EllipticCurveOID(object):
 
 class EllipticCurve(metaclass=abc.ABCMeta):
     @abc.abstractproperty
-    def name(self):
+    def name(self) -> str:
         """
         The name of the curve. e.g. secp256r1.
         """
 
     @abc.abstractproperty
-    def key_size(self):
+    def key_size(self) -> int:
         """
         Bit size of a secret scalar for the curve.
         """
@@ -50,7 +56,9 @@ class EllipticCurve(metaclass=abc.ABCMeta):
 
 class EllipticCurveSignatureAlgorithm(metaclass=abc.ABCMeta):
     @abc.abstractproperty
-    def algorithm(self):
+    def algorithm(
+        self,
+    ) -> typing.Union[asym_utils.Prehashed, hashes.HashAlgorithm]:
         """
         The digest algorithm used with this signature.
         """
@@ -58,50 +66,64 @@ class EllipticCurveSignatureAlgorithm(metaclass=abc.ABCMeta):
 
 class EllipticCurvePrivateKey(metaclass=abc.ABCMeta):
     @abc.abstractmethod
-    def signer(self, signature_algorithm):
+    def signer(
+        self,
+        signature_algorithm: EllipticCurveSignatureAlgorithm,
+    ) -> AsymmetricSignatureContext:
         """
         Returns an AsymmetricSignatureContext used for signing data.
         """
 
     @abc.abstractmethod
-    def exchange(self, algorithm, peer_public_key):
+    def exchange(
+        self, algorithm: "ECDH", peer_public_key: "EllipticCurvePublicKey"
+    ) -> bytes:
         """
         Performs a key exchange operation using the provided algorithm with the
         provided peer's public key.
         """
 
     @abc.abstractmethod
-    def public_key(self):
+    def public_key(self) -> "EllipticCurvePublicKey":
         """
         The EllipticCurvePublicKey for this private key.
         """
 
     @abc.abstractproperty
-    def curve(self):
+    def curve(self) -> EllipticCurve:
         """
         The EllipticCurve that this key is on.
         """
 
     @abc.abstractproperty
-    def key_size(self):
+    def key_size(self) -> int:
         """
         Bit size of a secret scalar for the curve.
         """
 
     @abc.abstractmethod
-    def sign(self, data, signature_algorithm):
+    def sign(
+        self,
+        data,
+        signature_algorithm: EllipticCurveSignatureAlgorithm,
+    ) -> bytes:
         """
         Signs the data
         """
 
     @abc.abstractmethod
-    def private_numbers(self):
+    def private_numbers(self) -> "EllipticCurvePrivateNumbers":
         """
         Returns an EllipticCurvePrivateNumbers.
         """
 
     @abc.abstractmethod
-    def private_bytes(self, encoding, format, encryption_algorithm):
+    def private_bytes(
+        self,
+        encoding: _serialization.Encoding,
+        format: _serialization.PrivateFormat,
+        encryption_algorithm: _serialization.KeySerializationEncryption,
+    ) -> bytes:
         """
         Returns the key serialized as bytes.
         """
@@ -112,43 +134,58 @@ EllipticCurvePrivateKeyWithSerialization = EllipticCurvePrivateKey
 
 class EllipticCurvePublicKey(metaclass=abc.ABCMeta):
     @abc.abstractmethod
-    def verifier(self, signature, signature_algorithm):
+    def verifier(
+        self,
+        signature: bytes,
+        signature_algorithm: EllipticCurveSignatureAlgorithm,
+    ) -> AsymmetricVerificationContext:
         """
         Returns an AsymmetricVerificationContext used for signing data.
         """
 
     @abc.abstractproperty
-    def curve(self):
+    def curve(self) -> EllipticCurve:
         """
         The EllipticCurve that this key is on.
         """
 
     @abc.abstractproperty
-    def key_size(self):
+    def key_size(self) -> int:
         """
         Bit size of a secret scalar for the curve.
         """
 
     @abc.abstractmethod
-    def public_numbers(self):
+    def public_numbers(self) -> "EllipticCurvePublicNumbers":
         """
         Returns an EllipticCurvePublicNumbers.
         """
 
     @abc.abstractmethod
-    def public_bytes(self, encoding, format):
+    def public_bytes(
+        self,
+        encoding: _serialization.Encoding,
+        format: _serialization.PublicFormat,
+    ) -> bytes:
         """
         Returns the key serialized as bytes.
         """
 
     @abc.abstractmethod
-    def verify(self, signature, data, signature_algorithm):
+    def verify(
+        self,
+        signature: bytes,
+        data: bytes,
+        algorithm: EllipticCurveSignatureAlgorithm,
+    ) -> None:
         """
         Verifies the signature of the data.
         """
 
     @classmethod
-    def from_encoded_point(cls, curve, data):
+    def from_encoded_point(
+        cls, curve: EllipticCurve, data: bytes
+    ) -> "EllipticCurvePublicKey":
         utils._check_bytes("data", data)
 
         if not isinstance(curve, EllipticCurve):
@@ -288,20 +325,23 @@ _CURVE_TYPES: typing.Dict[str, typing.Type[EllipticCurve]] = {
 }
 
 
-@utils.register_interface(EllipticCurveSignatureAlgorithm)
-class ECDSA(object):
+class ECDSA(EllipticCurveSignatureAlgorithm):
     def __init__(self, algorithm):
         self._algorithm = algorithm
 
     algorithm = utils.read_only_property("_algorithm")
 
 
-def generate_private_key(curve, backend=None):
+def generate_private_key(
+    curve: EllipticCurve, backend=None
+) -> EllipticCurvePrivateKey:
     backend = _get_backend(backend)
     return backend.generate_elliptic_curve_private_key(curve)
 
 
-def derive_private_key(private_value, curve, backend=None):
+def derive_private_key(
+    private_value: int, curve: EllipticCurve, backend=None
+) -> EllipticCurvePrivateKey:
     backend = _get_backend(backend)
     if not isinstance(private_value, int):
         raise TypeError("private_value must be an integer type.")
@@ -316,7 +356,7 @@ def derive_private_key(private_value, curve, backend=None):
 
 
 class EllipticCurvePublicNumbers(object):
-    def __init__(self, x, y, curve):
+    def __init__(self, x: int, y: int, curve: EllipticCurve):
         if not isinstance(x, int) or not isinstance(y, int):
             raise TypeError("x and y must be integers.")
 
@@ -327,11 +367,11 @@ class EllipticCurvePublicNumbers(object):
         self._x = x
         self._curve = curve
 
-    def public_key(self, backend=None):
+    def public_key(self, backend=None) -> EllipticCurvePublicKey:
         backend = _get_backend(backend)
         return backend.load_elliptic_curve_public_numbers(self)
 
-    def encode_point(self):
+    def encode_point(self) -> bytes:
         warnings.warn(
             "encode_point has been deprecated on EllipticCurvePublicNumbers"
             " and will be removed in a future version. Please use "
@@ -349,7 +389,9 @@ class EllipticCurvePublicNumbers(object):
         )
 
     @classmethod
-    def from_encoded_point(cls, curve, data):
+    def from_encoded_point(
+        cls, curve: EllipticCurve, data: bytes
+    ) -> "EllipticCurvePublicNumbers":
         if not isinstance(curve, EllipticCurve):
             raise TypeError("curve must be an EllipticCurve instance")
 
@@ -402,7 +444,9 @@ class EllipticCurvePublicNumbers(object):
 
 
 class EllipticCurvePrivateNumbers(object):
-    def __init__(self, private_value, public_numbers):
+    def __init__(
+        self, private_value: int, public_numbers: EllipticCurvePublicNumbers
+    ):
         if not isinstance(private_value, int):
             raise TypeError("private_value must be an integer.")
 
@@ -415,7 +459,7 @@ class EllipticCurvePrivateNumbers(object):
         self._private_value = private_value
         self._public_numbers = public_numbers
 
-    def private_key(self, backend=None):
+    def private_key(self, backend=None) -> EllipticCurvePrivateKey:
         backend = _get_backend(backend)
         return backend.load_elliptic_curve_private_numbers(self)
 
diff --git a/tests/hazmat/primitives/test_ec.py b/tests/hazmat/primitives/test_ec.py
index b69475a12..c089adc43 100644
--- a/tests/hazmat/primitives/test_ec.py
+++ b/tests/hazmat/primitives/test_ec.py
@@ -88,7 +88,7 @@ class DummyCurve(ec.EllipticCurve):
 
 
 class DummySignatureAlgorithm(ec.EllipticCurveSignatureAlgorithm):
-    algorithm = None
+    algorithm = hashes.SHA256()
 
 
 @pytest.mark.requires_backend_interface(interface=EllipticCurveBackend)
@@ -129,10 +129,10 @@ def test_derive_private_key_errors(backend):
     _skip_curve_unsupported(backend, curve)
 
     with pytest.raises(TypeError):
-        ec.derive_private_key("one", curve, backend)
+        ec.derive_private_key("one", curve, backend)  # type: ignore[arg-type]
 
     with pytest.raises(TypeError):
-        ec.derive_private_key(10, "five", backend)
+        ec.derive_private_key(10, "five", backend)  # type: ignore[arg-type]
 
     with pytest.raises(ValueError):
         ec.derive_private_key(-7, curve, backend)
@@ -167,7 +167,7 @@ def test_invalid_ec_numbers_args(private_value, x, y, curve):
 
 def test_invalid_private_numbers_public_numbers():
     with pytest.raises(TypeError):
-        ec.EllipticCurvePrivateNumbers(1, None)
+        ec.EllipticCurvePrivateNumbers(1, None)  # type: ignore[arg-type]
 
 
 def test_encode_point():
@@ -233,7 +233,7 @@ def test_from_encoded_point_not_a_curve():
     with pytest.raises(TypeError):
         with pytest.warns(CryptographyDeprecationWarning):
             ec.EllipticCurvePublicNumbers.from_encoded_point(
-                "notacurve", b"\x04data"
+                "notacurve", b"\x04data"  # type: ignore[arg-type]
             )
 
 
@@ -1128,7 +1128,7 @@ class TestEllipticCurvePEMPublicKeySerialization(object):
     def test_from_encoded_point_not_a_curve(self):
         with pytest.raises(TypeError):
             ec.EllipticCurvePublicKey.from_encoded_point(
-                "notacurve", b"\x04data"
+                "notacurve", b"\x04data"  # type: ignore[arg-type]
             )
 
     def test_from_encoded_point_unsupported_encoding(self):
@@ -1187,7 +1187,9 @@ class TestECDSAVerification(object):
         with pytest.raises(TypeError), pytest.warns(
             CryptographyDeprecationWarning
         ):
-            public_key.verifier(1234, ec.ECDSA(hashes.SHA256()))
+            public_key.verifier(
+                1234, ec.ECDSA(hashes.SHA256())  # type: ignore[arg-type]
+            )
 
 
 @pytest.mark.requires_backend_interface(interface=EllipticCurveBackend)
@@ -1241,15 +1243,15 @@ class TestECDH(object):
                     peer_pubkey = public_numbers.public_key(backend)
 
                 z = private_key.exchange(ec.ECDH(), peer_pubkey)
-                z = int(hexlify(z).decode("ascii"), 16)
+                zz = int(hexlify(z).decode("ascii"), 16)
                 # At this point fail indicates that one of the underlying keys
                 # was changed. This results in a non-matching derived key.
                 if vector["fail"]:
                     # Errno 8 indicates Z should be changed.
                     assert vector["errno"] == 8
-                    assert z != vector["Z"]
+                    assert zz != vector["Z"]
                 else:
-                    assert z == vector["Z"]
+                    assert zz == vector["Z"]
 
     @pytest.mark.parametrize(
         "vector",
diff --git a/tests/wycheproof/test_ecdh.py b/tests/wycheproof/test_ecdh.py
index a1a90c141..510ec9337 100644
--- a/tests/wycheproof/test_ecdh.py
+++ b/tests/wycheproof/test_ecdh.py
@@ -66,6 +66,7 @@ def test_ecdh(backend, wycheproof):
         public_key = serialization.load_der_public_key(
             binascii.unhexlify(wycheproof.testcase["public"]), backend
         )
+        assert isinstance(public_key, ec.EllipticCurvePublicKey)
     except NotImplementedError:
         assert wycheproof.has_flag("UnnamedCurve")
         return
@@ -93,6 +94,7 @@ def test_ecdh(backend, wycheproof):
 )
 def test_ecdh_ecpoint(backend, wycheproof):
     curve = _CURVES[wycheproof.testgroup["curve"]]
+    assert isinstance(curve, ec.EllipticCurve)
     _skip_exchange_algorithm_unsupported(backend, ec.ECDH(), curve)
 
     private_key = ec.derive_private_key(