Add new initialization arguments for mask class to support set mask with less code

src/ptf/mask.py

@@ -12,13 +12,28 @@ class MaskException(Exception):
 
 
 class Mask:
-    def __init__(self, exp_pkt, ignore_extra_bytes=False):
+    def __init__(self, exp_pkt, ignore_extra_bytes=False, dont_care_all=False):
         self.exp_pkt = exp_pkt
         self.size = len(exp_pkt)
         self.valid = True
-        self.mask = [0xFF] * self.size
+        self.mask = [0] * self.size if dont_care_all else [0xFF] * self.size
         self.ignore_extra_bytes = ignore_extra_bytes
 
+    def set_care(self, offset, bitwidth):
+        for idx in range(offset, offset + bitwidth):
+            offsetB = idx // 8
+            offsetb = idx % 8
+            self.mask[offsetB] = self.mask[offsetB] | (1 << (7 - offsetb))
+
+    def set_care_all(self):
+        self.mask = [0xFF] * self.size
+
+    def set_care_packet(self, hdr_type, field_name):
+        offset, bitwidth = self._calculate_fields_offset_and_bitwidth(
+            hdr_type, field_name
+        )
+        self.set_care(offset, bitwidth)
+
     def set_do_not_care(self, offset, bitwidth):
         # a very naive but simple method
         # we do it bit by bit :)
@@ -27,7 +42,45 @@ def set_do_not_care(self, offset, bitwidth):
             offsetb = idx % 8
             self.mask[offsetB] = self.mask[offsetB] & (~(1 << (7 - offsetb)))
 
+    def set_do_not_care_all(self):
+        self.mask = [0] * self.size
+
     def set_do_not_care_packet(self, hdr_type, field_name):
+        offset, bitwidth = self._calculate_fields_offset_and_bitwidth(
+            hdr_type, field_name
+        )
+        self.set_do_not_care(offset, bitwidth)
+
+    def set_do_not_care_scapy(self, hdr_type, field_name):
+        warnings.warn(
+            '"set_do_not_care_scapy" is going to be deprecated, please '
+            'switch to the new one: "set_do_not_care_packet"',
+            DeprecationWarning,
+        )
+        self.set_do_not_care_packet(hdr_type, field_name)
+
+    def set_ignore_extra_bytes(self):
+        self.ignore_extra_bytes = True
+
+    def is_valid(self):
+        return self.valid
+
+    def pkt_match(self, pkt):
+        # just to be on the safe side
+        pkt = bytearray(bytes(pkt))
+        # we fail if we don't match on sizes, or if ignore_extra_bytes is set,
+        # fail if we have not received at least size bytes
+        if (not self.ignore_extra_bytes and len(pkt) != self.size) or len(
+            pkt
+        ) < self.size:
+            return False
+        exp_pkt = bytearray(bytes(self.exp_pkt))
+        for i in range(self.size):
+            if (exp_pkt[i] & self.mask[i]) != (pkt[i] & self.mask[i]):
+                return False
+        return True
+
+    def _calculate_fields_offset_and_bitwidth(self, hdr_type, field_name):
         if hdr_type not in self.exp_pkt:
             self.valid = False
             raise MaskException("Unknown header type")
@@ -64,36 +117,7 @@ def set_do_not_care_packet(self, hdr_type, field_name):
                 break
             else:
                 offset += bits
-        self.set_do_not_care(hdr_offset * 8 + offset, bitwidth)
-
-    def set_do_not_care_scapy(self, hdr_type, field_name):
-        warnings.warn(
-            '"set_do_not_care_scapy" is going to be deprecated, please '
-            'switch to the new one: "set_do_not_care_packet"',
-            DeprecationWarning,
-        )
-        self.set_do_not_care_packet(hdr_type, field_name)
-
-    def set_ignore_extra_bytes(self):
-        self.ignore_extra_bytes = True
-
-    def is_valid(self):
-        return self.valid
-
-    def pkt_match(self, pkt):
-        # just to be on the safe side
-        pkt = bytearray(bytes(pkt))
-        # we fail if we don't match on sizes, or if ignore_extra_bytes is set,
-        # fail if we have not received at least size bytes
-        if (not self.ignore_extra_bytes and len(pkt) != self.size) or len(
-            pkt
-        ) < self.size:
-            return False
-        exp_pkt = bytearray(bytes(self.exp_pkt))
-        for i in range(self.size):
-            if (exp_pkt[i] & self.mask[i]) != (pkt[i] & self.mask[i]):
-                return False
-        return True
+        return hdr_offset * 8 + offset, bitwidth
 
     def __str__(self):
         old_stdout = sys.stdout

utests/tests/ptf/test_mask.py

@@ -23,13 +23,46 @@ def test_mask__mask_simple_packet(self, scapy_simple_tcp_packet):
         mask_packet.set_do_not_care_packet(TCP, "chksum")
         assert mask_packet.pkt_match(modified_packet)
 
+    def test_mask__set_do_not_care_all(self):
+        expected_packet = "\x01\x02\x03\x04\x05\x06"
+        packet = "\x08\x07\x06\x05\x04\x03\x02\x01"
+        mask = Mask(expected_packet.encode(), ignore_extra_bytes=True)
+        mask.set_do_not_care_all()
+        assert mask.pkt_match(packet.encode())
+
     def test_mask__set_do_not_care(self):
         expected_packet = "\x01\x02\x03\x04\x05\x06"
         packet = "\x01\x00\x00\x04\x05\x06\x07\x08"
         mask = Mask(expected_packet.encode(), ignore_extra_bytes=True)
         mask.set_do_not_care(8, 16)
         assert mask.pkt_match(packet.encode())
 
+    def test_mask__set_care_all(self):
+        expected_packet = "\x01\x02\x03\x04\x05\x06"
+        packet = "\x00\x02\x03\x04\x05\x06"
+        mask = Mask(expected_packet.encode(), ignore_extra_bytes=True)
+        mask.set_care_all()
+        assert not mask.pkt_match(packet.encode())
+
+    def test_mask__set_care(self):
+        expected_packet = "\x01\x02\x03\x04\x05\x06"
+        packet = "\x01\x02\x00\x04\x05\x06\x07\x08"
+        mask = Mask(
+            expected_packet.encode(), ignore_extra_bytes=True, dont_care_all=True
+        )
+        mask.set_care(0, 16)
+        assert mask.pkt_match(packet.encode())
+        mask.set_care(16, 8)
+        assert not mask.pkt_match(packet.encode())
+
+    def test_mask__set_care_packet(self):
+        packet = IP(src="1.1.1.1")
+        mask = Mask(packet.copy(), ignore_extra_bytes=True, dont_care_all=True)
+        packet[IP].src = "2.2.2.2"
+        assert mask.pkt_match(packet)
+        mask.set_care_packet(IP, "src")
+        assert not mask.pkt_match(packet)
+
     def test_mask__check_masking_conditional_field(self, scapy_simple_vxlan_packet):
         simple_vxlan = scapy_simple_vxlan_packet
         simple_vxlan[VXLAN].flags = "G"