Merge pull request #4139 from gpotter2/bitlenfield

Add BitLenEnumField
secdev · Oct 18, 2023 · afe8eec · afe8eec
2 parents 2a4c4ae + 86a953b
commit afe8eec
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 6 deletions.
diff --git a/scapy/fields.py b/scapy/fields.py
@@ -2427,7 +2427,7 @@ class BitField(_BitField[int]):
     __doc__ = _BitField.__doc__
 
 
-class BitFixedLenField(BitField):
+class BitLenField(BitField):
     __slots__ = ["length_from"]
 
     def __init__(self,
@@ -2437,15 +2437,15 @@ def __init__(self,
                  ):
         # type: (...) -> None
         self.length_from = length_from
-        super(BitFixedLenField, self).__init__(name, default, 0)
+        super(BitLenField, self).__init__(name, default, 0)
 
     def getfield(self,  # type: ignore
                  pkt,  # type: Packet
                  s,  # type: Union[Tuple[bytes, int], bytes]
                  ):
         # type: (...) -> Union[Tuple[Tuple[bytes, int], int], Tuple[bytes, int]]  # noqa: E501
         self.size = self.length_from(pkt)
-        return super(BitFixedLenField, self).getfield(pkt, s)
+        return super(BitLenField, self).getfield(pkt, s)
 
     def addfield(self,  # type: ignore
                  pkt,  # type: Packet
@@ -2454,7 +2454,7 @@ def addfield(self,  # type: ignore
                  ):
         # type: (...) -> Union[Tuple[bytes, int, int], bytes]
         self.size = self.length_from(pkt)
-        return super(BitFixedLenField, self).addfield(pkt, s, val)
+        return super(BitLenField, self).addfield(pkt, s, val)
 
 
 class BitFieldLenField(BitField):
@@ -2661,6 +2661,32 @@ def i2repr(self,
         return _EnumField.i2repr(self, pkt, x)
 
 
+class BitLenEnumField(BitLenField, _EnumField[int]):
+    __slots__ = EnumField.__slots__
+
+    def __init__(self,
+                 name,  # type: str
+                 default,  # type: Optional[int]
+                 length_from,  # type: Callable[[Packet], int]
+                 enum,  # type: Dict[int, str]
+                 **kwargs,  # type: Any
+                 ):
+        # type: (...) -> None
+        _EnumField.__init__(self, name, default, enum)
+        BitLenField.__init__(self, name, default, length_from, **kwargs)
+
+    def any2i(self, pkt, x):
+        # type: (Optional[Packet], Any) -> int
+        return _EnumField.any2i(self, pkt, x)  # type: ignore
+
+    def i2repr(self,
+               pkt,  # type: Optional[Packet]
+               x,  # type: Union[List[int], int]
+               ):
+        # type: (...) -> Any
+        return _EnumField.i2repr(self, pkt, x)
+
+
 class ShortEnumField(EnumField[int]):
     __slots__ = EnumField.__slots__
 

diff --git a/scapy/layers/sixlowpan.py b/scapy/layers/sixlowpan.py
@@ -58,7 +58,7 @@
 from scapy.fields import (
     BitEnumField,
     BitField,
-    BitFixedLenField,
+    BitLenField,
     BitScalingField,
     ByteEnumField,
     ByteField,
@@ -282,7 +282,7 @@ class LoWPAN_HC1(Packet):
             lambda pkt: pkt.nh == 1 and pkt.hc2
         ),
         # Out of spec
-        BitFixedLenField("pad", 0, _get_hc1_pad)
+        BitLenField("pad", 0, _get_hc1_pad)
     ]
 
     def post_dissect(self, data):

diff --git a/test/fields.uts b/test/fields.uts
@@ -336,7 +336,26 @@ p = TestBFLenF(b"\xff\x6fabcdeFGH")
 p
 assert p.len == 6 and p.str == b"abcde" and Raw in p and p[Raw].load == b"FGH"
 
+= Test BitLenField
+~ field
+
+SIZES = {0: 6, 1: 6, 2: 14, 3: 22}
+
+class TestBitLenField(Packet):
+    fields_desc = [
+        BitField("mode", 0, 2),
+        BitLenField("value", 0, length_from=lambda pkt: SIZES[pkt.mode])
+    ]
+
+p = TestBitLenField(mode=1, value=50)
+assert bytes(p) == b"r"
+
+p = TestBitLenField(mode=2, value=5000)
+assert bytes(p) == b'\x93\x88'
 
+p = TestBitLenField(b'\xc0\x01\xf4')
+assert p.mode == 3
+assert p.value == 500
 
 ############
 ############