diff --git a/scapy/fields.py b/scapy/fields.py
index 0dbe6910eba412c44c8f5ead1712f173fc499e3b..53f84727ca737ee33fc52d11438303d2aefc7554 100644
--- a/scapy/fields.py
+++ b/scapy/fields.py
@@ -798,32 +798,68 @@ class XBitField(BitField):
class _EnumField(Field):
def __init__(self, name, default, enum, fmt = "H"):
- i2s = self.i2s = {}
- s2i = self.s2i = {}
- if type(enum) is list:
- keys = range(len(enum))
+ """ Initializes enum fields.
+
+ @param name: name of this field
+ @param default: default value of this field
+ @param enum: either a dict or a tuple of two callables. Dict keys are
+ the internal values, while the dict values are the
+ user-friendly representations. If the tuple is provided,
+ the first callable receives the internal value as
+ parameter and returns the user-friendly representation
+ and the second callable does the converse. The first
+ callable may return None to default to a literal string
+ (repr()) representation.
+ @param fmt: struct.pack format used to parse and serialize the
+ internal value from and to machine representation.
+ """
+ if isinstance(enum, tuple):
+ self.i2s_cb = enum[0]
+ self.s2i_cb = enum[1]
+ self.i2s = None
+ self.s2i = None
else:
- keys = enum.keys()
- if any(type(x) is str for x in keys):
- i2s, s2i = s2i, i2s
- for k in keys:
- i2s[k] = enum[k]
- s2i[enum[k]] = k
+ i2s = self.i2s = {}
+ s2i = self.s2i = {}
+ self.i2s_cb = None
+ self.s2i_cb = None
+ if type(enum) is list:
+ keys = range(len(enum))
+ else:
+ keys = enum.keys()
+ if any(type(x) is str for x in keys):
+ i2s, s2i = s2i, i2s
+ for k in keys:
+ i2s[k] = enum[k]
+ s2i[enum[k]] = k
Field.__init__(self, name, default, fmt)
+
def any2i_one(self, pkt, x):
if type(x) is str:
- x = self.s2i[x]
+ try:
+ x = self.s2i[x]
+ except TypeError:
+ x = self.s2i_cb(x)
return x
+
def i2repr_one(self, pkt, x):
- if self not in conf.noenum and not isinstance(x,VolatileValue) and x in self.i2s:
- return self.i2s[x]
+ if self not in conf.noenum and not isinstance(x,VolatileValue):
+ try:
+ return self.i2s[x]
+ except KeyError:
+ pass
+ except TypeError:
+ ret = self.i2s_cb(x)
+ if ret is not None:
+ return ret
return repr(x)
def any2i(self, pkt, x):
if type(x) is list:
return map(lambda z,pkt=pkt:self.any2i_one(pkt,z), x)
else:
- return self.any2i_one(pkt,x)
+ return self.any2i_one(pkt,x)
+
def i2repr(self, pkt, x):
if type(x) is list:
return map(lambda z,pkt=pkt:self.i2repr_one(pkt,z), x)
@@ -831,17 +867,21 @@ class _EnumField(Field):
return self.i2repr_one(pkt,x)
class EnumField(_EnumField):
- __slots__ = ["i2s", "s2i"]
+ __slots__ = ["i2s", "s2i", "s2i_cb", "i2s_cb"]
class CharEnumField(EnumField):
def __init__(self, name, default, enum, fmt = "1s"):
EnumField.__init__(self, name, default, enum, fmt)
- k = self.i2s.keys()
- if k and len(k[0]) != 1:
- self.i2s,self.s2i = self.s2i,self.i2s
+ if self.i2s is not None:
+ k = self.i2s.keys()
+ if k and len(k[0]) != 1:
+ self.i2s,self.s2i = self.s2i,self.i2s
def any2i_one(self, pkt, x):
if len(x) != 1:
- x = self.s2i[x]
+ if self.s2i is None:
+ x = self.s2i_cb(x)
+ else:
+ x = self.s2i[x]
return x
class BitEnumField(BitField, _EnumField):
@@ -856,6 +896,7 @@ class BitEnumField(BitField, _EnumField):
return _EnumField.i2repr(self, pkt, x)
class ShortEnumField(EnumField):
+ __slots__ = EnumField.__slots__
def __init__(self, name, default, enum):
EnumField.__init__(self, name, default, enum, "H")
@@ -883,10 +924,18 @@ class LEIntEnumField(EnumField):
class XShortEnumField(ShortEnumField):
def i2repr_one(self, pkt, x):
- if self not in conf.noenum and not isinstance(x,VolatileValue) and x in self.i2s:
- return self.i2s[x]
+ if self not in conf.noenum and not isinstance(x,VolatileValue):
+ try:
+ return self.i2s[x]
+ except KeyError:
+ pass
+ except TypeError:
+ ret = self.i2s_cb(x)
+ if ret is not None:
+ return ret
return lhex(x)
+
class _MultiEnumField(_EnumField):
def __init__(self, name, default, enum, depends_on, fmt = "H"):
diff --git a/test/fields.uts b/test/fields.uts
index 653bf60d3e00a6c62c040a88be9c3833e6f07fb3..2c683f85cad881cbdd10dfd44b371ade811ccb2c 100644
--- a/test/fields.uts
+++ b/test/fields.uts
@@ -507,4 +507,265 @@ 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)
+############
+############
++ EnumField tests
+
+= EnumField tests initialization
+
+# Basic EnumField
+f = EnumField('test', 0, {0: 'Foo', 1: 'Bar'})
+# Reverse i2s/s2i
+rf = EnumField('test', 0, {'Foo': 0, 'Bar': 1})
+# EnumField initialized with a list
+lf = EnumField('test', 0, ['Foo', 'Bar'])
+# EnumField with i2s_cb/s2i_cb
+fcb = EnumField('test', 0, (
+ lambda x: 'Foo' if x == 0 else 'Bar' if 1 <= x <= 10 else repr(x),
+ lambda x: 0 if x == 'Foo' else 1 if x == 'Bar' else int(x),
+ )
+)
+
+def expect_exception(e, c):
+ try:
+ eval(c)
+ return False
+ except e:
+ return True
+
+
+= EnumField.any2i_one
+~ field enumfield
+
+assert(f.any2i_one(None, 'Foo') == 0)
+assert(f.any2i_one(None, 'Bar') == 1)
+assert(f.any2i_one(None, 2) == 2)
+expect_exception(KeyError, 'f.any2i_one(None, "Baz")')
+
+assert(rf.any2i_one(None, 'Foo') == 0)
+assert(rf.any2i_one(None, 'Bar') == 1)
+assert(rf.any2i_one(None, 2) == 2)
+expect_exception(KeyError, 'rf.any2i_one(None, "Baz")')
+
+assert(lf.any2i_one(None, 'Foo') == 0)
+assert(lf.any2i_one(None, 'Bar') == 1)
+assert(lf.any2i_one(None, 2) == 2)
+expect_exception(KeyError, 'lf.any2i_one(None, "Baz")')
+
+assert(fcb.any2i_one(None, 'Foo') == 0)
+assert(fcb.any2i_one(None, 'Bar') == 1)
+assert(fcb.any2i_one(None, 5) == 5)
+expect_exception(ValueError, 'fcb.any2i_one(None, "Baz")')
+
+True
+
+= EnumField.any2i
+~ field enumfield
+
+assert(f.any2i(None, 'Foo') == 0)
+assert(f.any2i(None, 'Bar') == 1)
+assert(f.any2i(None, 2) == 2)
+expect_exception(KeyError, 'f.any2i(None, "Baz")')
+assert(f.any2i(None, ['Foo', 'Bar', 2]) == [0, 1, 2])
+
+assert(rf.any2i(None, 'Foo') == 0)
+assert(rf.any2i(None, 'Bar') == 1)
+assert(rf.any2i(None, 2) == 2)
+expect_exception(KeyError, 'rf.any2i(None, "Baz")')
+assert(rf.any2i(None, ['Foo', 'Bar', 2]) == [0, 1, 2])
+
+assert(lf.any2i(None, 'Foo') == 0)
+assert(lf.any2i(None, 'Bar') == 1)
+assert(lf.any2i(None, 2) == 2)
+expect_exception(KeyError, 'lf.any2i(None, "Baz")')
+assert(lf.any2i(None, ['Foo', 'Bar', 2]) == [0, 1, 2])
+
+assert(fcb.any2i(None, 'Foo') == 0)
+assert(fcb.any2i(None, 'Bar') == 1)
+assert(fcb.any2i(None, 5) == 5)
+expect_exception(ValueError, 'fcb.any2i(None, "Baz")')
+assert(f.any2i(None, ['Foo', 'Bar', 5]) == [0, 1, 5])
+
+True
+
+= EnumField.i2repr_one
+~ field enumfield
+
+assert(f.i2repr_one(None, 0) == 'Foo')
+assert(f.i2repr_one(None, 1) == 'Bar')
+expect_exception(KeyError, 'f.i2repr_one(None, 2)')
+
+assert(rf.i2repr_one(None, 0) == 'Foo')
+assert(rf.i2repr_one(None, 1) == 'Bar')
+expect_exception(KeyError, 'rf.i2repr_one(None, 2)')
+
+assert(lf.i2repr_one(None, 0) == 'Foo')
+assert(lf.i2repr_one(None, 1) == 'Bar')
+expect_exception(KeyError, 'lf.i2repr_one(None, 2)')
+
+assert(fcb.i2repr_one(None, 0) == 'Foo')
+assert(fcb.i2repr_one(None, 1) == 'Bar')
+assert(fcb.i2repr_one(None, 5) == 'Bar')
+assert(fcb.i2repr_one(None, 11) == repr(11))
+
+conf.noenum.add(f, rf, lf, fcb)
+
+assert(f.i2repr_one(None, 0) == repr(0))
+assert(f.i2repr_one(None, 1) == repr(1))
+assert(f.i2repr_one(None, 2) == repr(2))
+
+assert(rf.i2repr_one(None, 0) == repr(0))
+assert(rf.i2repr_one(None, 1) == repr(1))
+assert(rf.i2repr_one(None, 2) == repr(2))
+
+assert(lf.i2repr_one(None, 0) == repr(0))
+assert(lf.i2repr_one(None, 1) == repr(1))
+assert(lf.i2repr_one(None, 2) == repr(2))
+
+assert(fcb.i2repr_one(None, 0) == repr(0))
+assert(fcb.i2repr_one(None, 1) == repr(1))
+assert(fcb.i2repr_one(None, 5) == repr(5))
+assert(fcb.i2repr_one(None, 11) == repr(11))
+
+conf.noenum.remove(f, rf, lf, fcb)
+
+assert(f.i2repr_one(None, RandNum(0, 10)) == '<RandNum>')
+assert(rf.i2repr_one(None, RandNum(0, 10)) == '<RandNum>')
+assert(lf.i2repr_one(None, RandNum(0, 10)) == '<RandNum>')
+assert(fcb.i2repr_one(None, RandNum(0, 10)) == '<RandNum>')
+
+True
+
+= EnumField.i2repr
+~ field enumfield
+
+assert(f.i2repr(None, 0) == 'Foo')
+assert(f.i2repr(None, 1) == 'Bar')
+expect_exception(KeyError, 'f.i2repr(None, 2)')
+assert(f.i2repr(None, [0, 1]) == ['Foo', 'Bar'])
+
+assert(rf.i2repr(None, 0) == 'Foo')
+assert(rf.i2repr(None, 1) == 'Bar')
+expect_exception(KeyError, 'rf.i2repr(None, 2)')
+assert(rf.i2repr(None, [0, 1]) == ['Foo', 'Bar'])
+
+assert(lf.i2repr(None, 0) == 'Foo')
+assert(lf.i2repr(None, 1) == 'Bar')
+expect_exception(KeyError, 'lf.i2repr(None, 2)')
+assert(lf.i2repr(None, [0, 1]) == ['Foo', 'Bar'])
+
+assert(fcb.i2repr(None, 0) == 'Foo')
+assert(fcb.i2repr(None, 1) == 'Bar')
+assert(fcb.i2repr(None, 5) == 'Bar')
+assert(fcb.i2repr(None, 11) == repr(11))
+assert(fcb.i2repr(None, [0, 1, 5, 11]) == ['Foo', 'Bar', 'Bar', repr(11)])
+
+conf.noenum.add(f, rf, lf, fcb)
+
+assert(f.i2repr(None, 0) == repr(0))
+assert(f.i2repr(None, 1) == repr(1))
+assert(f.i2repr(None, 2) == repr(2))
+assert(f.i2repr(None, [0, 1, 2]) == [repr(0), repr(1), repr(2)])
+
+assert(rf.i2repr(None, 0) == repr(0))
+assert(rf.i2repr(None, 1) == repr(1))
+assert(rf.i2repr(None, 2) == repr(2))
+assert(rf.i2repr(None, [0, 1, 2]) == [repr(0), repr(1), repr(2)])
+
+assert(lf.i2repr(None, 0) == repr(0))
+assert(lf.i2repr(None, 1) == repr(1))
+assert(lf.i2repr(None, 2) == repr(2))
+assert(lf.i2repr(None, [0, 1, 2]) == [repr(0), repr(1), repr(2)])
+
+assert(fcb.i2repr(None, 0) == repr(0))
+assert(fcb.i2repr(None, 1) == repr(1))
+assert(fcb.i2repr(None, 5) == repr(5))
+assert(fcb.i2repr(None, 11) == repr(11))
+assert(fcb.i2repr(None, [0, 1, 5, 11]) == [repr(0), repr(1), repr(5), repr(11)])
+
+conf.noenum.remove(f, rf, lf, fcb)
+
+assert(f.i2repr_one(None, RandNum(0, 10)) == '<RandNum>')
+assert(rf.i2repr_one(None, RandNum(0, 10)) == '<RandNum>')
+assert(lf.i2repr_one(None, RandNum(0, 10)) == '<RandNum>')
+assert(fcb.i2repr_one(None, RandNum(0, 10)) == '<RandNum>')
+
+True
+
+############
+############
++ CharEnumField tests
+
+= Building expect_exception handler
+~ field charenumfield
+
+def expect_exception(e, c):
+ try:
+ eval(c)
+ return False
+ except e:
+ return True
+
+
+= CharEnumField tests initialization
+~ field charenumfield
+
+fc = CharEnumField('test', 'f', {'f': 'Foo', 'b': 'Bar'})
+fcb = CharEnumField('test', 'a', (
+ lambda x: 'Foo' if x == 'a' else 'Bar' if x == 'b' else 'Baz',
+ lambda x: 'a' if x == 'Foo' else 'b' if x == 'Bar' else ''
+))
+
+True
+
+= CharEnumField.any2i_one
+~ field charenumfield
+
+assert(fc.any2i_one(None, 'Foo') == 'f')
+assert(fc.any2i_one(None, 'Bar') == 'b')
+expect_exception(KeyError, 'fc.any2i_one(None, "Baz")')
+
+assert(fcb.any2i_one(None, 'Foo') == 'a')
+assert(fcb.any2i_one(None, 'Bar') == 'b')
+assert(fcb.any2i_one(None, 'Baz') == '')
+
+True
+
+############
+############
++ XShortEnumField tests
+
+= Building expect_exception handler
+~ field xshortenumfield
+
+def expect_exception(e, c):
+ try:
+ eval(c)
+ return False
+ except e:
+ return True
+
+
+= XShortEnumField tests initialization
+~ field xshortenumfield
+
+f = XShortEnumField('test', 0, {0: 'Foo', 1: 'Bar'})
+fcb = XShortEnumField('test', 0, (
+ lambda x: 'Foo' if x == 0 else 'Bar' if x == 1 else lhex(x),
+ lambda x: x
+))
+
+True
+
+= XShortEnumField.i2repr_one
+~ field xshortenumfield
+
+assert(f.i2repr_one(None, 0) == 'Foo')
+assert(f.i2repr_one(None, 1) == 'Bar')
+assert(f.i2repr_one(None, 0xff) == '0xff')
+
+assert(f.i2repr_one(None, 0) == 'Foo')
+assert(f.i2repr_one(None, 1) == 'Bar')
+assert(f.i2repr_one(None, 0xff) == '0xff')
+True