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

1import struct 

2from dataclasses import dataclass 

3from enum import IntEnum 

4from typing import Final, Tuple, Union 

5 

6from zenmqtt.exceptions import MalformedPacketError 

7 

8 

9class PacketType(IntEnum): 

10 RESERVED = 0x00 

11 

12 CONNECT = 0x01 # Client to Server | Connection request 

13 CONNACK = 0x02 # Server to Client | Connect acknowledgment 

14 

15 PUBLISH = 0x03 # Client to Server or Server to Client | Publish message 

16 PUBACK = ( 

17 0x04 # Client to Server or Server to Client | Publish acknowledgment (QoS 1) 

18 ) 

19 PUBREC = 0x05 # Client to Server or Server to Client | Publish received (QoS 2 delivery part 1) 

20 PUBREL = 0x06 # Client to Server or Server to Client | Publish release (QoS 2 delivery part 2) 

21 PUBCOMP = 0x07 # Client to Server or Server to Client | Publish complete (QoS 2 delivery part 3) 

22 

23 SUBSCRIBE = 0x08 # Client to Server | Subscribe request 

24 SUBACK = 0x09 # Server to Client | Subscribe acknowledgment 

25 UNSUBSCRIBE = 0x0A # Client to Server | Unsubscribe request 

26 UNSUBACK = 0x0B # Server to Client | Unsubscribe acknowledgment 

27 

28 PINGREQ = 0x0C # Client to Server | PING request 

29 PINGRESP = 0x0D # Server to Client | PONG request 

30 

31 DISCONNECT = 0x0E # Client to Server or Server to Client | Disconnect notification 

32 

33 AUTH = 0x0F # Client to Server or Server to Client | Authentication exchange 

34 

35 

36# MQTT Control Packet 

37# +-------------------------------------------------------+ 

38# | Fixed Header, present in all MQTT Control Packets | 

39# +-------------------------------------------------------+ 

40# | Variable Header, present in some MQTT Control Packets | 

41# +-------------------------------------------------------+ 

42# | Payload, present in some MQTT Control Packets | 

43# +-------------------------------------------------------+ 

44 

45 

46FIXED_HEADER_SIZE: Final[int] = 1 

47 

48# "Remaining Length" is encoded in at most 4 bytes 

49_MAX_VARIABLE_BYTE_INTEGER_SIZE: Final[int] = 4 

50 

51_UINT16: Final[struct.Struct] = struct.Struct("!H") 

52_UINT32: Final[struct.Struct] = struct.Struct("!L") 

53 

54 

55@dataclass(frozen=True) 

56class FixedHeader: 

57 __slots__ = ("packet_type", "flags", "length") 

58 

59 packet_type: PacketType 

60 flags: int 

61 length: int 

62 

63 

64def parse_variable_byte_integer( 

65 data: Union[bytes, bytearray], offset: int = 0 

66) -> Tuple[int, int]: 

67 """ 

68 Returns the value and the number of bytes it takes. 

69 

70 :raises IndexError: data ends before the end of the value 

71 :raises MalformedPacketError: the value takes more than 4 bytes 

72 """ 

73 value = 0 

74 multiplier = 1 

75 

76 for size in range(1, _MAX_VARIABLE_BYTE_INTEGER_SIZE + 1): 

77 byte = data[offset + size - 1] 

78 value += (byte & 0x7F) * multiplier 

79 

80 if byte & 0x80 == 0: 

81 return value, size 

82 

83 multiplier *= 128 

84 

85 raise MalformedPacketError("Variable byte integer is longer than 4 bytes") 

86 

87 

88def parse_fixed_header(data: bytes) -> Tuple[FixedHeader, int]: 

89 """Returns the fixed header and its size.""" 

90 try: 

91 length, size = parse_variable_byte_integer(data, 1) 

92 except IndexError as exc: 

93 raise MalformedPacketError("Incomplete fixed header") from exc 

94 

95 header = FixedHeader( 

96 packet_type=PacketType(data[0] >> 4), flags=data[0] & 0xF, length=length 

97 ) 

98 

99 return header, FIXED_HEADER_SIZE + size 

100 

101 

102def split_packet(data: bytes) -> Tuple[FixedHeader, "BytesReader"]: 

103 """Splits a whole packet into the fixed header and a reader of the body.""" 

104 header, header_size = parse_fixed_header(data) 

105 

106 if len(data) != header_size + header.length: 

107 raise MalformedPacketError( 

108 f"Packet size {len(data)} doesn't match {header_size + header.length}" 

109 ) 

110 

111 return header, BytesReader(data[header_size:]) 

112 

113 

114class BytesReader: 

115 """ 

116 Reads values from the body of a packet. Reading past the end means that 

117 the packet is malformed. 

118 """ 

119 

120 __slots__ = ("_data", "_offset") 

121 

122 def __init__(self, data: bytes) -> None: 

123 self._data = data 

124 self._offset = 0 

125 

126 def remaining(self) -> int: 

127 return len(self._data) - self._offset 

128 

129 def read(self, size: int) -> bytes: 

130 start = self._offset 

131 end = start + size 

132 

133 if end > len(self._data): 

134 raise MalformedPacketError("Unexpected end of the packet") 

135 

136 self._offset = end 

137 return self._data[start:end] 

138 

139 def read_byte(self) -> int: 

140 if self._offset >= len(self._data): 

141 raise MalformedPacketError("Unexpected end of the packet") 

142 

143 value = self._data[self._offset] 

144 self._offset += 1 

145 return value 

146 

147 def read_uint16(self) -> int: 

148 value: int = _UINT16.unpack(self.read(2))[0] 

149 return value 

150 

151 def read_uint32(self) -> int: 

152 value: int = _UINT32.unpack(self.read(4))[0] 

153 return value 

154 

155 def read_variable_byte_integer(self) -> int: 

156 try: 

157 value, size = parse_variable_byte_integer(self._data, self._offset) 

158 except IndexError as exc: 

159 raise MalformedPacketError("Unexpected end of the packet") from exc 

160 

161 self._offset += size 

162 return value 

163 

164 def read_binary(self) -> bytes: 

165 """Binary data with 2 bytes of length.""" 

166 return self.read(self.read_uint16()) 

167 

168 def read_str(self) -> str: 

169 """UTF-8 string with 2 bytes of length.""" 

170 try: 

171 return self.read_binary().decode() 

172 except UnicodeDecodeError as exc: 

173 raise MalformedPacketError("Invalid UTF-8 string") from exc 

174 

175 def read_rest(self) -> bytes: 

176 return self.read(self.remaining()) 

177 

178 def ensure_end(self) -> None: 

179 if self.remaining(): 

180 raise MalformedPacketError( 

181 f"{self.remaining()} unexpected bytes at the end of the packet" 

182 )