src: Add trunk/access connection utilities between Switch and Devices

Author: roykuper13

src/__init__.py

@@ -0,0 +1,119 @@
+import socket
+import asyncio
+import threading
+from scapy.all import Ether
+from collections import namedtuple
+from contextlib import suppress
+
+from src.util import PortType, add_vlan_tag, fix_checksums
+
+
+DeviceEntry = namedtuple('DeviceEntry', [
+    'dev',
+    'sock',
+    'vlan',
+    # Trunk/Access
+    'port_type',
+])
+
+IP_MAX_SIZE = 65535
+
+
+class Connections(object):
+
+    def __init__(self):
+        self._connections_thread = threading.Thread(target=self._start_event_loop)
+        self._event_loop = None
+        self._devices = list()
+        self._message_queue = None
+
+    async def _read_message_queue(self):
+        while True:
+            packet, device_entry = await self._message_queue.get()
+            await self._process_packet(packet, device_entry)
+
+    async def _process_packet(self, packet, src_device_entry):
+        dst_mac = Ether(packet).dst
+
+        # Looking for the destination device
+        for dst_device in self._devices:
+            if dst_device.dev.get_mac == dst_mac and dst_device.vlan == src_device_entry.vlan:
+
+                if dst_device.port_type == PortType.TRUNK:
+                    packet = add_vlan_tag(packet, dst_device.vlan)
+
+                packet = fix_checksums(packet)
+                await self._event_loop.sock_sendall(dst_device.sock, packet)
+
+    def _read_raw_packet(self, device_entry):
+        try:
+            packet = device_entry.sock.recv(IP_MAX_SIZE)
+            self._message_queue.put_nowait((packet, device_entry))
+        except OSError:
+            # TODO: Currently the interface is closed successfully, but the cb
+            # raises an exception.
+            pass
+
+    def _start_event_loop(self):
+        self._event_loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(self._event_loop)
+
+        self._message_queue = asyncio.Queue()
+        asyncio.ensure_future(self._read_message_queue())
+
+        try:
+            self._event_loop.run_forever()
+
+            for task in asyncio.Task.all_tasks():
+                task.cancel()
+
+                with suppress(asyncio.CancelledError):
+                    # await until task is completed if it is currently running/pending.
+                    self._event_loop.run_until_complete(task)
+        finally:
+            self._event_loop.close()
+
+    def _update_arp_tables(self, new_dev):
+        for curr_dev in self._devices:
+            # If both devices in the same vlan, they should update each other with their addresses.
+            if new_dev.vlan == curr_dev.vlan:
+                new_dev.dev.run_from_namespace('arp -s {ip} {mac}'.format(
+                    ip=curr_dev.dev.get_ip, mac=curr_dev.dev.get_mac))
+
+                curr_dev.dev.run_from_namespace('arp -s {ip} {mac}'.format(
+                    ip=new_dev.dev.get_ip, mac=new_dev.dev.get_mac))
+
+    def append_device(self, dev, vlan, port_type):
+        def _append_device_entry(self, new_dev_entry):
+            self._devices.append(new_dev_entry)
+
+        sock = socket.socket(socket.AF_PACKET, socket.SOCK_RAW, socket.htons(0x0003))
+        sock.bind(('br-{}'.format(dev.get_name), 0))
+
+        new_dev_entry = DeviceEntry(dev, sock, vlan, port_type)
+
+        self._update_arp_tables(new_dev_entry)
+
+        self._event_loop.add_reader(sock, self._read_raw_packet, new_dev_entry)
+
+        # We're pushing the job into the event loop, so no lock is needed.
+        # threadsafe since we're calling it from the main thread,
+        # and not from the eventloop thread.
+        self._event_loop.call_soon_threadsafe(_append_device_entry, self, new_dev_entry)
+
+    def remove_device(self, dev_to_remove):
+        def _remove_device_entry(self, dev_to_remove):
+            for dev_entry in self._devices:
+                if dev_entry.dev == dev_to_remove:
+                    dev_entry.sock.close()
+                    self._event_loop.remove_reader(dev_entry.sock)
+                    self._devices.remove(dev_entry)
+                    break
+
+        self._event_loop.call_soon_threadsafe(_remove_device_entry, self, dev_to_remove)
+
+    def start_connections_thread(self):
+        self._connections_thread.start()
+
+    def stop_connections_thread(self):
+        self._event_loop.call_soon_threadsafe(self._event_loop.stop)

src/connections.py

@@ -0,0 +1,109 @@
+from src.util import shell_run_and_check, convert_subnetmask_to_cidr, run_shell_cmd
+from src.exception import LongInterfaceException, NamespaceCreationException
+
+
+class Device(object):
+
+    # IFNAMSIZ is 15 (excluding the null terminator).
+    # We're appending at most 8 bytes to the given interface name, so 7
+    # characters are left for the user.
+    # Note that we prefer to limit the user's input, instead of generating
+    # names by ourselves, as the user should know which resources this program
+    # is creating (for example, the user might manually delete interfaces
+    # by himself).
+    INTERFACE_NAME_MAX_SIZE = 7
+
+    netns_commands = [
+        # Create veth pair
+        'ip link add veth-{name} type veth peer name br-veth-{name}',
+        'ip link set br-veth-{name} up',
+        # Set the new network namespace's veth interface
+        'ip link set veth-{name} netns {name}',
+        'ip netns exec {name} ip link set veth-{name} up',
+    ]
+
+    def __init__(self, name, ip, subnet_mask):
+        if len(name) > self.INTERFACE_NAME_MAX_SIZE:
+            raise LongInterfaceException(
+                "Device name length must be at most {} characters({})".format(
+                    self.INTERFACE_NAME_MAX_SIZE, name))
+
+        # Create the network namespace for the device
+        if not shell_run_and_check('ip netns add {}'.format(name)):
+            raise NamespaceCreationException(
+                "Can not create network namespace '{}' "
+                "(Make sure to run the script with root priviliages)".format(name))
+
+        self._dev_name = name
+        self._ip = ip
+        self._cidr = convert_subnetmask_to_cidr(subnet_mask)
+        self._mac = None
+
+        self._ns_is_set = False
+
+    def _get_ns_veth_mac(self):
+        """
+        In case the veth interface of the new namespace is set (setup_namespace),
+        The function returns the mac address of the veth interface (in utf-8).
+        Otherwise, the function returns None.
+        """
+        mac, err = run_shell_cmd(
+            'ip netns exec {ns} cat /sys/class/net/veth-{ns}/address'.format(ns=self._dev_name))
+
+        return mac.decode('utf-8') if err == b'' else None
+
+    @property
+    def get_name(self):
+        return self._dev_name
+
+    @property
+    def get_ip(self):
+        return self._ip
+
+    @property
+    def get_cidr(self):
+        return self._cidr
+
+    @property
+    def get_mac(self):
+        if self._mac is None:
+            self._mac = self._get_ns_veth_mac()
+
+        return self._mac
+
+    def get_switch_veth_name(self):
+        # This class creates veth pair. One veth is set to the new network namespace
+        # in `setup_namespace`. The second veth should be used by the default namespace,
+        # later on (See switch.py). Therefore, the default namespace should be able
+        # to get the second veth's name.
+        return 'br-veth-{}'.format(self._dev_name)
+
+    def get_device_veth_name(self):
+        # That is the veth interface in the "device" side (i.e - in the new network
+        # namespace side).
+        return 'veth-{}'.format(self._dev_name)
+
+    def run_from_namespace(self, cmd):
+        out, err = run_shell_cmd('ip netns exec {ns} {cmd}'.format(ns=self._dev_name, cmd=cmd))
+        return out.decode('utf-8') if err == b'' else err.decode('utf-8')
+
+    def setup_namespace(self):
+        ''' Safe to call more then once. '''
+        if self._ns_is_set:
+            return True
+
+        if not shell_run_and_check(' && '.join(self.netns_commands).format(name=self._dev_name)):
+            return False
+
+        self._ns_is_set = True
+        return True
+
+    def term(self):
+        ''' Safe to call more then once. '''
+        if self._ns_is_set:
+            shell_run_and_check('ip netns del {}'.format(self._dev_name))
+            # We've created the switch's veth interface, so we are responsible for
+            # deleting it.
+            shell_run_and_check('ip link del {}'.format(self.get_switch_veth_name()))
+
+            self._ns_is_set = False

src/device.py

@@ -0,0 +1,14 @@
+class LongInterfaceException(Exception):
+    pass
+
+
+class NamespaceCreationException(Exception):
+    pass
+
+
+class BridgeInterfaceCreationException(Exception):
+    pass
+
+
+class NamespaceConnectionException(Exception):
+    pass

src/exception.py

@@ -0,0 +1,92 @@
+from src.connections import Connections
+from src.util import shell_run_and_check, PortType
+from src.exception import (NamespaceCreationException,
+                           BridgeInterfaceCreationException, NamespaceConnectionException)
+
+
+class Switch(object):
+
+    def __init__(self):
+        self._connections = Connections()
+        self._connections.start_connections_thread()
+
+    def _set_bridge_interface(self, dev):
+        set_bridge_cmds = [
+            'ip link add name {br} type bridge',
+            'ip link set {br} up',
+            'ip link set {br} promisc on',
+            'ip link set {veth} master {br}',
+        ]
+
+        return shell_run_and_check(' && '.join(set_bridge_cmds).format(
+            br='br-{}'.format(dev.get_name), veth=dev.get_switch_veth_name()))
+
+    def _set_access_connection(self, dev):
+        # Set IP to veth interface (within the device's network namespace)
+        if not shell_run_and_check(
+            'ip netns exec {name} ip addr add {ip}/{cidr} dev veth-{name}'.format(
+                name=dev.get_name, ip=dev.get_ip, cidr=dev.get_cidr)):
+            return False
+
+        return True
+
+    def _set_trunk_connection(self, dev, vlan):
+        vlan_interface_cmds = [
+            'ip link add link {iface} name {iface}.v type vlan id {vlan}',
+            'ip addr add {ip}/{cidr} brd + dev {iface}.v',
+            'ip link set {iface}.v up',
+        ]
+
+        # Configure vlan interface inside the network namespace
+        concat_cmds = ' && '.join(
+            ['ip netns exec {} '.format(dev.get_name) + cmd for cmd in vlan_interface_cmds]).format(
+            iface=dev.get_device_veth_name(),
+            ip=dev.get_ip,
+            cidr=dev.get_cidr,
+            vlan=vlan)
+
+        if not shell_run_and_check(concat_cmds):
+            return False
+
+        return True
+
+    def connect_device_access(self, dev, vlan):
+        if not dev.setup_namespace():
+            raise NamespaceCreationException("failed to create network namespace for {}".format(
+                dev.get_name))
+
+        if not self._set_bridge_interface(dev):
+            raise BridgeInterfaceCreationException(
+                "failed to create bridge interface for {}".format(dev.get_name))
+
+        if not self._set_access_connection(dev):
+            raise NamespaceConnectionException("failed to create connection to namespace {}".format(
+                dev.get_name))
+
+        self._connections.append_device(dev, vlan, PortType.ACCESS)
+
+    def connect_device_trunk(self, dev, vlan):
+        if not dev.setup_namespace():
+            raise NamespaceCreationException("failed to create network namespace for {}".format(
+                dev.get_name))
+
+        if not self._set_bridge_interface(dev):
+            raise BridgeInterfaceCreationException(
+                "failed to create bridge interface for {}".format(dev.get_name))
+
+        if not self._set_trunk_connection(dev, vlan):
+            raise NamespaceConnectionException("failed to create connection to namespace {}".format(
+                dev.get_name))
+
+        self._connections.append_device(dev, vlan, PortType.TRUNK)
+
+    def disconnect_device(self, dev):
+        # Delete switch's bridge interface
+        shell_run_and_check('ip link del br-{}'.format(dev.get_name))
+
+        self._connections.remove_device(dev)
+
+        dev.term()
+
+    def term(self):
+        self._connections.stop_connections_thread()