Dispatch mDNS from DNS to add custom functionality

src/packet/DNS.py

@@ -52,6 +52,15 @@ class DNS(Packet):
             layer_idx += 1
             question_record = question_records.getlayer(layer_idx)
 
+    
+    def get_field(self) -> str:
+        """
+        Randomly pick a DNS field to be modified.
+
+        :return: Field name.
+        """
+        return random.choice(self.fields)
+
 
     def tweak(self) -> dict:
         """
@@ -63,7 +72,8 @@ class DNS(Packet):
         :return: Dictionary containing tweak information.
         """
         # Get field which will be modified
-        field = random.choice(self.fields)
+        field = self.get_field()
+        
         # Get auxiliary fields
         qdcount = self.layer.getfieldval("qdcount")
         question_records = self.layer.getfieldval("qd") if qdcount > 0 else None

src/packet/Packet.py

@@ -116,9 +116,14 @@ class Packet:
             try:
                 protocol = layer.name.replace(" ", "_")
                 if protocol == "IP" and packet.getfieldval("version") == 4:
+                    # IPv4 packet
                     protocol = "IPv4"
                 elif protocol == "IP" and packet.getfieldval("version") == 6:
+                    # IPv6 packet
                     protocol = "IPv6"
+                elif protocol == "DNS" and packet.getfieldval("sport") == 5353 and packet.getfieldval("sport") == 5353:
+                    # mDNS packet
+                    protocol = "mDNS"
                 else:
                     protocol = Packet.protocols.get(protocol, protocol)
                 module = importlib.import_module(f"packet.{protocol}")

src/packet/mDNS.py

+import random
+import scapy.all as scapy
+from scapy.layers import dns
+from packet.DNS import DNS
+
+class mDNS(DNS):
+
+    # Class variables
+    name = "mDNS"
+
+    # Modifiable fields
+    fields = {
+        "query": [
+            "qr",
+            "qtype",
+            "qname"
+        ],
+        "response": [
+            "qr"
+        ]
+    }
+
+    
+    def __init__(self, packet: scapy.Packet, id: int = 0, last_layer_index: int = -1) -> None:
+        """
+        mDNS packet constructor.
+
+        :param packet: Scapy packet to be edited.
+        :param id: Packet integer identifier.
+        :param last_layer_index: [Optional] Index of the last layer of the packet.
+                                 If not specified, it will be calculated.
+        """
+        super().__init__(packet, id, last_layer_index)
+        qr = self.layer.getfieldval("qr")
+        self.qr_str = "query" if qr == 0 else "response"
+
+    
+    def get_field(self) -> str:
+        """
+        Randomly pick a DNS field to be modified.
+
+        :return: Field name.
+        """
+        return random.choice(self.fields[self.qr_str])