Coverage for zenmqtt/mqtt/connect.py: 100%

1import struct 

2from dataclasses import dataclass, field 

3from typing import Final, Optional, Sequence, Tuple, TypedDict, cast 

4 

5from zenmqtt.mqtt.packet import BytesReader, FixedHeader, PacketType 

6from zenmqtt.mqtt.properties import Properties, pack_properties, parse_properties 

7from zenmqtt.mqtt.reason_codes import ReasonCode, to_reason_code 

8from zenmqtt.mqtt.utils import pack_binary, pack_fixed_header, pack_str16 

9 

10 

11class ConnackProperties(TypedDict, total=False): 

12 session_expiry_interval: int 

13 assigned_client_identifier: str 

14 server_keep_alive: int 

15 authentication_method: str 

16 authentication_data: bytes 

17 response_information: str 

18 server_reference: str 

19 reason_string: str 

20 receive_maximum: int 

21 topic_alias_maximum: int 

22 maximum_qos: int 

23 retain_available: bool 

24 user_property: Sequence[Tuple[str, str]] 

25 maximum_packet_size: int 

26 wildcard_subscription_available: bool 

27 subscription_identifier_available: bool 

28 shared_subscription_available: bool 

29 

30 

31# CONNACK "Session Present" flag 

32SESSION_PRESENT_FLAG: Final[int] = 0x01 

33 

34 

35@dataclass(frozen=True) 

36class ConnectionResult: 

37 __slots__ = ("flags", "reason_code", "properties") 

38 flags: int 

39 reason_code: ReasonCode 

40 properties: ConnackProperties 

41 

42 @property 

43 def session_present(self) -> bool: 

44 """The server has a session of the client from a previous connection.""" 

45 return bool(self.flags & SESSION_PRESENT_FLAG) 

46 

47 

48def parse_connack_packet( 

49 fixed_header: FixedHeader, reader: BytesReader 

50) -> ConnectionResult: 

51 flags = reader.read_byte() 

52 reason_code = to_reason_code(reader.read_byte()) 

53 properties = parse_properties(reader) 

54 

55 reader.ensure_end() 

56 

57 return ConnectionResult(flags=flags, reason_code=reason_code, properties=properties) 

58 

59 

60class ConnectProperties(TypedDict, total=False): 

61 session_expiry_interval: int 

62 authentication_method: str 

63 authentication_data: bytes 

64 request_problem_information: bool 

65 request_response_information: bool 

66 receive_maximum: int 

67 topic_alias_maximum: int 

68 user_property: Sequence[Tuple[str, str]] 

69 maximum_packet_size: int 

70 

71 

72class WillProperties(TypedDict, total=False): 

73 will_delay_interval: int 

74 payload_format_indicator: bool 

75 message_expiry_interval: int 

76 content_type: str 

77 response_topic: str 

78 correlation_data: bytes 

79 user_property: Sequence[Tuple[str, str]] 

80 

81 

82def _no_will_properties() -> WillProperties: 

83 return {} 

84 

85 

86@dataclass(frozen=True, slots=True) 

87class WillMessage: 

88 """ 

89 The message which the server publishes when the connection is closed 

90 without DISCONNECT (or with reason 0x04 "Disconnect with Will Message"). 

91 """ 

92 

93 topic: str 

94 payload: bytes 

95 qos: int = 0 

96 retain: bool = False 

97 properties: WillProperties = field(default_factory=_no_will_properties) 

98 

99 

100MAX_KEEP_ALIVE: Final[int] = 2**16 - 1 

101 

102# CONNECT flags 

103_CLEAN_START_FLAG: Final[int] = 0x02 

104_WILL_FLAG: Final[int] = 0x04 

105_WILL_QOS_SHIFT: Final[int] = 3 

106_WILL_RETAIN_FLAG: Final[int] = 0x20 

107_PASSWORD_FLAG: Final[int] = 0x40 

108_USERNAME_FLAG: Final[int] = 0x80 

109 

110 

111def pack_connect_packet( 

112 client_id: str, 

113 username: Optional[str], 

114 password: Optional[str], 

115 *, 

116 clean_start: bool, 

117 keepalive: int, 

118 properties: ConnectProperties, 

119 will: Optional[WillMessage] = None, 

120) -> bytes: 

121 if not 0 <= keepalive <= MAX_KEEP_ALIVE: 

122 raise ValueError(f"Invalid keep alive: {keepalive}") 

123 

124 connect_flags = 0 

125 

126 if clean_start: 

127 connect_flags |= _CLEAN_START_FLAG 

128 

129 will_bytes = b"" 

130 

131 if will: 

132 if will.qos not in (0, 1, 2): 

133 raise ValueError(f"Invalid will QoS: {will.qos}") 

134 

135 connect_flags |= _WILL_FLAG | (will.qos << _WILL_QOS_SHIFT) 

136 

137 if will.retain: 

138 connect_flags |= _WILL_RETAIN_FLAG 

139 

140 will_bytes = b"".join( 

141 ( 

142 pack_properties(cast(Properties, will.properties)), 

143 pack_str16(will.topic), 

144 pack_binary(will.payload), 

145 ) 

146 ) 

147 

148 username_bytes = password_bytes = b"" 

149 

150 # MQTT 5 allows a password without a username, both may be empty 

151 if username is not None: 

152 connect_flags |= _USERNAME_FLAG 

153 username_bytes = pack_str16(username) 

154 

155 if password is not None: 

156 connect_flags |= _PASSWORD_FLAG 

157 password_bytes = pack_str16(password) 

158 

159 variable_header = struct.pack( 

160 "!H4sBBH", 4, b"MQTT", 5, connect_flags, keepalive 

161 ) + pack_properties(cast(Properties, properties)) 

162 

163 # the order of the payload fields is defined by the spec 

164 payload = b"".join( 

165 (pack_str16(client_id), will_bytes, username_bytes, password_bytes) 

166 ) 

167 

168 return ( 

169 pack_fixed_header(PacketType.CONNECT, 0x00, len(variable_header) + len(payload)) 

170 + variable_header 

171 + payload 

172 ) 

173 

174 

175class DisconnectProperties(TypedDict, total=False): 

176 session_expiry_interval: int 

177 server_reference: str 

178 reason_string: str 

179 user_property: Sequence[Tuple[str, str]] 

180 

181 

182def pack_disconnect_packet(reason_code: int, properties: DisconnectProperties) -> bytes: 

183 properties_bytes = pack_properties(cast(Properties, properties)) 

184 payload_length = 1 + len(properties_bytes) 

185 payload = struct.pack("!B", reason_code) + properties_bytes 

186 

187 return pack_fixed_header(PacketType.DISCONNECT, 0x00, payload_length) + payload 

188 

189 

190@dataclass(frozen=True) 

191class ServerDisconnect: 

192 __slots__ = ("reason_code", "properties") 

193 

194 reason_code: ReasonCode 

195 properties: DisconnectProperties 

196 

197 

198def parse_disconnect_packet( 

199 fixed_header: FixedHeader, reader: BytesReader 

200) -> ServerDisconnect: 

201 # both the reason code and properties may be omitted 

202 reason_code = to_reason_code(reader.read_byte() if reader.remaining() else 0x0) 

203 properties = parse_properties(reader) if reader.remaining() else {} 

204 

205 reader.ensure_end() 

206 

207 return ServerDisconnect( 

208 reason_code=reason_code, 

209 properties=cast(DisconnectProperties, properties), 

210 )