DNS: edit all question records if more than one

src/packet/DNS.py

-import logging
 import random
 import scapy.all as scapy
 from scapy.layers import dns
@@ -38,6 +37,22 @@ class DNS(Packet):
     ]
 
 
+    @staticmethod
+    def iter_question_records(question_records: dns.DNSQRField) -> iter:
+        """
+        Iterate over question records.
+
+        :param question_records: List of question records.
+        :return: Iterator over question records.
+        """
+        layer_idx = 0
+        question_record = question_records.getlayer(layer_idx)
+        while question_record is not None:
+            yield question_record
+            layer_idx += 1
+            question_record = question_records.getlayer(layer_idx)
+
+
     def tweak(self) -> dict:
         """
         Randomly edit one DNS field, among the following:
@@ -48,10 +63,11 @@ class DNS(Packet):
         :return: Dictionary containing tweak information.
         """
         # Get field which will be modified
-        field = random.choice(self.fields)
+        #field = random.choice(self.fields)
+        field = "qname"
         # Get auxiliary fields
         qdcount = self.layer.getfieldval("qdcount")
-        question_record = self.layer.getfieldval("qd") if qdcount > 0 else None
+        question_records = self.layer.getfieldval("qd") if qdcount > 0 else None
 
         # Initialize old and new values
         old_value = None
@@ -65,25 +81,34 @@ class DNS(Packet):
             self.layer.setfieldval("qr", new_value)
         
         # Field is query type
-        elif field == "qtype" and question_record is not None:
-            old_value = question_record.getfieldval("qtype")
+        elif field == "qtype" and question_records is not None:
+            old_value = question_records.getfieldval("qtype")
             # Randomly pick new query type
             new_value = old_value
             while new_value == old_value:
                 new_value = random.choice(self.qtypes)
-            question_record.setfieldval("qtype", new_value)
+            question_records.setfieldval("qtype", new_value)
         
         # Field is query name
-        elif field == "qname" and question_record is not None:
-            old_value = question_record.getfieldval("qname")
-            suffix = old_value[-1]
-            old_value_trimmed = old_value[:-1]
-            # Randomly change one character in query name
-            new_value_trimmed = old_value_trimmed
-            while new_value_trimmed == old_value_trimmed:
-                new_value_trimmed = Packet.bytes_edit_char(old_value_trimmed)
-            new_value = new_value_trimmed + bytes(chr(suffix), "utf-8")
-            question_record.setfieldval("qname", new_value)
+        elif field == "qname" and question_records is not None:
+            old_value = ""
+            new_value = ""
+            for question_record in DNS.iter_question_records(question_records):
+                if old_value != "":
+                    old_value += " + "
+                old_value_single = question_record.getfieldval("qname")
+                old_value += old_value_single.decode("utf-8")
+                suffix = old_value_single[-1]
+                old_value_trimmed = old_value_single[:-1]
+                # Randomly change one character in query name
+                new_value_trimmed = old_value_trimmed
+                while new_value_trimmed == old_value_trimmed:
+                    new_value_trimmed = Packet.bytes_edit_char(old_value_trimmed)
+                new_value_single = new_value_trimmed + bytes(chr(suffix), "utf-8")
+                if new_value != "":
+                    new_value += " + "
+                new_value += new_value_single.decode("utf-8")
+                question_record.setfieldval("qname", new_value_single)
         
         # Update checksums
         self.update_checksums()