Revert part of 2629 / FlagsField & SMB2 fixes

Author: gpotter2

scapy/fields.py

@@ -2152,27 +2152,54 @@ class FlagsField(BitField):
 
    Make sure all your flags have a label
 
-   Example:
+   Example (list):
        >>> from scapy.packet import Packet
        >>> class FlagsTest(Packet):
                fields_desc = [FlagsField("flags", 0, 8, ["f0", "f1", "f2", "f3", "f4", "f5", "f6", "f7"])]  # noqa: E501
        >>> FlagsTest(flags=9).show2()
        ###[ FlagsTest ]###
          flags     = f0+f3
-       >>> FlagsTest(flags=0).show2().strip()
+
+    Example (str):
+       >>> from scapy.packet import Packet
+       >>> class TCPTest(Packet):
+               fields_desc = [
+                   BitField("reserved", 0, 7),
+                   FlagsField("flags", 0x2, 9, "FSRPAUECN")
+               ]
+       >>> TCPTest(flags=3).show2()
        ###[ FlagsTest ]###
-         flags     =
+         reserved  = 0
+         flags     = FS
+
+    Example (dict):
+       >>> from scapy.packet import Packet
+       >>> class FlagsTest2(Packet):
+               fields_desc = [
+                   FlagsField("flags", 0x2, 16, {
+                       1: "1",  # 1st bit
+                       8: "2"   # 8th bit
+                   })
+               ]
 
    :param name: field's name
    :param default: default value for the field
    :param size: number of bits in the field (in bits)
-   :param names: (list or dict) label for each flag, Least Significant Bit tag's name is written first  # noqa: E501
+   :param names: (list or str or dict) label for each flag
+       If it's a str or a list, the least Significant Bit tag's name
+       is written first.
    """
     ismutable = True
-    __slots__ = ["multi", "names"]
+    __slots__ = ["names"]
 
     def __init__(self, name, default, size, names):
-        self.multi = isinstance(names, list)
+        # Convert the dict to a list
+        if isinstance(names, dict):
+            tmp = ["bit_%d" % i for i in range(size)]
+            for i, v in six.viewitems(names):
+                tmp[i] = v
+            names = tmp
+        # Store the names as str or list
         self.names = names
         BitField.__init__(self, name, default, size)
 

scapy/layers/smb2.py

@@ -9,11 +9,30 @@
 
 from scapy.config import conf
 from scapy.packet import Packet, bind_layers
-from scapy.fields import StrFixedLenField, LEIntField, LEShortEnumField, \
-    ShortEnumField, XLEIntField, LEShortField, FlagsField, LELongField, \
-    XLELongField, XNBytesField, FieldLenField, IntField, FieldListField, \
-    XStrLenField, ShortField, IntEnumField, StrFieldUtf16, XLEShortField, \
-    UUIDField, XLongField, PacketListField, PadField
+from scapy.fields import (
+    FieldLenField,
+    FieldListField,
+    FlagsField,
+    IntEnumField,
+    IntField,
+    LEIntField,
+    LELongField,
+    LEShortEnumField,
+    LEShortField,
+    PacketListField,
+    PadField,
+    ShortEnumField,
+    ShortField,
+    StrFieldUtf16,
+    StrFixedLenField,
+    UUIDField,
+    XLEIntField,
+    XLELongField,
+    XLEShortField,
+    XLongField,
+    XNBytesField,
+    XStrLenField,
+)
 
 
 # EnumField
@@ -35,13 +54,13 @@
 
 # FlagField
 SMB2_CAPABILITIES = {
-    30: "CapabilitiesEncryption",
-    29: "CapabilitiesDirectoryLeasing",
-    28: "CapabilitiesPersistentHandles",
-    27: "CapabilitiesMultiChannel",
-    26: "CapabilitiesLargeMTU",
-    25: "CapabilitiesLeasing",
-    24: "CapabilitiesDFS",
+    30: "Encryption",
+    29: "DirectoryLeasing",
+    28: "PersistentHandles",
+    27: "MultiChannel",
+    26: "LargeMTU",
+    25: "Leasing",
+    24: "DFS",
 }
 
 # EnumField
@@ -76,6 +95,13 @@ class SMB2_Header(Packet):
         XNBytesField("Signature", 0, 16),
     ]
 
+    def guess_payload_class(self, payload):
+        if self.Command == 0x0000:
+            if self.Flags.SMB2_FLAGS_SERVER_TO_REDIR:
+                return SMB2_Negociate_Protocol_Response_Header
+            return SMB2_Negociate_Protocol_Request_Header
+        return super(SMB2_Header, self).guess_payload_class(payload)
+
 
 class SMB2_Compression_Transform_Header(Packet):
     name = "SMB2 Compression Transform Header"
@@ -270,18 +296,6 @@ class SMB2_Negociate_Protocol_Response_Header(Packet):
 bind_layers(SMB2_Encryption_Capabilities, conf.padding_layer)
 bind_layers(SMB2_Compression_Capabilities, conf.padding_layer)
 bind_layers(SMB2_Netname_Negociate_Context_ID, conf.padding_layer)
-bind_layers(
-    SMB2_Header,
-    SMB2_Negociate_Protocol_Request_Header,
-    Command=0x0000,
-    Flags=lambda f: (f >> 24) & 1 == 0
-)
-bind_layers(
-    SMB2_Header,
-    SMB2_Negociate_Protocol_Response_Header,
-    Command=0x0000,
-    Flags=lambda f: (f >> 24) & 1 == 1
-)
 bind_layers(
     SMB2_Negociate_Context,
     SMB2_Preauth_Integrity_Capabilities,

test/fields.uts

@@ -1487,6 +1487,35 @@ assert "f2" in flags
 assert "f4" in flags
 assert "f0" in flags
 
+= FlagsField with str
+
+class TCPTest(Packet):
+    fields_desc = [
+        BitField("reserved", 0, 7),
+        FlagsField("flags", 0x2, 9, "FSRPAUECN")
+    ]
+
+a = TCPTest(flags=3)
+assert a.flags.F
+assert a.flags.S
+assert a.sprintf("%flags%") == "FS"
+
+= FlagsField with dict
+
+class FlagsTest2(Packet):
+    fields_desc = [
+        FlagsField("flags", 0x2, 16, {
+            0: "A",  # 0 bit
+            7: "B"   # 7 bit
+        })
+    ]
+
+a = FlagsTest2(flags=255)
+a.sprintf("%flags%")
+assert a.flags.A
+assert a.flags.B
+assert a.sprintf("%flags%") == "A+bit_1+bit_2+bit_3+bit_4+bit_5+bit_6+B"
+
 
 ########
 ########

test/smb2.uts

@@ -1,7 +1,3 @@
-############
-############
-~ SMB2
-
 + SMB2 Header
 
 = SMB2 Header dissecting
@@ -59,15 +55,6 @@ assert pkt[NBTSession].TYPE == 0x00         # session message
 smb2 = pkt[SMB2_Header]
 assert smb2.Start == b'\xfeSMB'
 
-
-
-
-
-
-
-
-
-
 + SMB2 Negociate Procotol Request Header dissecting
 
 = Common fields in header
@@ -146,28 +133,6 @@ assert netname.DataLength == 28
 assert netname.NetName == '192.168.178.21'
 
 
-
-
-
-
-
-
-
-
-+ test SMB2 Negociate Protocol Request Header - assembling
-
-pkt = IP() / TCP() / NBTSession() / SMB2_Header() / SMB2_Negociate_Protocol_Request_Header()
-assert SMB2_Negociate_Protocol_Request_Header in pkt
-
-
-
-
-
-
-
-
-
-
 + SMB2 Negociate Protocol Response Header dissecting
 
 = Common fields in header
@@ -229,16 +194,3 @@ assert comp.CompressionAlgorithmCount == 1
 assert len(comp.CompressionAlgorithms) == 1
 assert comp.CompressionAlgorithms[0] == 1
 
-
-
-
-
-
-
-
-
-
-+ SMB2 Negociate Protocol Response Header assembling
-
-pkt = IP() / TCP() / NBTSession() / SMB2_Header() / SMB2_Negociate_Protocol_Response_Header()
-assert SMB2_Negociate_Protocol_Response_Header in pkt