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

1import struct 

2from dataclasses import dataclass 

3from typing import Final, Sequence, Tuple, TypedDict, Union, 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_codes 

8from zenmqtt.mqtt.utils import pack_fixed_header, pack_str16 

9 

10 

11class SubscriptionProperties(TypedDict, total=False): 

12 subscription_identifier: int 

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

14 

15 

16class SubackProperties(TypedDict, total=False): 

17 reason_string: str 

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

19 

20 

21class UnsubscribeProperties(TypedDict, total=False): 

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

23 

24 

25@dataclass(frozen=True) 

26class SubscribeResult: 

27 __slots__ = ("packet_identifier", "properties", "reason_codes") 

28 

29 packet_identifier: int 

30 properties: SubackProperties 

31 reason_codes: Sequence[ReasonCode] 

32 

33 

34@dataclass(frozen=True) 

35class UnsubscribeResult: 

36 __slots__ = ("packet_identifier", "properties", "reason_codes") 

37 

38 packet_identifier: int 

39 properties: UnsubscribeProperties 

40 reason_codes: Sequence[ReasonCode] 

41 

42 

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

44class Subscription: 

45 """ 

46 Topic filter with MQTT 5 subscription options. 

47 

48 :param no_local: don't receive messages published by this client 

49 :param retain_as_published: keep the RETAIN flag of forwarded messages 

50 :param retain_handling: 0 - send retained messages on subscribe, 

51 1 - only if the subscription is new, 2 - don't send them 

52 """ 

53 

54 topic: str 

55 qos: int = 0 

56 no_local: bool = False 

57 retain_as_published: bool = False 

58 retain_handling: int = 0 

59 

60 

61# (topic, qos) is the short form of Subscription(topic, qos) 

62SubscriptionRequest = Union[Subscription, Tuple[str, int]] 

63 

64_SHARED_SUBSCRIPTION_PREFIX: Final[str] = "$share/" 

65 

66# subscription options byte 

67_NO_LOCAL_FLAG: Final[int] = 0x04 

68_RETAIN_AS_PUBLISHED_FLAG: Final[int] = 0x08 

69_RETAIN_HANDLING_SHIFT: Final[int] = 4 

70 

71 

72def pack_subscription_options(subscription: Subscription) -> int: 

73 if subscription.qos not in (0, 1, 2): 

74 raise ValueError(f"Invalid QoS: {subscription.qos}") 

75 

76 if subscription.retain_handling not in (0, 1, 2): 

77 raise ValueError(f"Invalid retain handling: {subscription.retain_handling}") 

78 

79 if subscription.no_local and subscription.topic.startswith( 

80 _SHARED_SUBSCRIPTION_PREFIX 

81 ): 

82 raise ValueError("No Local can't be used with a shared subscription") 

83 

84 options = subscription.qos | subscription.retain_handling << _RETAIN_HANDLING_SHIFT 

85 

86 if subscription.no_local: 

87 options |= _NO_LOCAL_FLAG 

88 

89 if subscription.retain_as_published: 

90 options |= _RETAIN_AS_PUBLISHED_FLAG 

91 

92 return options 

93 

94 

95def to_subscription(request: SubscriptionRequest) -> Subscription: 

96 return request if isinstance(request, Subscription) else Subscription(*request) 

97 

98 

99def pack_subscribe_packet( 

100 packet_identifier: int, 

101 topics: Sequence[SubscriptionRequest], 

102 properties: SubscriptionProperties, 

103) -> bytes: 

104 parts = [ 

105 struct.pack("!H", packet_identifier), 

106 pack_properties(cast(Properties, properties)), 

107 ] 

108 

109 for request in topics: 

110 subscription = to_subscription(request) 

111 parts += ( 

112 pack_str16(subscription.topic), 

113 struct.pack("!B", pack_subscription_options(subscription)), 

114 ) 

115 

116 body = b"".join(parts) 

117 

118 return pack_fixed_header(PacketType.SUBSCRIBE, 0x2, len(body)) + body 

119 

120 

121def parse_suback_packet( 

122 fixed_header: FixedHeader, reader: BytesReader 

123) -> SubscribeResult: 

124 packet_identifier, properties, reason_codes = _parse_packet(fixed_header, reader) 

125 

126 return SubscribeResult( 

127 packet_identifier=packet_identifier, 

128 properties=cast(SubackProperties, properties), 

129 reason_codes=reason_codes, 

130 ) 

131 

132 

133def pack_unsubscribe_packet( 

134 packet_identifier: int, topics: Sequence[str], properties: UnsubscribeProperties 

135) -> bytes: 

136 body = b"".join( 

137 ( 

138 struct.pack("!H", packet_identifier), 

139 pack_properties(cast(Properties, properties)), 

140 *map(pack_str16, topics), 

141 ) 

142 ) 

143 

144 return pack_fixed_header(PacketType.UNSUBSCRIBE, 0x2, len(body)) + body 

145 

146 

147def parse_unsuback_packet( 

148 fixed_header: FixedHeader, reader: BytesReader 

149) -> UnsubscribeResult: 

150 packet_identifier, properties, reason_codes = _parse_packet(fixed_header, reader) 

151 

152 return UnsubscribeResult( 

153 packet_identifier=packet_identifier, 

154 properties=cast(UnsubscribeProperties, properties), 

155 reason_codes=reason_codes, 

156 ) 

157 

158 

159def _parse_packet( 

160 fixed_header: FixedHeader, reader: BytesReader 

161) -> Tuple[int, Properties, Sequence[ReasonCode]]: 

162 packet_identifier = reader.read_uint16() 

163 properties = parse_properties(reader) 

164 # one reason code per topic 

165 reason_codes = to_reason_codes(reader.read_rest()) 

166 

167 return packet_identifier, properties, reason_codes