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

1import struct 

2from dataclasses import dataclass 

3from enum import IntEnum 

4from typing import Any, Final, Optional, Sequence, Tuple, TypedDict, cast 

5 

6from zenmqtt.exceptions import MalformedPacketError 

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

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

9from zenmqtt.mqtt.reason_codes import ReasonCode, to_reason_code 

10from zenmqtt.mqtt.utils import pack_fixed_header, pack_str16 

11 

12 

13class QoS(IntEnum): 

14 """Quality of service levels; a value is an int, plain 0, 1, 2 work too.""" 

15 

16 AT_MOST_ONCE = 0 

17 AT_LEAST_ONCE = 1 

18 EXACTLY_ONCE = 2 

19 

20 

21_QOS_BY_VALUE: Final[tuple[QoS, ...]] = tuple(QoS) 

22 

23 

24class PublishProperties(TypedDict, total=False): 

25 payload_format_indicator: bool 

26 message_expiry_interval: int 

27 content_type: str 

28 response_topic: str 

29 # identifiers of all subscriptions which matched the message 

30 # (MQTT 5, 3.3.2.3.8); the server sends them, a client must not 

31 subscription_identifier: Sequence[int] 

32 topic_alias: int 

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

34 

35 

36class PubackProperties(TypedDict, total=False): 

37 reason_string: str 

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

39 

40 

41class PubrecProperties(TypedDict, total=False): 

42 reason_string: str 

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

44 

45 

46class PubrelProperties(TypedDict, total=False): 

47 reason_string: str 

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

49 

50 

51class PubcompProperties(TypedDict, total=False): 

52 reason_string: str 

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

54 

55 

56def pack_publish_packet( 

57 packet_identifier: int, 

58 topic: str, 

59 payload: bytes, 

60 qos: int, 

61 retain: bool, 

62 dup: bool, 

63 properties: PublishProperties, 

64) -> bytes: 

65 if qos not in (0, 1, 2): 

66 raise ValueError(f"Invalid QoS: {qos}") 

67 

68 packed_topic = pack_str16(topic) 

69 packed_identifier = struct.pack("!H", packet_identifier) if qos else b"" 

70 packed_properties = pack_properties(cast(Properties, properties)) 

71 

72 fixed_header = pack_fixed_header( 

73 PacketType.PUBLISH, 

74 flags=(dup & 1) << 3 | (qos << 1) | retain & 1, 

75 length=len(packed_topic) 

76 + len(packed_identifier) 

77 + len(packed_properties) 

78 + len(payload), 

79 ) 

80 

81 # the payload is copied once 

82 return b"".join( 

83 (fixed_header, packed_topic, packed_identifier, packed_properties, payload) 

84 ) 

85 

86 

87@dataclass(frozen=True) 

88class MQTTMessage: 

89 __slots__ = ( 

90 "packet_identifier", 

91 "dup", 

92 "qos", 

93 "retain", 

94 "payload", 

95 "properties", 

96 "topic", 

97 ) 

98 dup: bool 

99 qos: QoS 

100 retain: bool 

101 

102 packet_identifier: Optional[int] 

103 topic: str 

104 payload: bytes 

105 

106 properties: PublishProperties 

107 

108 

109def parse_publish_packet(fixed_header: FixedHeader, reader: BytesReader) -> MQTTMessage: 

110 dup = bool(fixed_header.flags & 0x8) 

111 qos = (fixed_header.flags & 0x6) >> 1 

112 retain = bool(fixed_header.flags & 0x01) 

113 

114 if qos == 3: 

115 raise MalformedPacketError("PUBLISH with QoS 3") 

116 

117 topic = reader.read_str() 

118 packet_identifier = reader.read_uint16() if qos else None 

119 properties = parse_properties(reader) 

120 

121 return MQTTMessage( 

122 dup=dup, 

123 qos=_QOS_BY_VALUE[qos], 

124 retain=retain, 

125 packet_identifier=packet_identifier, 

126 topic=topic, 

127 payload=reader.read_rest(), 

128 properties=cast(PublishProperties, properties), 

129 ) 

130 

131 

132@dataclass(frozen=True) 

133class PubAckResult: 

134 __slots__ = ("packet_identifier", "reason_code", "properties") 

135 

136 packet_identifier: int 

137 reason_code: ReasonCode 

138 properties: PubackProperties 

139 

140 

141def parse_puback_packet(fixed_header: FixedHeader, reader: BytesReader) -> PubAckResult: 

142 packet_identifier, reason_code, properties = _parse_publish_response_packet( 

143 fixed_header, reader 

144 ) 

145 

146 return PubAckResult( 

147 packet_identifier=packet_identifier, 

148 reason_code=reason_code, 

149 properties=properties, 

150 ) 

151 

152 

153def pack_puback_packet( 

154 packet_identifier: int, reason_code: int, properties: PubackProperties 

155) -> bytes: 

156 return _pack_publish_response_packet( 

157 PacketType.PUBACK, 0x0, packet_identifier, reason_code, properties 

158 ) 

159 

160 

161@dataclass(frozen=True) 

162class PubRecResult: 

163 __slots__ = ("packet_identifier", "reason_code", "properties") 

164 packet_identifier: int 

165 reason_code: ReasonCode 

166 properties: PubrecProperties 

167 

168 

169def parse_pubrec_packet(fixed_header: FixedHeader, reader: BytesReader) -> PubRecResult: 

170 packet_identifier, reason_code, properties = _parse_publish_response_packet( 

171 fixed_header, reader 

172 ) 

173 

174 return PubRecResult( 

175 packet_identifier=packet_identifier, 

176 reason_code=reason_code, 

177 properties=properties, 

178 ) 

179 

180 

181def pack_pubrec_packet( 

182 packet_identifier: int, reason_code: int, properties: PubrecProperties 

183) -> bytes: 

184 return _pack_publish_response_packet( 

185 PacketType.PUBREC, 0x0, packet_identifier, reason_code, properties 

186 ) 

187 

188 

189@dataclass(frozen=True) 

190class PubRelResult: 

191 __slots__ = ("packet_identifier", "reason_code", "properties") 

192 

193 packet_identifier: int 

194 reason_code: ReasonCode 

195 properties: PubrelProperties 

196 

197 

198def parse_pubrel_packet(fixed_header: FixedHeader, reader: BytesReader) -> PubRelResult: 

199 packet_identifier, reason_code, properties = _parse_publish_response_packet( 

200 fixed_header, reader 

201 ) 

202 

203 return PubRelResult( 

204 packet_identifier=packet_identifier, 

205 reason_code=reason_code, 

206 properties=properties, 

207 ) 

208 

209 

210def pack_pubrel_packet( 

211 packet_identifier: int, reason_code: int, properties: PubrelProperties 

212) -> bytes: 

213 return _pack_publish_response_packet( 

214 PacketType.PUBREL, 0x2, packet_identifier, reason_code, properties 

215 ) 

216 

217 

218@dataclass(frozen=True) 

219class PubCompResult: 

220 __slots__ = ("packet_identifier", "reason_code", "properties") 

221 

222 packet_identifier: int 

223 reason_code: ReasonCode 

224 properties: PubcompProperties 

225 

226 

227def parse_pubcomp_packet( 

228 fixed_header: FixedHeader, reader: BytesReader 

229) -> PubCompResult: 

230 packet_identifier, reason_code, properties = _parse_publish_response_packet( 

231 fixed_header, reader 

232 ) 

233 

234 return PubCompResult( 

235 packet_identifier=packet_identifier, 

236 reason_code=reason_code, 

237 properties=properties, 

238 ) 

239 

240 

241def pack_pubcomp_packet( 

242 packet_identifier: int, reason_code: int, properties: PubcompProperties 

243) -> bytes: 

244 return _pack_publish_response_packet( 

245 PacketType.PUBCOMP, 0x0, packet_identifier, reason_code, properties 

246 ) 

247 

248 

249def _parse_publish_response_packet( 

250 fixed_header: FixedHeader, reader: BytesReader 

251) -> Tuple[int, ReasonCode, Properties]: 

252 packet_identifier = reader.read_uint16() 

253 # both the reason code and properties may be omitted 

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

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

256 

257 reader.ensure_end() 

258 

259 return packet_identifier, reason_code, properties 

260 

261 

262def _pack_publish_response_packet( 

263 packet_type: PacketType, 

264 flags: int, 

265 packet_identifier: int, 

266 reason_code: int, 

267 properties: Any, 

268) -> bytes: 

269 variable_header = struct.pack("!H", packet_identifier) 

270 

271 # the reason code and properties may be omitted if they are empty 

272 if reason_code or properties: 

273 variable_header += struct.pack("!B", reason_code) 

274 

275 if properties: 

276 variable_header += pack_properties(cast(Properties, properties)) 

277 

278 return pack_fixed_header(packet_type, flags, len(variable_header)) + variable_header 

279 

280 

281# PUBACK for QoS 1; PUBCOMP (or PUBREC with an error reason code) for QoS 2 

282PublishAcknowledgement = PubAckResult | PubRecResult | PubCompResult