scapy - demo of dispatch_hook method to dissect next layer based on payload data instead of lower layer attribute(s)

dispatch_hook_demo.py

"""

    This demo of the dispatch_hook uses a hypothetical protocol running on
    top of Ethernet using Ether-type 0xa5b4:

    Ethernet  [ Foo | Bar ]  Baz

    Foo : Type String
    Bar : Type Short String

    Foo and Bar share the same Ether-type and can only be distinguished by checking the first
    byte of the layer. This is the point where the dispatch_hook catches in...

    e.g.

    >>> frm = Ether()/Foo()/Baz()
    >>> frm.show()
    ###[ Ethernet ]###
      dst= ff:ff:ff:ff:ff:ff
      src= 78:24:af:8a:8a:4e
      type= 0xa5b4
    ###[ Foo ]###
         type= FooLayer
         value= 'foo'
    ###[ Baz ]###
            value= 'baz!'

"""
import struct

from scapy.config import conf
from scapy.layers.l2 import Ether
from scapy.packet import Packet, bind_layers, Raw
from scapy.fields import StrFixedLenField, ByteEnumField, ObservableDict, ShortField

FOO_BAR_BAZ_ETHER_TYPE = 0xa5b4


class FooBarBaz(Packet):
    """
    this class acts as a base for Foo and Bar
    """
    FOO_TYPE_ID = 0x00
    BAR_TYPE_ID = 0x01

    ETHER_TYPE = 0xa5b4

    types = ObservableDict(
        {FOO_TYPE_ID: 'FooLayer',
         BAR_TYPE_ID: 'BarLayer'
         })

    @classmethod
    def dispatch_hook(cls, _pkt=None, *args, **kargs):
        """
        this is called upon dissection and returns the class of the upper layer
        :param _pkt: payload
        :param args:
        :param kargs:
        :return: class representing the layer
        """

        foo_bar_type = struct.unpack("B", _pkt[0])[0]

        try:
            return {
                FooBarBaz.FOO_TYPE_ID: Foo,
                FooBarBaz.BAR_TYPE_ID: Bar
            }[foo_bar_type]
        except KeyError:
            # if there is no matching upper layer -> return raw (could also throw an Exception to signal broken frame)
            return Raw


class Foo(FooBarBaz):
    name = 'Foo'
    fields_desc = [ByteEnumField('type', FooBarBaz.FOO_TYPE_ID, FooBarBaz.types),
                   StrFixedLenField('value', 'foo', length=3)]


class Bar(FooBarBaz):
    name = 'Bar'
    fields_desc = [ByteEnumField('type', FooBarBaz.BAR_TYPE_ID, FooBarBaz.types),
                   ShortField('some_number', 42),
                   StrFixedLenField('value', 'bar', length=3)]


class Baz(FooBarBaz):
    name = 'Baz'
    fields_desc = [StrFixedLenField('value', 'baz!', length=4)]


# only the dispatcher is bound to the lower layer
bind_layers(Ether, FooBarBaz, type=FOO_BAR_BAZ_ETHER_TYPE)
bind_layers(Foo, Baz)
bind_layers(Bar, Baz)

conf.l2types.register(FOO_BAR_BAZ_ETHER_TYPE, Ether)

dispatch_hook_demo.uts

#
# execute test:
# > test/run_tests -P "load_contrib('dispatch_hook_demo')" -t scapy/contrib/dispatch_hook_demo.uts
#

% Test dispatch_hook demo - FooBarBaz protocol


#######################################################################
+ build and dissect FooBarBaz layers
#######################################################################

= foo

frm = Ether()/Foo()/Baz()
frm = Ether(frm.do_build())
foo_layer = frm.getlayer(Foo, nb=1)
assert(foo_layer)
frm.show()

= bar

frm = Ether()/Bar()/Baz()
frm = Ether(frm.do_build())
bar_layer = frm.getlayer(Bar, nb=1)
assert(bar_layer)
frm.show()

= unknown
frm = Ether(type=FOO_BAR_BAZ_ETHER_TYPE)/Raw(b'\x02\x01\x02\x03\x04\x05')
frm = Ether(frm.do_build())
raw_layer = frm.getlayer(Raw, nb=1)
assert(raw_layer)
frm.show()