test_pk: Add type annotations

The file now passes `mypy --strict`.

Author: Synss

src/mbedtls/pk.pyi

@@ -107,9 +107,13 @@ class RSA(CipherBase):
     @overload
     def export_key(self, format: Literal["PEM"]) -> _PEM: ...
     @overload
+    def export_key(self) -> Union[_DER, _PEM]: ...
+    @overload
     def export_public_key(self, format: Literal["DER"]) -> _DER: ...
     @overload
     def export_public_key(self, format: Literal["PEM"]) -> _PEM: ...
+    @overload
+    def export_public_key(self) -> Union[_DER, _PEM]: ...
 
 class ECPoint:
     def __init__(self, x: _MPI, y: _MPI, z: _MPI) -> None: ...
@@ -132,6 +136,8 @@ class ECC(CipherBase):
         key: Optional[bytes] = ...,
         password: Optional[bytes] = ...,
     ) -> None: ...
+    @property
+    def curve(self) -> Curve: ...
     def generate(self) -> bytes: ...
     @overload
     def export_key(self, format: Literal["DER"]) -> _DER: ...
@@ -140,11 +146,15 @@ class ECC(CipherBase):
     @overload
     def export_key(self, format: Literal["NUM"]) -> _NUM: ...
     @overload
+    def export_key(self) -> Union[_DER, _PEM, _NUM]: ...
+    @overload
     def export_public_key(self, format: Literal["DER"]) -> _DER: ...
     @overload
     def export_public_key(self, format: Literal["PEM"]) -> _PEM: ...
     @overload
     def export_public_key(self, format: Literal["POINT"]) -> ECPoint: ...
+    @overload
+    def export_public_key(self) -> Union[_DER, _PEM, ECPoint]: ...
 
 class DHServer:
     def __init__(self, modulus: _MPI, generator: _MPI) -> None: ...

tests/test_pk.py

@@ -2,7 +2,8 @@
 
 import numbers
 import pickle
-from typing import Any, List, Union, cast
+import sys
+from typing import Any, Callable, List, Tuple, Type, TypeVar, Union, cast
 
 import pytest
 
@@ -13,7 +14,6 @@
 from mbedtls.pk import (
     ECC,
     RSA,
-    CipherBase,
     Curve,
     DHClient,
     DHServer,
@@ -26,6 +26,14 @@
     get_supported_curves,
 )
 
+if sys.version_info < (3, 8):
+    from typing_extensions import Literal
+else:
+    from typing import Literal
+
+
+T = TypeVar("T")
+
 
 def test_supported_curves() -> None:
     assert sorted(get_supported_curves()) == [
@@ -55,14 +63,14 @@ def test_get_supported_ciphers() -> None:
     ]
 
 
-def test_rsa_encryp_decrypt(randbytes) -> None:
+def test_rsa_encryp_decrypt(randbytes: Callable[[int], bytes]) -> None:
     rsa = RSA()
     rsa.generate(1024)
     msg = randbytes(rsa.key_size - 11)
     assert rsa.decrypt(rsa.encrypt(msg)) == msg
 
 
-def do_generate(cipher: CipherBase):
+def do_generate(cipher: T) -> bytes:
     if isinstance(cipher, RSA):
         return cipher.generate(1024)
     if isinstance(cipher, ECC):
@@ -72,30 +80,38 @@ def do_generate(cipher: CipherBase):
 
 class TestECPoint:
     @pytest.fixture(params=[(MPI(1), MPI(2), MPI(3)), (1, 2, 3)])
-    def xyz(self, request):
-        return request.param
+    def xyz(self, request: Any) -> Tuple[int, int, int]:
+        return request.param  # type: ignore[no-any-return]
 
     @pytest.fixture()
-    def point(self, xyz):
+    def point(self, xyz: Tuple[int, int, int]) -> ECPoint:
         return ECPoint(*xyz)
 
-    @pytest.mark.parametrize("repr_", [repr, str], ids=lambda f: f.__name__)
-    def test_repr(self, repr_, point) -> None:
+    @pytest.mark.parametrize(
+        "repr_",
+        [repr, str],
+        ids=lambda f: f.__name__,  # type: ignore[no-any-return]
+    )
+    def test_repr(
+        self, repr_: Callable[[object], str], point: ECPoint
+    ) -> None:
         assert isinstance(repr_(point), str)
 
-    def test_pickle(self, point) -> None:
+    def test_pickle(self, point: ECPoint) -> None:
         assert point == pickle.loads(pickle.dumps(point))
 
-    def test_hash(self, point) -> None:
+    def test_hash(self, point: ECPoint) -> None:
         assert isinstance(hash(point), int)
 
-    def test_accessors(self, point, xyz) -> None:
+    def test_accessors(
+        self, point: ECPoint, xyz: Tuple[int, int, int]
+    ) -> None:
         x, y, z = xyz
         assert point.x == x
         assert point.y == y
         assert point.z == z
 
-    def test_eq_point(self, point, xyz) -> None:
+    def test_eq_point(self, point: ECPoint, xyz: Tuple[int, int, int]) -> None:
         assert (point == ECPoint(*xyz)) is True
         assert (point == ECPoint(0, 0, 0)) is False
 
@@ -105,7 +121,7 @@ def test_eq_zero(self) -> None:
         assert (zero == 0) is True
         assert (zero == ECPoint(0, 0, 0)) is True
 
-    def test_bool(self, point) -> None:
+    def test_bool(self, point: ECPoint) -> None:
         assert bool(point) is True
         assert bool(ECPoint(0, 0, 0)) is False
 
@@ -127,13 +143,13 @@ def cipher(self, request: Any) -> Union[ECC, RSA]:
             return cast(RSA, request.param())
         return ECC(request.param)
 
-    def test_pickle(self, cipher) -> None:
+    def test_pickle(self, cipher: Union[ECC, RSA]) -> None:
         assert cipher == pickle.loads(pickle.dumps(cipher))
 
-    def test_hash(self, cipher) -> None:
+    def test_hash(self, cipher: Union[ECC, RSA]) -> None:
         assert isinstance(hash(cipher), int)
 
-    def test_export_private_key(self, cipher) -> None:
+    def test_export_private_key(self, cipher: Union[ECC, RSA]) -> None:
         assert cipher.export_key("DER") == b""
         assert cipher.export_key("PEM") == ""
         assert cipher.key_size == 0
@@ -150,7 +166,7 @@ def test_export_private_key(self, cipher) -> None:
         assert type(cipher).from_PEM(cipher.export_key("PEM")) == cipher
         assert cipher.key_size > 0
 
-    def test_import_private_key(self, cipher) -> None:
+    def test_import_private_key(self, cipher: Union[ECC, RSA]) -> None:
         assert not cipher.export_key()
         assert not cipher.export_public_key()
 
@@ -164,7 +180,7 @@ def test_import_private_key(self, cipher) -> None:
         assert check_pair(other, other) is True
         assert cipher == other
 
-    def test_export_public_key(self, cipher) -> None:
+    def test_export_public_key(self, cipher: Union[ECC, RSA]) -> None:
         assert cipher.export_public_key("DER") == b""
         assert cipher.export_public_key("PEM") == ""
 
@@ -178,13 +194,13 @@ def test_export_public_key(self, cipher) -> None:
             cipher.export_public_key("PEM")
         ) == cipher.export_public_key("PEM")
 
-    def test_import_public_key(self, cipher) -> None:
+    def test_import_public_key(self, cipher: Union[ECC, RSA]) -> None:
         assert not cipher.export_key()
         assert not cipher.export_public_key()
 
         do_generate(cipher)
 
-        pub = cipher.export_public_key()
+        pub = cipher.export_public_key("DER")
         other = type(cipher).from_buffer(pub)
         assert not other.export_key()
         assert other.export_public_key()
@@ -196,9 +212,14 @@ def test_import_public_key(self, cipher) -> None:
     @pytest.mark.parametrize(
         "digestmod",
         [_get_md_alg(name) for name in hashlib.algorithms_guaranteed],
-        ids=lambda dm: dm().name,
+        ids=lambda dm: dm().name,  # type: ignore[no-any-return]
     )
-    def test_sign_verify(self, cipher, digestmod, randbytes) -> None:
+    def test_sign_verify(
+        self,
+        cipher: Union[ECC, RSA],
+        digestmod: str,
+        randbytes: Callable[[int], bytes],
+    ) -> None:
         msg = randbytes(4096)
         assert cipher.sign(msg, digestmod) is None
 
@@ -212,17 +233,18 @@ def test_sign_verify(self, cipher, digestmod, randbytes) -> None:
 
 class TestECCExportKey:
     @pytest.fixture(params=get_supported_curves())
-    def curve(self, request):
+    def curve(self, request: Any) -> Curve:
+        assert isinstance(request.param, Curve)
         return request.param
 
-    def test_export_private_key(self, curve) -> None:
+    def test_export_private_key(self, curve: Curve) -> None:
         ecc = ECC(curve)
         assert ecc.export_key("NUM") == 0
 
         ecc.generate()
         assert ecc.export_key("NUM") != 0
 
-    def test_export_public_key_to_point(self, curve) -> None:
+    def test_export_public_key_to_point(self, curve: Curve) -> None:
         ecc = ECC(curve)
         assert ecc.export_public_key("POINT") == 0
         assert ecc.export_public_key("POINT") == ECPoint(0, 0, 0)
@@ -244,7 +266,7 @@ def test_export_public_key_to_point(self, curve) -> None:
 
 class TestECCtoECDH:
     @pytest.mark.parametrize("curve", get_supported_curves())
-    def test_exchange(self, curve) -> None:
+    def test_exchange(self, curve: Curve) -> None:
         ecp = ECC(curve)
         ecp.generate()
 
@@ -263,7 +285,9 @@ def test_exchange(self, curve) -> None:
 
 class TestDH:
     @pytest.mark.parametrize("dh_cls", [DHServer, DHClient])
-    def test_pickle(self, dh_cls) -> None:
+    def test_pickle(
+        self, dh_cls: Union[Type[DHServer], Type[DHClient]]
+    ) -> None:
         dhentity = dh_cls(MPI.prime(64), MPI.prime(20))
 
         with pytest.raises(TypeError) as excinfo:
@@ -272,7 +296,9 @@ def test_pickle(self, dh_cls) -> None:
         assert str(excinfo.value).startswith("cannot pickle")
 
     @pytest.mark.parametrize("dh_cls", [DHServer, DHClient])
-    def test_accessors(self, dh_cls) -> None:
+    def test_accessors(
+        self, dh_cls: Union[Type[DHServer], Type[DHClient]]
+    ) -> None:
         modulus = MPI.prime(64)
         generator = MPI.prime(20)
 
@@ -301,12 +327,14 @@ def test_exchange(self) -> None:
 
 class TestECDH:
     @pytest.fixture(params=get_supported_curves())
-    def key(self, request):
+    def key(self, request: Any) -> ECC:
         key = ECC(request.param)
         return key
 
     @pytest.mark.parametrize("peer_cls", [ECDHServer, ECDHClient])
-    def test_pickle(self, key, peer_cls) -> None:
+    def test_pickle(
+        self, key: ECC, peer_cls: Union[Type[ECDHServer], Type[ECDHClient]]
+    ) -> None:
         key.generate()
         peer = peer_cls(key)
 
@@ -316,7 +344,9 @@ def test_pickle(self, key, peer_cls) -> None:
         assert str(excinfo.value).startswith("cannot pickle")
 
     @pytest.mark.parametrize("peer_cls", [ECDHServer, ECDHClient])
-    def test_key_accessors_without_key(self, key, peer_cls) -> None:
+    def test_key_accessors_without_key(
+        self, key: ECC, peer_cls: Union[Type[ECDHServer], Type[ECDHClient]]
+    ) -> None:
         peer = peer_cls(key)
 
         assert not peer._has_private()
@@ -327,9 +357,9 @@ def test_key_accessors_without_key(self, key, peer_cls) -> None:
         assert peer.peers_public_key == 0
         assert peer.shared_secret == 0
 
-    def test_client_accessors_with_key(self, key) -> None:
+    def test_client_accessors_with_key(self, key: ECC) -> None:
         der = key.generate()
-        format = (
+        format: Literal["NUM", "DER"] = (
             "NUM" if key.curve in (Curve.CURVE25519, Curve.CURVE448) else "DER"
         )
         assert der == key.export_key(format)
@@ -344,9 +374,9 @@ def test_client_accessors_with_key(self, key) -> None:
         assert peer.peers_public_key == key.export_public_key("POINT")
         assert peer.shared_secret == 0
 
-    def test_server_accessors_with_key(self, key) -> None:
+    def test_server_accessors_with_key(self, key: ECC) -> None:
         der = key.generate()
-        format = (
+        format: Literal["NUM", "DER"] = (
             "NUM" if key.curve in (Curve.CURVE25519, Curve.CURVE448) else "DER"
         )
         assert der == key.export_key(format)
@@ -361,7 +391,7 @@ def test_server_accessors_with_key(self, key) -> None:
         assert peer.peers_public_key == 0
         assert peer.shared_secret == 0
 
-    def test_exchange(self, key) -> None:
+    def test_exchange(self, key: ECC) -> None:
         srv, cli = ECDHServer(key), ECDHClient(key)
 
         ske = srv.generate()
@@ -383,7 +413,7 @@ def test_exchange(self, key) -> None:
         assert srv_sec == cli_sec
         assert srv.shared_secret == cli.shared_secret
 
-    def test_generate_public(self, key) -> None:
+    def test_generate_public(self, key: ECC) -> None:
         srv, cli = ECDHServer(key), ECDHClient(key)
 
         srv.generate()
@@ -393,7 +423,7 @@ def test_generate_public(self, key) -> None:
         assert cli.public_key == srv.public_key
 
 
-def do_exchange(alice, bob) -> None:
+def do_exchange(alice: ECDHNaive, bob: ECDHNaive) -> None:
     alice_to_bob = alice.generate()
     bob_to_alice = bob.generate()
     alice.import_peers_public(bob_to_alice)
@@ -404,10 +434,11 @@ def do_exchange(alice, bob) -> None:
 
 class TestECDHNaive:
     @pytest.fixture(params=[Curve.CURVE448, Curve.CURVE25519])
-    def curve(self, request):
+    def curve(self, request: Any) -> Curve:
+        assert isinstance(request.param, Curve)
         return request.param
 
-    def test_key_accessors_without_key(self, curve) -> None:
+    def test_key_accessors_without_key(self, curve: Curve) -> None:
         peer = ECDHNaive(curve)
 
         assert not peer._has_private()
@@ -417,7 +448,7 @@ def test_key_accessors_without_key(self, curve) -> None:
         assert peer.public_key == 0
         assert peer.peers_public_key == 0
 
-    def test_exchange(self, curve) -> None:
+    def test_exchange(self, curve: Curve) -> None:
         alice, bob = ECDHNaive(curve), ECDHNaive(curve)
 
         alice_to_bob = alice.generate()
@@ -448,7 +479,7 @@ def test_exchange(self, curve) -> None:
         assert alice_secret == alice.shared_secret
         assert bob_secret == bob.shared_secret
 
-    def test_attacker_fails_with_public_keys(self, curve) -> None:
+    def test_attacker_fails_with_public_keys(self, curve: Curve) -> None:
         alice, bob = ECDHNaive(curve), ECDHNaive(curve)
         do_exchange(alice, bob)
 
@@ -458,7 +489,7 @@ def test_attacker_fails_with_public_keys(self, curve) -> None:
         with pytest.raises(TLSError):
             eve.generate_secret()
 
-    def test_attacker_succeeds_with_private_key(self, curve) -> None:
+    def test_attacker_succeeds_with_private_key(self, curve: Curve) -> None:
         alice, bob = ECDHNaive(curve), ECDHNaive(curve)
         do_exchange(alice, bob)