Add duplicate option for punt-policies, and introduce manipulator actions

Combination of two things:
1. `duplicate` option when setting the manipulator and punt-policies
2. manipulator actions - drop, inject raw or handle encapsultion

Also added tests for manipulator actions,
and clearify set_manipulation's cb argument

Author: roykuper13

src/connections.py

@@ -7,7 +7,7 @@
 import concurrent.futures
 
 from src.util import PortType, add_vlan_tag, fix_checksums, apply_bpf_filter
-from src.manipulation import ManipulateArgs
+from src.manipulation import ManipulateArgs, ManipulateActions
 
 
 DeviceEntry = namedtuple('DeviceEntry', [
@@ -35,6 +35,7 @@ def __init__(self):
 
         self._manipulate_cb = None
         self._manipulate_filter = NO_FILTER
+        self._manipulate_dup = False
 
     async def _read_message_queue(self):
         while True:
@@ -46,29 +47,48 @@ async def _process_packet(self, packet, src_device_entry):
         # to tag the packet.
         should_inject_raw = False
 
+        packets = [packet]
+
         if self._manipulate_cb is not None:
             # We'll manipulate if there's no filter or if there's
             # filter and the packet matches the filter
             if self._manipulate_filter == NO_FILTER or \
                     apply_bpf_filter(packet, self._manipulate_filter):
-                packet, should_inject_raw = await self._event_loop.run_in_executor(
+
+                # In case the packet matches the filter, and `duplicate` is not
+                # set, we remove the packet from the list, since the manipulator
+                # might drop the packet.
+                if not self._manipulate_dup:
+                    packets.remove(packet)
+
+                packet, action = await self._event_loop.run_in_executor(
                     self._executers,
                     self._manipulate_cb,
                     ManipulateArgs(packet, src_device_entry.port_type, src_device_entry.vlan))
 
-        src_vlan = self._get_vlan_by_mac(Ether(packet).src)
-        if src_vlan is None:
-            return None
+                if action == ManipulateActions.INJECT_RAW:
+                    should_inject_raw = True
+                    packets.append(packet)
+                elif action == ManipulateActions.HANDLE_ENCAP:
+                    should_inject_raw = False
+                    packets.append(packet)
+                # else - action is DROP, so we don't add the packet to the list.
+
+        for packet in packets:
+
+            src_vlan = self._get_vlan_by_mac(Ether(packet).src)
+            if src_vlan is None:
+                return None
 
-        # Looking for the destination device
-        for dst_device in self._devices:
-            if dst_device.dev.get_mac == Ether(packet).dst and dst_device.vlan == src_vlan:
+            # Looking for the destination device
+            for dst_device in self._devices:
+                if dst_device.dev.get_mac == Ether(packet).dst and dst_device.vlan == src_vlan:
 
-                if dst_device.port_type == PortType.TRUNK and not should_inject_raw:
-                    packet = add_vlan_tag(packet, dst_device.vlan)
+                    if dst_device.port_type == PortType.TRUNK and not should_inject_raw:
+                        packet = add_vlan_tag(packet, dst_device.vlan)
 
-                packet = fix_checksums(packet)
-                await self._event_loop.sock_sendall(dst_device.sock, packet)
+                    packet = fix_checksums(packet)
+                    await self._event_loop.sock_sendall(dst_device.sock, packet)
 
     def _read_raw_packet(self, device_entry):
         try:
@@ -146,11 +166,12 @@ def _remove_device_entry(self, dev_to_remove):
 
         self._event_loop.call_soon_threadsafe(_remove_device_entry, self, dev_to_remove)
 
-    def set_manipulation(self, cb, bpf_filter):
+    def set_manipulation(self, cb, bpf_filter, duplicate):
         ''' Assumes cb is in valid form '''
         def __set_manipulation(cb, bpf_filter):
             self._manipulate_cb = cb
             self._manipulate_filter = bpf_filter
+            self._manipulate_dup = duplicate
 
         # set_manipulation is called from the main thread.
         # Since we're affecting the event loop thread, we should call_soon_threadsafe,

src/manipulation.py

@@ -1,6 +1,16 @@
+import enum
 import inspect
 from collections import namedtuple
 
+
+class ManipulateActions(enum.Enum):
+    DROP = 0
+    # The manipulator will deal with tagging, if needed.
+    INJECT_RAW = 1
+    # The switch should deal with tagging, if needed.
+    HANDLE_ENCAP = 2
+
+
 ManipulateArgs = namedtuple('ManipulateInfo', [
     'packet',
     # can be used if the manipulator wants to deal with tagging.
@@ -10,9 +20,8 @@
 
 ManipulateRet = namedtuple('ManipulateRet', [
     'packet',
-    # In case you want to deal with tagging (add dot1q layer),
-    # this should be set to True.
-    'should_inject_raw',
+    # From ManipulateActions
+    'action',
 ])
 
 
@@ -21,7 +30,7 @@ def default_manipulation_cb(manipulate_args: ManipulateArgs) -> ManipulateRet:
     Can be used as an example of a valid signature of manipulation callback.
     This default callback just returns the packet.
     '''
-    return ManipulateRet(manipulate_args.packet, False)
+    return ManipulateRet(manipulate_args.packet, ManipulateActions.HANDLE_ENCAP)
 
 
 def validate_manipulation_cb(cb):

src/switch.py

@@ -51,7 +51,7 @@ def _set_trunk_connection(self, dev, vlan):
 
         return True
 
-    def set_manipulation(self, cb, punt_policies_bpf=NO_FILTER):
+    def set_manipulation(self, cb, punt_policies_bpf=NO_FILTER, duplicate=False):
         """
         Sets the manipulation routine.
 
@@ -64,13 +64,32 @@ def set_manipulation(self, cb, punt_policies_bpf=NO_FILTER):
         In addition, In case you want to emulate "Punt-Policies" (decide what packets
         should be forwarded to the manipulation callback), you can pass a bpf filter.
         As mentioned, only packets that are filtered by this bpf filter will be forwarded
-        to the manipultion routine before continuing the switch's flow.
+        to the manipultion routine before they get processed by the switch.
 
         :param cb: A callback the the manipulation routine.
+                   The callback must have arguments and return annotations,
+                   since this function validates this callback using them.
+
+                   A valid callback has to match the following form:
+                   1. Recieve one argument. Its type is ManipulateArgs.
+                   2. Return ManipulateRet.
+
+                   An example of a valid callback (taken from manipulation.py:
+                   `default_manipulation_cb`):
+
+                def default_manipulation_cb(manipulate_args: ManipulateArgs) -> ManipulateRet:
+                    return ManipulateRet(manipulate_args.packet, ManipulateActions.HANDLE_ENCAP)
+
+                    This callback for example, just returns the packet.
+
         :param punt_policies_bpf: The "Punt-Policies". Default is no-filter.
+        :param duplicate: True if the packet should be processed by both switch
+                          and manipulator, False if the packet should be processed
+                          by the manipulator only.
+        :return: True if the callback is in a valid form, False otherwise.
         """
         if validate_manipulation_cb(cb):
-            self._connections.set_manipulation(cb, punt_policies_bpf)
+            self._connections.set_manipulation(cb, punt_policies_bpf, duplicate)
             return True
 
         return False

tests/test_manipulation.py

@@ -2,7 +2,8 @@
 from scapy.all import Ether, IP, ICMP, Raw
 
 from src.device import Device
-from src.manipulation import ManipulateArgs, ManipulateRet
+from src.manipulation import ManipulateArgs, ManipulateRet, ManipulateActions
+from src.util import add_vlan_tag
 
 
 D1_ADDR = '192.168.250.1'
@@ -29,7 +30,19 @@ def forward_d1_to_d3(arg: ManipulateArgs) -> ManipulateRet:
     packet[Ether].src = GLOBAL_D4_MAC
     packet[Ether].dst = GLOBAL_D3_MAC
 
-    return ManipulateRet(Raw(packet).load, False)
+    return ManipulateRet(Raw(packet).load, ManipulateActions.HANDLE_ENCAP)
+
+
+def drop_packet(arg: ManipulateArgs) -> ManipulateRet:
+    return ManipulateRet(None, ManipulateActions.DROP)
+
+
+def tag_packet_inject_raw(arg: ManipulateArgs) -> ManipulateRet:
+    return ManipulateRet(add_vlan_tag(arg.packet, arg.src_vlan), ManipulateActions.INJECT_RAW)
+
+
+def tag_packet_handle_encap(arg: ManipulateArgs) -> ManipulateRet:
+    return ManipulateRet(add_vlan_tag(arg.packet, arg.src_vlan), ManipulateActions.HANDLE_ENCAP)
 
 
 def test_vlan_hopping_and_punt_policies(switch):
@@ -65,7 +78,7 @@ def test_vlan_hopping_and_punt_policies(switch):
 
     # Setting the manipulation routine, and punt-policies
     switch.set_manipulation(forward_d1_to_d3, 'src host {} && dst host {}'.format(
-        D1_ADDR, D2_ADDR))
+        D1_ADDR, D2_ADDR), False)
 
     def _run_listening_nc(dev):
         out = dev.run_from_namespace('timeout 7s nc -lu 0.0.0.0 4444')
@@ -89,3 +102,62 @@ def _run_listening_nc(dev):
     switch.disconnect_device(d2)
     switch.disconnect_device(d3)
     switch.disconnect_device(d4)
+
+
+def test_manipulator_drop_packets(switch):
+    d1 = Device('a', '192.168.250.1', '255.255.255.0')
+    d2 = Device('b', '192.168.250.2', '255.255.255.0')
+
+    switch.connect_device_access(d1, 20)
+    switch.connect_device_access(d2, 20)
+
+    switch.set_manipulation(drop_packet, duplicate=False)
+
+    # No duplication, and the manipulator drops everything
+    # so we expect to recieve no response.
+    out = d1.run_from_namespace('ping -c 1 192.168.250.2 -W 2')
+    assert '1 packets transmitted, 0 received' in out
+
+    # Now we set `duplicate` to True, which means that even
+    # if we drop packets, the original packet will still be processed
+    # by the switch.
+    switch.set_manipulation(drop_packet, duplicate=True)
+
+    # So now we expect to recieve the response
+    out = d1.run_from_namespace('ping -c 1 192.168.250.2')
+    assert '1 packets transmitted, 1 received' in out
+
+    switch.disconnect_device(d1)
+    switch.disconnect_device(d2)
+
+
+def test_manipulator_inject_raw(switch):
+    d1 = Device('a', '192.168.250.1', '255.255.255.0')
+    d2 = Device('b', '192.168.250.2', '255.255.255.0')
+
+    switch.connect_device_trunk(d1, 20)
+    switch.connect_device_trunk(d2, 20)
+
+    # The manipulator will tag all packets, and will set INJECT_RAW
+    # as the action. The original packets won't get processed by switch,
+    # So if we get a valid connection, it means that the manipulator
+    # successfully tagged the packet, and the switch didn't deal with
+    # with tagging; it just send the packet as is (raw).
+    switch.set_manipulation(tag_packet_inject_raw, duplicate=False)
+
+    out = d1.run_from_namespace('ping -c 1 192.168.250.2')
+    assert '1 packets transmitted, 1 received' in out
+
+    out = d2.run_from_namespace('ping -c 1 192.168.250.1')
+    assert '1 packets transmitted, 1 received' in out
+
+    switch.set_manipulation(tag_packet_handle_encap, duplicate=False)
+
+    out = d1.run_from_namespace('ping -c 1 192.168.250.2 -W 2')
+    assert '1 packets transmitted, 0 received' in out
+
+    out = d2.run_from_namespace('ping -c 1 192.168.250.1 -W 2')
+    assert '1 packets transmitted, 0 received' in out
+
+    switch.disconnect_device(d1)
+    switch.disconnect_device(d2)