MQTTSubscribe now supports multiple topic subscriptions in the payload. (#2759)

Author: luisgar1990

scapy/contrib/mqtt.py

@@ -220,14 +220,26 @@ class MQTTPubcomp(Packet):
     ]
 
 
+class MQTTTopic(Packet):
+    name = "MQTT topic"
+    fields_desc = [
+        FieldLenField("length", None, length_of="topic"),
+        StrLenField("topic", "", length_from=lambda pkt:pkt.length)
+    ]
+
+    def guess_payload_class(self, payload):
+        return conf.padding_layer
+
+
+class MQTTTopicQOS(MQTTTopic):
+    fields_desc = MQTTTopic.fields_desc + [ByteEnumField("QOS", 0, QOS_LEVEL)]
+
+
 class MQTTSubscribe(Packet):
     name = "MQTT subscribe"
     fields_desc = [
         ShortField("msgid", None),
-        FieldLenField("length", None, length_of="topic"),
-        StrLenField("topic", "",
-                    length_from=lambda pkt: pkt.length),
-        ByteEnumField("QOS", 0, QOS_LEVEL),
+        PacketListField("topics", [], cls=MQTTTopicQOS)
     ]
 
 
@@ -247,32 +259,11 @@ class MQTTSuback(Packet):
     ]
 
 
-class MQTTTopic(Packet):
-    name = "MQTT topic"
-    fields_desc = [
-        FieldLenField("len", None, length_of="topic"),
-        StrLenField("topic", "", length_from=lambda pkt:pkt.len)
-    ]
-
-    def guess_payload_class(self, payload):
-        return conf.padding_layer
-
-
-def cb_topic(pkt, lst, cur, remain):
-    """
-    Decode the remaining bytes as a MQTT topic
-    """
-    if len(remain) > 3:
-        return MQTTTopic
-    else:
-        return conf.raw_layer
-
-
 class MQTTUnsubscribe(Packet):
     name = "MQTT unsubscribe"
     fields_desc = [
         ShortField("msgid", None),
-        PacketListField("topics", [], next_cls_cb=cb_topic)
+        PacketListField("topics", [], cls=MQTTTopic)
     ]
 
 

test/contrib/mqtt.uts

@@ -70,20 +70,20 @@ assert(connack.retcode == 0)
 
 
 = MQTTSubscribe, packet instantiation
-sb = MQTT()/MQTTSubscribe(msgid=1,topic='newtopic',QOS=0,length=0)
+sb = MQTT()/MQTTSubscribe(msgid=1, topics=[MQTTTopicQOS(topic='newtopic', QOS=1, length=0)])
 assert(sb.type == 8)
 assert(sb.msgid == 1)
-assert(sb.topic == b'newtopic')
-assert(sb.length == 0)
-assert(sb[MQTTSubscribe].QOS == 0)
+assert(sb.topics[0].topic == b'newtopic')
+assert(sb.topics[0].length == 0)
+assert(sb[MQTTSubscribe][MQTTTopicQOS].QOS == 1)
 
 = MQTTSubscribe, packet dissection
-s = b'\x82\t\x00\x01\x00\x04test\x00'
+s = b'\x82\t\x00\x01\x00\x04test\x01'
 subscribe = MQTT(s)
 assert(subscribe.msgid == 1)
-assert(subscribe.length == 4)
-assert(subscribe.topic == b'test')
-assert(subscribe.QOS == 1)
+assert(subscribe.topics[0].length == 4)
+assert(subscribe.topics[0].topic == b'test')
+assert(subscribe.topics[0].QOS == 1)
 
 
 = MQTTSuback, packet instantiation
@@ -98,6 +98,30 @@ suback = MQTT(s)
 assert(suback.msgid == 1)
 assert(suback.retcode == 0)
 
+= MQTTUnsubscribe, packet instantiation
+unsb = MQTT()/MQTTUnsubscribe(msgid=1, topics=[MQTTTopic(topic='newtopic',length=0)])
+assert(unsb.type == 10)
+assert(unsb.msgid == 1)
+assert(unsb.topics[0].topic == b'newtopic')
+assert(unsb.topics[0].length == 0)
+
+= MQTTUnsubscribe, packet dissection
+u = b'\xA2\x09\x00\x01\x00\x03\x61\x2F\x62'
+unsubscribe = MQTT(u)
+assert(unsubscribe.msgid == 1)
+assert(unsubscribe.topics[0].length == 3)
+assert(unsubscribe.topics[0].topic == b'a/b')
+
+= MQTTUnsuback, packet instantiation
+unsk = MQTT()/MQTTUnsuback(msgid=1)
+assert(unsk.type == 11)
+assert(unsk.msgid == 1)
+
+= MQTTUnsuback, packet dissection
+u = b'\xb0\x02\x00\x01'
+unsuback = MQTT(u)
+assert(unsuback.type == 11)
+assert(unsuback.msgid == 1)
 
 = MQTTPubrec, packet instantiation
 pc = MQTT()/MQTTPubrec(msgid=1)
@@ -129,3 +153,7 @@ assert(type(MQTT().fieldtype['len'].randval() + 0) == int)
 = MQTTUnsubscribe
 u = MQTT(b'\xA2\x0C\x00\x01\x00\x03\x61\x2F\x62\x00\x03\x63\x2F\x64')
 assert MQTTUnsubscribe in u and len(u.topics) == 2 and u.topics[1].topic == b"c/d"
+
+= MQTTSubscribe
+u = MQTT(b'\x82\x10\x00\x01\x00\x03\x61\x2F\x62\x02\x00\x03\x63\x2F\x64\x00')
+assert MQTTSubscribe in u and len(u.topics) == 2 and u.topics[1].topic == b"c/d"