connections: Add an option to manipulate packets before forwarding them

Author: roykuper13

README.md

@@ -7,7 +7,7 @@ as Python objects.
 
 Moreover, linux-switch let you manipulate packets before the network switch forwards
 them (see example below). Thus, External binaries/applications that performs
-any logic on packets (for example, NAT) can be tested using linux-switch.
+any logic on packets (for example - VLAN Hopping, NAT, etc) can be tested using linux-switch.
 
 
 ## Description
@@ -26,6 +26,7 @@ When using access - they'll send untagged packets.
 
 ## Example
 
+### Basic
 ```python
 
 from switch import Switch
@@ -83,3 +84,7 @@ switch.disconnect_device(dev2)
 
 switch.term()
 ```
+
+### Manipulations
+
+For VLAN-Hopping example, see tests/test_manipulation.py (test_vlan_hopping)

src/connections.py

@@ -4,8 +4,10 @@
 from scapy.all import Ether
 from collections import namedtuple
 from contextlib import suppress
+import concurrent.futures
 
 from src.util import PortType, add_vlan_tag, fix_checksums
+from src.manipulation import ManipulateArgs, default_manipulation_cb
 
 
 DeviceEntry = namedtuple('DeviceEntry', [
@@ -18,6 +20,8 @@
 
 IP_MAX_SIZE = 65535
 
+MAX_WORKERS = 5
+
 
 class Connections(object):
 
@@ -26,25 +30,34 @@ def __init__(self):
         self._event_loop = None
         self._devices = list()
         self._message_queue = None
+        self._executers = None
+        self.manipulate_cb = default_manipulation_cb
 
     async def _read_message_queue(self):
         while True:
-            packet = await self._message_queue.get()
-            await self._process_packet(packet)
+            packet, dev_entry = await self._message_queue.get()
+            await self._process_packet(packet, dev_entry)
+
+    async def _process_packet(self, packet, src_device_entry):
+        # Should we deal with tagging. We left an option for the manipulator
+        # to tag the packet.
+        should_inject_raw = False
 
-    async def _process_packet(self, packet):
-        dst_mac = Ether(packet).dst
-        src_mac = Ether(packet).src
+        if self.manipulate_cb is not None:
+            packet, should_inject_raw = 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(src_mac)
+        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 == dst_mac and dst_device.vlan == src_vlan:
+            if dst_device.dev.get_mac == Ether(packet).dst and dst_device.vlan == src_vlan:
 
-                if dst_device.port_type == PortType.TRUNK:
+                if dst_device.port_type == PortType.TRUNK and not should_inject_raw:
                     packet = add_vlan_tag(packet, dst_device.vlan)
 
                 packet = fix_checksums(packet)
@@ -53,7 +66,7 @@ async def _process_packet(self, packet):
     def _read_raw_packet(self, device_entry):
         try:
             packet = device_entry.sock.recv(IP_MAX_SIZE)
-            self._message_queue.put_nowait(packet)
+            self._message_queue.put_nowait((packet, device_entry))
         except OSError:
             # TODO: Currently the interface is closed successfully, but the cb
             # raises an exception.
@@ -63,6 +76,8 @@ def _start_event_loop(self):
         self._event_loop = asyncio.new_event_loop()
         asyncio.set_event_loop(self._event_loop)
 
+        self._executers = concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS)
+
         self._message_queue = asyncio.Queue()
         asyncio.ensure_future(self._read_message_queue())
 

src/manipulation.py

@@ -0,0 +1,37 @@
+import inspect
+from collections import namedtuple
+
+ManipulateArgs = namedtuple('ManipulateInfo', [
+    'packet',
+    # can be used if the manipulator wants to deal with tagging.
+    'is_trunk_port',
+    'src_vlan',
+])
+
+ManipulateRet = namedtuple('ManipulateRet', [
+    'packet',
+    # In case you want to deal with tagging (add dot1q layer),
+    # this should be set to True.
+    'should_inject_raw',
+])
+
+
+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)
+
+
+def validate_manipulation_cb(cb):
+    # TODO: Currently annotations are required in order to check
+    # function's signature
+    sig = inspect.signature(cb)
+
+    if len(sig.parameters) != 1:
+        return False
+
+    param = list(sig.parameters.values())[0]
+
+    return param.annotation == ManipulateArgs and sig.return_annotation == ManipulateRet

src/switch.py

@@ -1,5 +1,6 @@
 from src.connections import Connections
 from src.util import shell_run_and_check, PortType
+from src.manipulation import validate_manipulation_cb
 from src.exception import (NamespaceCreationException,
                            BridgeInterfaceCreationException, NamespaceConnectionException)
 
@@ -88,5 +89,12 @@ def disconnect_device(self, dev):
 
         dev.term()
 
+    def set_manipulation(self, cb):
+        if validate_manipulation_cb(cb):
+            self._connections.manipulate_cb = cb
+            return True
+
+        return False
+
     def term(self):
         self._connections.stop_connections_thread()