Coverage for zenmqtt/connection/__init__.py: 100%

1from ssl import SSLContext 

2from typing import Awaitable, Callable, Optional 

3from urllib.parse import ParseResult, urlparse 

4 

5from .base import ConnectionOptions, MQTTConnection, MQTTConnectionTransport 

6from .tcp import build_tcp_connection_transport, build_tls_connection_transport 

7 

8__all__ = [ 

9 "ConnectionOptions", 

10 "MQTTConnection", 

11 "MQTTConnectionTransport", 

12 "create_connection", 

13 "register_implementation", 

14] 

15 

16ConnectionFactory = Callable[ 

17 [ParseResult, ConnectionOptions], Awaitable[MQTTConnectionTransport] 

18] 

19 

20_CONNECTION_IMPLEMENTATION_MAP: dict[str, ConnectionFactory] = { 

21 "tcp": build_tcp_connection_transport, 

22 "mqtt": build_tcp_connection_transport, 

23 "mqtts": build_tls_connection_transport, 

24 "ssl": build_tls_connection_transport, 

25} 

26 

27# schemes which encrypt the connection, the only ones accepting an SSL context 

28_TLS_SCHEMES = {"mqtts", "ssl"} 

29 

30 

31def register_implementation(scheme: str, factory: ConnectionFactory) -> None: 

32 _CONNECTION_IMPLEMENTATION_MAP[scheme] = factory 

33 

34 

35async def create_connection( 

36 url: str, ssl: Optional[SSLContext] = None 

37) -> MQTTConnection: 

38 """ 

39 :param url: e.g. tcp://localhost:1883 or mqtts://broker:8883, the port is 

40 optional: 1883 for tcp:// and mqtt://, 8883 for mqtts:// and ssl:// 

41 :param ssl: context for mqtts:// and ssl://, the default one verifies the 

42 server certificate with the system CA certificates 

43 """ 

44 parsed_url = urlparse(url) 

45 

46 if not parsed_url.hostname: 

47 raise ValueError(f"Invalid URL, host is missing: {url}") 

48 

49 if not (factory := _CONNECTION_IMPLEMENTATION_MAP.get(parsed_url.scheme)): 

50 raise ValueError(f"Unsupported URL scheme: {parsed_url.scheme}") 

51 

52 options: ConnectionOptions = {} 

53 

54 if ssl is not None: 

55 # a context with plain tcp:// would silently send data unencrypted 

56 if parsed_url.scheme not in _TLS_SCHEMES: 

57 raise ValueError( 

58 f"SSL context requires mqtts:// or ssl:// URL, got {parsed_url.scheme}" 

59 ) 

60 

61 options["ssl"] = ssl 

62 

63 transport = await factory(parsed_url, options) 

64 

65 return MQTTConnection(transport)