Add MultiFlagsField

scapy/fields.py

@@ -7,7 +7,7 @@
 Fields: basic data structures that make up parts of packets.
 """
 
-import struct,copy,socket
+import struct,copy,socket,collections
 from scapy.config import conf
 from scapy.volatile import *
 from scapy.data import *
@@ -990,7 +990,94 @@ class FlagsField(BitField):
             r = "+".join(r)
         return r
 
-            
+
+MultiFlagsEntry = collections.namedtuple('MultiFlagEntry', ['short', 'long'])
+
+
+class MultiFlagsField(BitField):
+    __slots__ = FlagsField.__slots__ + ["depends_on"]
+
+    def __init__(self, name, default, size, names, depends_on):
+        self.names = names
+        self.depends_on = depends_on
+        super(MultiFlagsField, self).__init__(name, default, size)
+
+    def any2i(self, pkt, x):
+        assert isinstance(x, (int, long, set)), 'set expected'
+
+        if pkt is not None:
+            if isinstance(x, (int, long)):
+                x = self.m2i(pkt, x)
+            else:
+                v = self.depends_on(pkt)
+                if v is not None:
+                    assert self.names.has_key(v), 'invalid dependency'
+                    these_names = self.names[v]
+                    s = set()
+                    for i in x:
+                        for j in these_names.keys():
+                            if these_names[j].short == i:
+                                s.add(i)
+                                break
+                        else:
+                            assert False, 'Unknown flag "{}" with this dependency'.format(i)
+                            continue
+                    x = s
+        return x
+
+    def i2m(self, pkt, x):
+        v = self.depends_on(pkt)
+        if v in self.names:
+            these_names = self.names[v]
+        else:
+            these_names = {}
+
+        r = 0
+        for flag_set in x:
+            for i in these_names.keys():
+                if these_names[i].short == flag_set:
+                    r |= 1 << i
+                    break
+            else:
+                r |= 1 << int(flag_set[len('bit '):])
+        return r
+
+    def m2i(self, pkt, x):
+        v = self.depends_on(pkt)
+        if v in self.names:
+            these_names = self.names[v]
+        else:
+            these_names = {}
+
+        r = set()
+        i = 0
+
+        while x:
+            if x & 1:
+                if i in these_names:
+                    r.add(these_names[i].short)
+                else:
+                    r.add('bit {}'.format(i))
+            x >>= 1
+            i += 1
+        return r
+
+    def i2repr(self, pkt, x):
+        v = self.depends_on(pkt)
+        if self.names.has_key(v):
+            these_names = self.names[v]
+        else:
+            these_names = {}
+
+        r = set()
+        for flag_set in x:
+            for i in these_names.itervalues():
+                if i.short == flag_set:
+                    r.add("{} ({})".format(i.long, i.short))
+                    break
+            else:
+                r.add(flag_set)
+        return repr(r)
 
 
 class FixedPointField(BitField):

test/regression.uts

@@ -6304,3 +6304,188 @@ except:
 assert(ret)
 p = ss.recv()
 assert(p.data == 3)
+
++ Tests on MultiFlagsField
+
+= Test calls on MultiFlagsField.any2i
+~ multiflagsfield
+
+import collections
+MockPacket = collections.namedtuple('MockPacket', ['type'])
+
+f = MultiFlagsField('flags', set(), 3, {
+        0: {
+            0: MultiFlagsEntry('A', 'OptionA'),
+            1: MultiFlagsEntry('B', 'OptionB'),
+        },
+        1: {
+            0: MultiFlagsEntry('+', 'Plus'),
+            1: MultiFlagsEntry('*', 'Star'),
+        },
+    },
+    depends_on=lambda x: x.type
+)
+
+mp = MockPacket(0)
+x = f.any2i(mp, set())
+assert(isinstance(x, set))
+assert(len(x) == 0)
+x = f.any2i(mp, {'A'})
+assert(isinstance(x, set))
+assert(len(x) == 1)
+assert('A' in x)
+assert('B' not in x)
+assert('+' not in x)
+x = f.any2i(mp, {'A', 'B'})
+assert(isinstance(x, set))
+assert(len(x) == 2)
+assert('A' in x)
+assert('B' in x)
+assert('+' not in x)
+assert('*' not in x)
+x = f.any2i(mp, 3)
+assert(isinstance(x, set))
+assert(len(x) == 2)
+assert('A' in x)
+assert('B' in x)
+assert('+' not in x)
+assert('*' not in x)
+x = f.any2i(mp, 7)
+assert(isinstance(x, set))
+assert(len(x) == 3)
+assert('A' in x)
+assert('B' in x)
+assert('bit 2' in x)
+assert('+' not in x)
+assert('*' not in x)
+mp = MockPacket(1)
+x = f.any2i(mp, {'+', '*'})
+assert(isinstance(x, set))
+assert(len(x) == 2)
+assert('+' in x)
+assert('*' in x)
+assert('A' not in x)
+assert('B' not in x)
+try:
+    x = f.any2i(mp, {'A'})
+    ret = False
+except AssertionError:
+    ret = True
+
+assert(ret)
+#Following test demonstrate a non-sensical yet acceptable usage :(
+x = f.any2i(None, {'Toto'})
+assert('Toto' in x)
+
+= Test calls on MultiFlagsField.i2m
+~ multiflagsfield
+
+import collections
+MockPacket = collections.namedtuple('MockPacket', ['type'])
+
+f = MultiFlagsField('flags', set(), 3, {
+        0: {
+            0: MultiFlagsEntry('A', 'OptionA'),
+            1: MultiFlagsEntry('B', 'OptionB'),
+        },
+        1: {
+            0: MultiFlagsEntry('+', 'Plus'),
+            1: MultiFlagsEntry('*', 'Star'),
+        },
+    },
+    depends_on=lambda x: x.type
+)
+
+mp = MockPacket(0)
+x = f.i2m(mp, set())
+assert(isinstance(x, (int, long)))
+assert(x == 0)
+x = f.i2m(mp, {'A'})
+assert(isinstance(x, (int, long)))
+assert(x == 1)
+x = f.i2m(mp, {'A', 'B'})
+assert(isinstance(x, (int, long)))
+assert(x == 3)
+x = f.i2m(mp, {'A', 'B', 'bit 2'})
+assert(isinstance(x, (int, long)))
+assert(x == 7)
+try:
+    x = f.i2m(mp, {'+'})
+    ret = False
+except:
+    ret = True
+
+assert(ret)
+
+= Test calls on MultiFlagsField.m2i
+~ multiflagsfield
+
+import collections
+MockPacket = collections.namedtuple('MockPacket', ['type'])
+
+f = MultiFlagsField('flags', set(), 3, {
+        0: {
+            0: MultiFlagsEntry('A', 'OptionA'),
+            1: MultiFlagsEntry('B', 'OptionB'),
+        },
+        1: {
+            0: MultiFlagsEntry('+', 'Plus'),
+            1: MultiFlagsEntry('*', 'Star'),
+        },
+    },
+    depends_on=lambda x: x.type
+)
+
+mp = MockPacket(0)
+x = f.m2i(mp, 2)
+assert(isinstance(x, set))
+assert(len(x) == 1)
+assert('B' in x)
+assert('A' not in x)
+assert('*' not in x)
+
+x = f.m2i(mp, 7)
+assert(isinstance(x, set))
+assert('B' in x)
+assert('A' in x)
+assert('bit 2' in x)
+assert('*' not in x)
+assert('+' not in x)
+x = f.m2i(mp, 0)
+assert(len(x) == 0)
+mp = MockPacket(1)
+x = f.m2i(mp, 2)
+assert(isinstance(x, set))
+assert(len(x) == 1)
+assert('*' in x)
+assert('+' not in x)
+assert('B' not in x)
+
+= Test calls on MultiFlagsField.i2repr
+~ multiflagsfield
+
+import collections, re
+MockPacket = collections.namedtuple('MockPacket', ['type'])
+
+f = MultiFlagsField('flags', set(), 3, {
+        0: {
+            0: MultiFlagsEntry('A', 'OptionA'),
+            1: MultiFlagsEntry('B', 'OptionB'),
+        },
+        1: {
+            0: MultiFlagsEntry('+', 'Plus'),
+            1: MultiFlagsEntry('*', 'Star'),
+        },
+    },
+    depends_on=lambda x: x.type
+)
+
+mp = MockPacket(0)
+x = f.i2repr(mp, {'A', 'B'})
+assert(re.match(r'^.*OptionA \(A\).*$', x) is not None)
+assert(re.match(r'^.*OptionB \(B\).*$', x) is not None)
+mp = MockPacket(1)
+x = f.i2repr(mp, {'*', '+', 'bit 2'})
+assert(re.match(r'^.*Star \(\*\).*$', x) is not None)
+assert(re.match(r'^.*Plus \(\+\).*$', x) is not None)
+assert(re.match(r'^.*bit 2.*$', x) is not None)