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

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

## 243 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# SECUREAUTH LABS. Copyright 2018 SecureAuth Corporation. All rights reserved.

3#

4# This software is provided under under a slightly modified version

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

6# for more information.

7#

8# Author: Alberto Solino (@agsolino)

9#

10# Description:

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

12# messages on topics.

13#

14# References:

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

16#

17# ToDo:

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

19# [ ] Implement QoS = QOS_ASSURED_DELIVERY when publishing messages

20#

21 from __future__ import print_function

22 import logging

23 import struct

24 import socket

25 from impacket.structure import Structure

26 try:

27 from OpenSSL import SSL

28 except ImportError:

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

30 raise

31

32# Packet Types

33 PACKET_CONNECT = 1 << 4

34 PACKET_CONNACK = 2 << 4

35 PACKET_PUBLISH = 3 << 4

36 PACKET_PUBACK = 4 << 4

37 PACKET_PUBREC = 5 << 4

38 PACKET_PUBREL = 6 << 4

39 PACKET_PUBCOMP = 7 << 4

40 PACKET_SUBSCRIBE = 8 << 4

41 PACKET_SUBSCRIBEACK = 9 << 4

42 PACKET_UNSUBSCRIBE = 10 << 4

43 PACKET_UNSUBACK = 11 << 4

44 PACKET_PINGREQ = 12 << 4

45 PACKET_PINGRESP = 13 << 4

46 PACKET_DISCONNECT = 14 << 4

47

48# CONNECT Flags

49 CONNECT_USERNAME = 0x80

50 CONNECT_PASSWORD = 0x40

51 CONNECT_CLEAN_SESSION = 0x2

52

53# CONNECT_ACK Return Errors

54 CONNECT_ACK_ERROR_MSGS = {

55 0x00: 'Connection Accepted',

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

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

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

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

60 0x05: 'Connection Refused, not authorized'

61}

62

63# QoS Levels

64 QOS_FIRE_AND_FORGET = 0

65 QOS_ACK_DELIVERY = 1

66 QOS_ASSURED_DELIVERY= 2

67

68 class MQTT_Packet(Structure):

69 commonHdr= (

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

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

72 )

73 structure = (

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

75 ('VariablePart', ':'),

76 )

77 def setQoS(self, QoS):

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

79

80 def fromString(self, data):

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

82 # Get the Length

83 index = 1

84 multiplier = 1

85 value = 0

86 encodedByte = 128

87 packetType = data[0]

88 while (encodedByte & 128) != 0:

89 encodedByte = ord(data[index])

90 value += (encodedByte & 127) * multiplier

91 multiplier *= 128

92 index += 1

93 if multiplier > 128 * 128 * 128:

94 raise Exception('Malformed Remaining Length')

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

96 return Structure.fromString(self, data)

97 raise Exception('Dont know')

98

99 def getData(self):

100 packetType = self['PacketType']

101 self.commonHdr = ()

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

103 output = ''

104 while packetLen > 0:

105 encodedByte = packetLen % 128

106 packetLen /= 128

107 if packetLen > 0:

108 encodedByte |= 128

109 output += chr(encodedByte)

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

111 self['PacketType'] = packetType

112 self['MessageLength'] = output

113 if output == '':

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

115

116 return Structure.getData(self)

117

118

119 class MQTT_String(Structure):

120 structure = (

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

122 ('Name',':'),

123 )

124

125 class MQTT_Connect(MQTT_Packet):

126 structure = (

127 ('ProtocolName',':', MQTT_String),

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

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

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

131 ('ClientID',':', MQTT_String),

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

133 )

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

135 MQTT_Packet.__init__(self, data, alignment)

136 if data is None:

137 self['PacketType'] = PACKET_CONNECT

138

139 class MQTT_ConnectAck(MQTT_Packet):

140 structure = (

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

142 )

143

144 class MQTT_Publish(MQTT_Packet):

145 structure = (

146 ('Topic',':', MQTT_String),

147 ('Message',':'),

148 )

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

150 MQTT_Packet.__init__(self, data, alignment)

151 if data is None:

152 self['PacketType'] = PACKET_PUBLISH

153

154 def getData(self):

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

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

157 self.structure = (

158 ('Topic', ':', MQTT_String),

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

160 ('Message', ':'),

161 )

162 return MQTT_Packet.getData(self)

163

164 class MQTT_Disconnect(MQTT_Packet):

165 structure = (

166 )

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

168 MQTT_Packet.__init__(self, data, alignment)

169 if data is None:

170 self['PacketType'] = PACKET_DISCONNECT

171

172 class MQTT_Subscribe(MQTT_Packet):

173 structure = (

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

175 ('Topic',':', MQTT_String),

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

177 )

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

179 MQTT_Packet.__init__(self, data, alignment)

180 if data is None:

181 self['PacketType'] = PACKET_SUBSCRIBE

182

183 class MQTT_SubscribeACK(MQTT_Packet):

184 structure = (

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

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

187 )

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

189 MQTT_Packet.__init__(self, data, alignment)

190 if data is None:

191 self['PacketType'] = PACKET_SUBSCRIBEACK

192

193 class MQTT_UnSubscribe(MQTT_Packet):

194 structure = (

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

196 ('Topics',':'),

197 )

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

199 MQTT_Packet.__init__(self, data, alignment)

200 if data is None:

201 self['PacketType'] = PACKET_UNSUBSCRIBE

202

203 class MQTTSessionError(Exception):

204 """

205 This is the exception every client should catch

206 """

207

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

209 Exception.__init__(self)

210 self.error = error

211 self.packet = packet

212 self.errorString = errorString

213

214 def getErrorCode(self):

215 return self.error

216

217 def getErrorPacket(self):

218 return self.packet

219

220 def getErrorString(self):

221 return self.errorString

222

223 def __str__(self):

224 return self.errorString

225

226 class MQTTConnection:

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

228 self._targetHost = host

229 self._targetPort = port

230 self._isSSL = isSSL

231 self._socket = None

232 self._messageId = 1

233 self.connectSocket()

234

235 def getSocket(self):

236 return self._socket

237

238 def connectSocket(self):

239 s = socket.socket()

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

241

242 if self._isSSL is True:

243 ctx = SSL.Context(SSL.TLSv1_METHOD)

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

245 self._socket.set_connect_state()

246 self._socket.do_handshake()

247 else:

248 self._socket = s

249

250 def send(self, request):

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

252

253 def sendReceive(self, request):

254 self.send(request)

255 return self.recv()

256

257 def recv(self):

258 REQUEST_SIZE = 8192

259 data = ''

260 done = False

261 while not done:

262 recvData = self._socket.recv(REQUEST_SIZE)

263 if len(recvData) < REQUEST_SIZE:

264 done = True

265 data += recvData

266

267 response = []

268 while len(data) > 0:

269 try:

270 message = MQTT_Packet(data)

271 remaining = data[len(message):]

272 except Exception:

273 # We need more data

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

275 else:

276 response.append(message)

277 data = remaining

278

279 self._messageId += 1

280 return response

281

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

283 """

284

285 :param clientId: Whatever cliend Id that represents you

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

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

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

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

290 :param flags:

291 :param keepAlive: default 60

292 :return: True or MQTTSessionError if something went wrong

293 """

294

295 # Let's build the packet

296 connectPacket = MQTT_Connect()

297 connectPacket['Version'] = version

298 connectPacket['Flags'] = flags

299 connectPacket['KeepAlive'] = keepAlive

300 connectPacket['ProtocolName'] = MQTT_String()

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

302

303 connectPacket['ClientID'] = MQTT_String()

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

305

306 if username is not None:

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

308 if username is None:

309 user = ''

310 else:

311 user = username

312 if password is None:

313 pwd = ''

314 else:

315 pwd = password

316

317 username = MQTT_String()

318 username['Name'] = user

319 password = MQTT_String()

320 password['Name'] = pwd

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

322

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

324

325 response = MQTT_ConnectAck(str(data))

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

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

328

329 return True

330

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

332 """

333

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

335 :param messageID: optional messageId

336 :param flags: Message flags

337 :param QoS: define the QoS requested

338 :return: True or MQTTSessionError if something went wrong

339 """

340 subscribePacket = MQTT_Subscribe()

341 subscribePacket['MessageID'] = messageID

342 subscribePacket['Topic'] = MQTT_String()

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

344 subscribePacket['Flags'] = flags

345 subscribePacket.setQoS(QoS)

346

347 try:

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

349 except Exception as e:

350 raise MQTTSessionError(errorString=str(e))

351

352 subAck = MQTT_SubscribeACK(str(data))

353

354 if subAck['ReturnCode'] > 2:

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

356

357 return True

358

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

360 """

361 Unsubscribes from a topic

362

363 :param topic:

364 :param messageID:

365 :param QoS: define the QoS requested

366 :return:

367 """

368 # ToDo: Support more than one topic

369 packet = MQTT_UnSubscribe()

370 packet['MessageID'] = messageID

371 packet['Topics'] = MQTT_String()

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

373 packet.setQoS( QoS )

374

375 return self.sendReceive(packet)

376

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

378

379 packet = MQTT_Publish()

380 packet['Topic'] = MQTT_String()

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

382 packet['Message'] = message

383 packet['MessageID'] = messageID

384 packet.setQoS( QoS )

385

386 return self.sendReceive(packet)

387

388 def disconnect(self):

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

390

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

392 HOST = '192.168.45.162'

393 USER = 'test'

394 PASS = 'test'

395

396 mqtt = MQTTConnection(HOST, 1883, False)

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

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

399 #mqtt.connect()

400

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

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

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

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

405

406

407 while True:

408

409 packets = mqtt.recv()

410 for packet in packets:

411 publish = MQTT_Publish(str(packet))

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

413

414 mqtt.disconnect()