Coverage for /root/GitHubProjects/impacket/impacket/mqtt.py: 26%

# Coverage for /root/GitHubProjects/impacket/impacket/mqtt.py : 26%

## 244 statements

Hot-keys on this page

r m x p toggle line displays

j k next/prev highlighted chunk

0 (zero) top of page

1 (one) first highlighted chunk

1#!/usr/bin/env python

2# Impacket - Collection of Python classes for working with network protocols.

3#

4# SECUREAUTH LABS. Copyright (C) 2018 SecureAuth Corporation. All rights reserved.

5#

6# This software is provided under a slightly modified version

7# of the Apache Software License. See the accompanying LICENSE file

8# for more information.

9#

10# Description:

11# Minimalistic MQTT implementation, just focused on connecting, subscribing and publishing basic

12# messages on topics.

13#

14# Author:

15# Alberto Solino (@agsolino)

16#

17# References:

18# - https://docs.oasis-open.org/mqtt/mqtt/v3.1.1/mqtt-v3.1.1.html

19#

20# ToDo:

21# [ ] Implement all the MQTT Control Packets and operations

22# [ ] Implement QoS = QOS_ASSURED_DELIVERY when publishing messages

23#

24

25 from __future__ import print_function

26 import logging

27 import struct

28 import socket

29 from impacket.structure import Structure

30 try:

31 from OpenSSL import SSL

32 except ImportError:

33 logging.critical("pyOpenSSL is not installed, can't continue")

34 raise

35

36# Packet Types

37 PACKET_CONNECT = 1 << 4

38 PACKET_CONNACK = 2 << 4

39 PACKET_PUBLISH = 3 << 4

40 PACKET_PUBACK = 4 << 4

41 PACKET_PUBREC = 5 << 4

42 PACKET_PUBREL = 6 << 4

43 PACKET_PUBCOMP = 7 << 4

44 PACKET_SUBSCRIBE = 8 << 4

45 PACKET_SUBSCRIBEACK = 9 << 4

46 PACKET_UNSUBSCRIBE = 10 << 4

47 PACKET_UNSUBACK = 11 << 4

48 PACKET_PINGREQ = 12 << 4

49 PACKET_PINGRESP = 13 << 4

50 PACKET_DISCONNECT = 14 << 4

51

52# CONNECT Flags

53 CONNECT_USERNAME = 0x80

54 CONNECT_PASSWORD = 0x40

55 CONNECT_CLEAN_SESSION = 0x2

56

57# CONNECT_ACK Return Errors

58 CONNECT_ACK_ERROR_MSGS = {

59 0x00: 'Connection Accepted',

60 0x01: 'Connection Refused, unacceptable protocol version',

61 0x02: 'Connection Refused, identifier rejected',

62 0x03: 'Connection Refused, Server unavailable',

63 0x04: 'Connection Refused, bad user name or password',

64 0x05: 'Connection Refused, not authorized'

65}

66

67# QoS Levels

68 QOS_FIRE_AND_FORGET = 0

69 QOS_ACK_DELIVERY = 1

70 QOS_ASSURED_DELIVERY= 2

71

72 class MQTT_Packet(Structure):

73 commonHdr= (

74 ('PacketType','B=0'),

75 ('MessageLength','<L=0'),

76 )

77 structure = (

78 ('_VariablePart', '_-VariablePart', 'self["MessageLength"]'),

79 ('VariablePart', ':'),

80 )

81 def setQoS(self, QoS):

82 self['PacketType'] |= (QoS << 1)

83

84 def fromString(self, data):

85 if data is not None and len(data) > 2:

86 # Get the Length

87 index = 1

88 multiplier = 1

89 value = 0

90 encodedByte = 128

91 packetType = data[0]

92 while (encodedByte & 128) != 0:

93 encodedByte = ord(data[index])

94 value += (encodedByte & 127) * multiplier

95 multiplier *= 128

96 index += 1

97 if multiplier > 128 * 128 * 128:

98 raise Exception('Malformed Remaining Length')

99 data = packetType + struct.pack('<L', value) + data[index: value+index]

100 return Structure.fromString(self, data)

101 raise Exception('Dont know')

102

103 def getData(self):

104 packetType = self['PacketType']

105 self.commonHdr = ()

106 packetLen = len(Structure.getData(self))

107 output = ''

108 while packetLen > 0:

109 encodedByte = packetLen % 128

110 packetLen /= 128

111 if packetLen > 0:

112 encodedByte |= 128

113 output += chr(encodedByte)

114 self.commonHdr = ( ('PacketType','B=0'), ('MessageLength',':'), )

115 self['PacketType'] = packetType

116 self['MessageLength'] = output

117 if output == '':

118 self['MessageLength'] = chr(00)

119

120 return Structure.getData(self)

121

122

123 class MQTT_String(Structure):

124 structure = (

125 ('Length','>H-Name'),

126 ('Name',':'),

127 )

128

129 class MQTT_Connect(MQTT_Packet):

130 structure = (

131 ('ProtocolName',':', MQTT_String),

132 ('Version','B=3'),

133 ('Flags','B=2'),

134 ('KeepAlive','>H=60'),

135 ('ClientID',':', MQTT_String),

136 ('Payload',':=""'),

137 )

138 def __init__(self, data = None, alignment = 0):

139 MQTT_Packet.__init__(self, data, alignment)

140 if data is None:

141 self['PacketType'] = PACKET_CONNECT

142

143 class MQTT_ConnectAck(MQTT_Packet):

144 structure = (

145 ('ReturnCode', '>H=0'),

146 )

147

148 class MQTT_Publish(MQTT_Packet):

149 structure = (

150 ('Topic',':', MQTT_String),

151 ('Message',':'),

152 )

153 def __init__(self, data = None, alignment = 0):

154 MQTT_Packet.__init__(self, data, alignment)

155 if data is None:

156 self['PacketType'] = PACKET_PUBLISH

157

158 def getData(self):

159 if self['PacketType'] & 6 > 0:

160 # We have QoS enabled, we need to have a MessageID field

161 self.structure = (

162 ('Topic', ':', MQTT_String),

163 ('MessageID', '>H=0'),

164 ('Message', ':'),

165 )

166 return MQTT_Packet.getData(self)

167

168 class MQTT_Disconnect(MQTT_Packet):

169 structure = (

170 )

171 def __init__(self, data=None, alignment=0):

172 MQTT_Packet.__init__(self, data, alignment)

173 if data is None:

174 self['PacketType'] = PACKET_DISCONNECT

175

176 class MQTT_Subscribe(MQTT_Packet):

177 structure = (

178 ('MessageID','>H=1'),

179 ('Topic',':', MQTT_String),

180 ('Flags','B=0'),

181 )

182 def __init__(self, data = None, alignment = 0):

183 MQTT_Packet.__init__(self, data, alignment)

184 if data is None:

185 self['PacketType'] = PACKET_SUBSCRIBE

186

187 class MQTT_SubscribeACK(MQTT_Packet):

188 structure = (

189 ('MessageID','>H=0'),

190 ('ReturnCode','B=0'),

191 )

192 def __init__(self, data = None, alignment = 0):

193 MQTT_Packet.__init__(self, data, alignment)

194 if data is None:

195 self['PacketType'] = PACKET_SUBSCRIBEACK

196

197 class MQTT_UnSubscribe(MQTT_Packet):

198 structure = (

199 ('MessageID','>H=1'),

200 ('Topics',':'),

201 )

202 def __init__(self, data = None, alignment = 0):

203 MQTT_Packet.__init__(self, data, alignment)

204 if data is None:

205 self['PacketType'] = PACKET_UNSUBSCRIBE

206

207 class MQTTSessionError(Exception):

208 """

209 This is the exception every client should catch

210 """

211

212 def __init__(self, error=0, packet=0, errorString=''):

213 Exception.__init__(self)

214 self.error = error

215 self.packet = packet

216 self.errorString = errorString

217

218 def getErrorCode(self):

219 return self.error

220

221 def getErrorPacket(self):

222 return self.packet

223

224 def getErrorString(self):

225 return self.errorString

226

227 def __str__(self):

228 return self.errorString

229

230 class MQTTConnection:

231 def __init__(self, host, port, isSSL=False):

232 self._targetHost = host

233 self._targetPort = port

234 self._isSSL = isSSL

235 self._socket = None

236 self._messageId = 1

237 self.connectSocket()

238

239 def getSocket(self):

240 return self._socket

241

242 def connectSocket(self):

243 s = socket.socket()

244 s.connect((self._targetHost, int(self._targetPort)))

245

246 if self._isSSL is True:

247 ctx = SSL.Context(SSL.TLSv1_METHOD)

248 self._socket = SSL.Connection(ctx, s)

249 self._socket.set_connect_state()

250 self._socket.do_handshake()

251 else:

252 self._socket = s

253

254 def send(self, request):

255 return self._socket.sendall(str(request))

256

257 def sendReceive(self, request):

258 self.send(request)

259 return self.recv()

260

261 def recv(self):

262 REQUEST_SIZE = 8192

263 data = ''

264 done = False

265 while not done:

266 recvData = self._socket.recv(REQUEST_SIZE)

267 if len(recvData) < REQUEST_SIZE:

268 done = True

269 data += recvData

270

271 response = []

272 while len(data) > 0:

273 try:

274 message = MQTT_Packet(data)

275 remaining = data[len(message):]

276 except Exception:

277 # We need more data

278 remaining = data + self._socket.recv(REQUEST_SIZE)

279 else:

280 response.append(message)

281 data = remaining

282

283 self._messageId += 1

284 return response

285

286 def connect(self, clientId = ' ', username = None, password = None, protocolName = 'MQIsdp', version = 3, flags = CONNECT_CLEAN_SESSION, keepAlive = 60):

287 """

288

289 :param clientId: Whatever cliend Id that represents you

290 :param username: if None, anonymous connection will be attempted

291 :param password: if None, anonymous connection will be attempted

292 :param protocolName: specification states default should be 'MQTT' but some brokers might expect 'MQIsdp'

293 :param version: Allowed versions are 3 or 4 (some brokers might like 4)

294 :param flags:

295 :param keepAlive: default 60

296 :return: True or MQTTSessionError if something went wrong

297 """

298

299 # Let's build the packet

300 connectPacket = MQTT_Connect()

301 connectPacket['Version'] = version

302 connectPacket['Flags'] = flags

303 connectPacket['KeepAlive'] = keepAlive

304 connectPacket['ProtocolName'] = MQTT_String()

305 connectPacket['ProtocolName']['Name'] = protocolName

306

307 connectPacket['ClientID'] = MQTT_String()

308 connectPacket['ClientID']['Name'] = clientId

309

310 if username is not None:

311 connectPacket['Flags'] |= CONNECT_USERNAME | CONNECT_PASSWORD

312 if username is None:

313 user = ''

314 else:

315 user = username

316 if password is None:

317 pwd = ''

318 else:

319 pwd = password

320

321 username = MQTT_String()

322 username['Name'] = user

323 password = MQTT_String()

324 password['Name'] = pwd

325 connectPacket['Payload'] = str(username) + str(password)

326

327 data= self.sendReceive(connectPacket)[0]

328

329 response = MQTT_ConnectAck(str(data))

330 if response['ReturnCode'] != 0:

331 raise MQTTSessionError(error = response['ReturnCode'], errorString = CONNECT_ACK_ERROR_MSGS[response['ReturnCode']] )

332

333 return True

334

335 def subscribe(self, topic, messageID = 1, flags = 0, QoS = 1):

336 """

337

338 :param topic: Topic name you want to subscribe to

339 :param messageID: optional messageId

340 :param flags: Message flags

341 :param QoS: define the QoS requested

342 :return: True or MQTTSessionError if something went wrong

343 """

344 subscribePacket = MQTT_Subscribe()

345 subscribePacket['MessageID'] = messageID

346 subscribePacket['Topic'] = MQTT_String()

347 subscribePacket['Topic']['Name'] = topic

348 subscribePacket['Flags'] = flags

349 subscribePacket.setQoS(QoS)

350

351 try:

352 data = self.sendReceive(subscribePacket)[0]

353 except Exception as e:

354 raise MQTTSessionError(errorString=str(e))

355

356 subAck = MQTT_SubscribeACK(str(data))

357

358 if subAck['ReturnCode'] > 2:

359 raise MQTTSessionError(errorString = 'Failure to subscribe')

360

361 return True

362

363 def unSubscribe(self, topic, messageID = 1, QoS = 0):

364 """

365 Unsubscribes from a topic

366

367 :param topic:

368 :param messageID:

369 :param QoS: define the QoS requested

370 :return:

371 """

372 # ToDo: Support more than one topic

373 packet = MQTT_UnSubscribe()

374 packet['MessageID'] = messageID

375 packet['Topics'] = MQTT_String()

376 packet['Topics']['Name'] = topic

377 packet.setQoS( QoS )

378

379 return self.sendReceive(packet)

380

381 def publish(self, topic, message, messageID = 1, QoS=0):

382

383 packet = MQTT_Publish()

384 packet['Topic'] = MQTT_String()

385 packet['Topic']['Name'] = topic

386 packet['Message'] = message

387 packet['MessageID'] = messageID

388 packet.setQoS( QoS )

389

390 return self.sendReceive(packet)

391

392 def disconnect(self):

393 return self.send(str(MQTT_Disconnect()))

394

395 if __name__ == '__main__': 395 ↛ 396line 395 didn't jump to line 396, because the condition on line 395 was never true

396 HOST = '192.168.45.162'

397 USER = 'test'

398 PASS = 'test'

399

400 mqtt = MQTTConnection(HOST, 1883, False)

401 mqtt.connect('secure-', username=USER, password=PASS, version = 3)

402 #mqtt.connect(protocolName='MQTT', version=4)

403 #mqtt.connect()

404

405 #mqtt.subscribe('/test/beto')

406 #mqtt.unSubscribe('/test/beto')

407 #mqtt.publish('/test/beto', 'Hey There, I"d like to talk to you', QoS=1)

408 mqtt.subscribe('$SYS/#')

409

410

411 while True:

412

413 packets = mqtt.recv()

414 for packet in packets:

415 publish = MQTT_Publish(str(packet))

416 print('%s -> %s' % (publish['Topic']['Name'], publish['Message']))

417

418 mqtt.disconnect()