Skip to content

Commit f0e2825

Browse files
committed
test_tls: Add type annotations
The file now passes `mypy --strict`.
1 parent 018b16f commit f0e2825

File tree

4 files changed

+335
-167
lines changed

4 files changed

+335
-167
lines changed

src/mbedtls/_tls.pyi

Lines changed: 69 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -5,17 +5,19 @@ from __future__ import annotations
55

66
import enum
77
import sys
8+
from collections import abc
89
from pathlib import Path
910
from typing import Mapping, Optional, Sequence, Tuple, Union, overload
1011

12+
from mbedtls.pk import ECC, RSA
1113
from mbedtls.x509 import CRT
1214

1315
if sys.version_info < (3, 8):
1416
from typing_extensions import Literal
1517
else:
1618
from typing import Literal
1719

18-
def ciphers_available() -> Sequence[bytes]: ...
20+
def ciphers_available() -> Sequence[str]: ...
1921
@enum.unique
2022
class NextProtocol(enum.Enum):
2123
H2: bytes
@@ -77,8 +79,10 @@ class WantReadError(TLSError): ...
7779
class RaggedEOF(TLSError): ...
7880
class HelloVerifyRequest(TLSError): ...
7981

80-
class TrustStore:
81-
def __init__(self, db: Optional[Sequence[CRT]] = ...) -> None: ...
82+
class TrustStore(abc.Sequence[CRT]):
83+
def __init__(
84+
self, db: Optional[Union[Sequence[CRT], TrustStore]] = ...
85+
) -> None: ...
8286
@classmethod
8387
def system(cls) -> TrustStore: ...
8488
@classmethod
@@ -96,10 +100,10 @@ class Purpose(enum.IntEnum):
96100
SERVER_AUTH: int
97101
CLIENT_AUTH: int
98102

99-
# TODO: Type these.
100-
Certificate = object
101-
PrivateKey = object
102-
CipherSuite = object
103+
_Key = Union[RSA, ECC]
104+
Certificate = CRT
105+
PrivateKey = _Key
106+
CipherSuite = str
103107
ServerNameCallback = object
104108

105109
class TLSConfiguration:
@@ -109,10 +113,25 @@ class TLSConfiguration:
109113
cls,
110114
validate_certificates: Optional[bool] = ...,
111115
certificate_chain: Optional[
112-
Tuple[Tuple[Certificate], PrivateKey]
116+
Tuple[Tuple[Certificate, ...], PrivateKey]
117+
] = ...,
118+
ciphers: Optional[Sequence[Union[CipherSuite, int]]] = ...,
119+
inner_protocols: Optional[Sequence[Union[NextProtocol, bytes]]] = ...,
120+
lowest_supported_version: Optional[TLSVersion] = ...,
121+
highest_supported_version: Optional[TLSVersion] = ...,
122+
trust_store: Optional[TrustStore] = ...,
123+
sni_callback: Optional[ServerNameCallback] = ...,
124+
pre_shared_key: Optional[Tuple[str, bytes]] = ...,
125+
pre_shared_key_store: Optional[Mapping[str, bytes]] = ...,
126+
) -> TLSConfiguration: ...
127+
def update(
128+
self,
129+
validate_certificates: Optional[bool] = ...,
130+
certificate_chain: Optional[
131+
Tuple[Tuple[Certificate, ...], PrivateKey]
113132
] = ...,
114-
ciphers: Optional[Tuple[Union[CipherSuite, int]]] = ...,
115-
inner_protocols: Optional[Tuple[Union[NextProtocol, bytes]]] = ...,
133+
ciphers: Optional[Sequence[Union[CipherSuite, int]]] = ...,
134+
inner_protocols: Optional[Sequence[Union[NextProtocol, bytes]]] = ...,
116135
lowest_supported_version: Optional[TLSVersion] = ...,
117136
highest_supported_version: Optional[TLSVersion] = ...,
118137
trust_store: Optional[TrustStore] = ...,
@@ -121,26 +140,49 @@ class TLSConfiguration:
121140
pre_shared_key_store: Optional[Mapping[str, bytes]] = ...,
122141
) -> TLSConfiguration: ...
123142
validate_certificates: Optional[bool]
124-
certificate_chain: Optional[Tuple[Tuple[Certificate], PrivateKey]]
143+
certificate_chain: Optional[Tuple[Tuple[Certificate, ...], PrivateKey]]
125144
ciphers: Optional[Tuple[Union[CipherSuite, int]]]
126145
inner_protocols: Optional[Tuple[Union[NextProtocol, bytes]]]
127146
lowest_supported_version: Optional[TLSVersion]
128147
highest_supported_version: Optional[TLSVersion]
129148
trust_store: Optional[TrustStore]
130149
sni_callback: Optional[ServerNameCallback]
150+
pre_shared_key: Optional[Tuple[str, bytes]]
151+
pre_shared_key_store: Optional[Mapping[str, bytes]]
131152

132153
class DTLSConfiguration:
133154
def __new__(
134155
cls,
135156
validate_certificates: Optional[bool] = ...,
136157
certificate_chain: Optional[
137-
Tuple[Tuple[Certificate], PrivateKey]
158+
Tuple[Tuple[Certificate, ...], PrivateKey]
138159
] = ...,
139-
ciphers: Optional[Tuple[Union[CipherSuite, int]]] = ...,
140-
inner_protocols: Optional[Tuple[Union[NextProtocol, bytes]]] = ...,
141-
lowest_supported_version: Optional[TLSVersion] = ...,
142-
highest_supported_version: Optional[TLSVersion] = ...,
160+
ciphers: Optional[Sequence[Union[CipherSuite, int]]] = ...,
161+
inner_protocols: Optional[Sequence[Union[NextProtocol, bytes]]] = ...,
162+
lowest_supported_version: Optional[DTLSVersion] = ...,
163+
highest_supported_version: Optional[DTLSVersion] = ...,
143164
trust_store: Optional[TrustStore] = ...,
165+
anti_replay: Optional[bool] = ...,
166+
handshake_timeout_min: Optional[int] = ...,
167+
handshake_timeout_max: Optional[int] = ...,
168+
sni_callback: Optional[ServerNameCallback] = ...,
169+
pre_shared_key: Optional[Tuple[str, bytes]] = ...,
170+
pre_shared_key_store: Optional[Mapping[str, bytes]] = ...,
171+
) -> DTLSConfiguration: ...
172+
def update(
173+
self,
174+
validate_certificates: Optional[bool] = ...,
175+
certificate_chain: Optional[
176+
Tuple[Tuple[Certificate, ...], PrivateKey]
177+
] = ...,
178+
ciphers: Optional[Sequence[Union[CipherSuite, int]]] = ...,
179+
inner_protocols: Optional[Sequence[Union[NextProtocol, bytes]]] = ...,
180+
lowest_supported_version: Optional[DTLSVersion] = ...,
181+
highest_supported_version: Optional[DTLSVersion] = ...,
182+
trust_store: Optional[TrustStore] = ...,
183+
anti_replay: Optional[bool] = ...,
184+
handshake_timeout_min: Optional[int] = ...,
185+
handshake_timeout_max: Optional[int] = ...,
144186
sni_callback: Optional[ServerNameCallback] = ...,
145187
pre_shared_key: Optional[Tuple[str, bytes]] = ...,
146188
pre_shared_key_store: Optional[Mapping[str, bytes]] = ...,
@@ -149,10 +191,15 @@ class DTLSConfiguration:
149191
certificate_chain: Optional[Tuple[Tuple[Certificate], PrivateKey]]
150192
ciphers: Optional[Tuple[Union[CipherSuite, int]]]
151193
inner_protocols: Optional[Tuple[Union[NextProtocol, bytes]]]
152-
lowest_supported_version: Optional[TLSVersion]
153-
highest_supported_version: Optional[TLSVersion]
194+
lowest_supported_version: Optional[DTLSVersion]
195+
highest_supported_version: Optional[DTLSVersion]
154196
trust_store: Optional[TrustStore]
197+
anti_replay: Optional[bool]
198+
handshake_timeout_min: Optional[int]
199+
handshake_timeout_max: Optional[int]
155200
sni_callback: Optional[ServerNameCallback]
201+
pre_shared_key: Optional[Tuple[str, bytes]]
202+
pre_shared_key_store: Optional[Mapping[str, bytes]]
156203

157204
class _BaseContext:
158205
def __init__(
@@ -169,6 +216,10 @@ class MbedTLSBuffer:
169216
self, context: _BaseContext, server_hostname: Optional[str] = None
170217
) -> None: ...
171218
@property
219+
def _input_buffer(self) -> bytes: ...
220+
@property
221+
def _output_buffer(self) -> bytes: ...
222+
@property
172223
def context(self) -> _BaseContext: ...
173224
@property
174225
def _server_hostname(self) -> str: ...

src/mbedtls/_tls.pyx

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -259,7 +259,7 @@ class HelloVerifyRequest(_exc.TLSError):
259259
pass
260260

261261

262-
class TrustStore:
262+
class TrustStore(abc.Sequence):
263263
def __init__(self, db=None):
264264
if db is None:
265265
db = []

src/mbedtls/tls.py

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,9 @@
99
import sys
1010
from typing import Any, NoReturn, Optional, Tuple, Union, cast
1111

12-
from ._tls import (
13-
DTLSConfiguration,
14-
DTLSVersion,
15-
HandshakeStep,
16-
HelloVerifyRequest,
17-
)
12+
from ._tls import DTLSConfiguration, DTLSVersion
13+
from ._tls import HandshakeStep as HandshakeStep
14+
from ._tls import HelloVerifyRequest
1815
from ._tls import MbedTLSBuffer as TLSWrappedBuffer
1916
from ._tls import (
2017
NextProtocol,

0 commit comments

Comments
 (0)