switch.py

#!/usr/bin/python3
import sys
import struct
import wrapper
import threading
import time
from wrapper import recv_from_any_link, send_to_link, get_switch_mac, get_interface_name

def parse_ethernet_header(data):
    # Unpack the header fields from the byte array
    #dest_mac, src_mac, ethertype = struct.unpack('!6s6sH', data[:14])
    dest_mac = data[0:6]
    src_mac = data[6:12]
    
    # Extract ethertype. Under 802.1Q, this may be the bytes from the VLAN TAG
    ether_type = (data[12] << 8) + data[13]

    vlan_id = -1
    # Check for VLAN tag (0x8100 in network byte order is b'\x81\x00')
    if ether_type == 0x8200:
        vlan_tci = int.from_bytes(data[14:16], byteorder='big')
        vlan_id = vlan_tci & 0x0FFF  # extract the 12-bit VLAN ID
        ether_type = (data[16] << 8) + data[17]

    return dest_mac, src_mac, ether_type, vlan_id

def create_vlan_tag(vlan_id):
    # 0x8100 for the Ethertype for 802.1Q
    # vlan_id & 0x0FFF ensures that only the last 12 bits are used
    return struct.pack('!H', 0x8200) + struct.pack('!H', vlan_id & 0x0FFF)

def remove_vlan_tag(data):
    # Remove the VLAN tag from the frame
    return data[0:12] + data[16:]

def open_bpdu_packet(data):
    # Unpack the header fields from the byte array
    BPDU_Config = data[6 + 6 + 2 + 3 + 4:]
    # Unpack the BPDU fields from the byte array
    BPDU_root_bridge_ID, BPDU_sender_bridge_ID, BPDU_root_path_cost = struct.unpack('!3i', BPDU_Config)
    return BPDU_root_bridge_ID, BPDU_sender_bridge_ID, BPDU_root_path_cost

def send_bdpu_every_sec():
    # Create and send a BPDU every second if the switch is the root to all switchs
    while True:
        if (is_root):
            for i in interfaces:
                if (get_interface_name(i) in switchs):
                    data = create_bpdu_packet()
                    send_to_link(i, len(data), data)
        time.sleep(1)

def create_bpdu_packet():
    # Create a BPDU packet with the following fields
    BPDU_Header = struct.pack('!i', 0)
    LCC_Header = struct.pack('!3b', 0x42, 0x42, 0x03)
    BPDU_Config = struct.pack('!3i', root_bridge_ID, own_bridge_ID, root_path_cost)
    ethernet_header = struct.pack('!6s6sH', b'\x01\x80\xc2\x00\x00\x00', get_switch_mac(), 
                                                len(LCC_Header + BPDU_Header + BPDU_Config))
    # Return the created BPDU packet
    return ethernet_header + LCC_Header + BPDU_Header + BPDU_Config

def read_data_vlan_config(switch_id):
    # Read VLAN configuration from the file
    filepath = f"configs/switch{switch_id}.cfg"
    with open(filepath, 'r') as f:
        vlan_configs = f.readlines()

    # Extract the \n character from each line
    vlan_configs = [line.strip() for line in vlan_configs]
    return vlan_configs

def is_switch(interface, switchs):
    # Check if the interface is a switch
    return interface in switchs

def send_data(data, port, dest_port):
    # Verify if the source port is a switch or a host
    if (is_switch(get_interface_name(port), switchs)):
        # If the destination port is a switch, send the data without changing it
        if (is_switch(get_interface_name(dest_port), switchs)):
            print(f"From switch {get_interface_name(port)} to switch {get_interface_name(dest_port)}")
            send_to_link(dest_port, len(data), data)
        else:
            # If the destination port is a host, remove the VLAN tag and send the data
            data = remove_vlan_tag(data)
            dest_vlan_id = hosts[hosts.index(get_interface_name(dest_port)) + 1]
            print(f"VLAN ID: {vlan_id}")
            # Check if the VLAN ID of the source is the same as the destination
            if (vlan_id == int(dest_vlan_id)):
                print(f"From switch {get_interface_name(port)} to host {get_interface_name(dest_port)}")
                send_to_link(dest_port, len(data), data)
    # If the source port is a host, check if the destination port is a switch or a host
    else:
        # If the destination port is a switch, add the VLAN tag and send the data
        if (is_switch(get_interface_name(dest_port), switchs)):
            dest_vlan_id = hosts[hosts.index(get_interface_name(port)) + 1]
            data = data[0:12] + create_vlan_tag(int(dest_vlan_id)) + data[12:]
            print(f"From host {get_interface_name(port)} to switch {get_interface_name(dest_port)}")
            send_to_link(dest_port, len(data), data)
        else:
            # If the destination port is a host, check if the VLAN IDs are the same and send the data
            send_vlan_id = hosts[hosts.index(get_interface_name(port)) + 1]
            dest_vlan_id = hosts[hosts.index(get_interface_name(dest_port)) + 1]
            if (send_vlan_id == dest_vlan_id):
                print(f"From host {get_interface_name(port)} to host {get_interface_name(dest_port)}")
                send_to_link(dest_port, len(data), data)

def handle_stp_frame(port, data):
    global root_bridge_ID, root_path_cost, root_port, ports
    print(f"Received BPDU from {src_mac}")
    # Unpack the BPDU fields from the data
    BPDU_root_bridge_ID, BPDU_sender_bridge_ID, BPDU_root_path_cost = open_bpdu_packet(data)

    # Check if the root bridge ID of the BPDU is less than our root bridge ID
    if (BPDU_root_bridge_ID < root_bridge_ID):
        print(f"BPDU root bridge ID {BPDU_root_bridge_ID} is less than root bridge ID {root_bridge_ID}")
        # Update the root bridge ID, root path cost, and root port
        root_bridge_ID = BPDU_root_bridge_ID
        root_path_cost = BPDU_root_path_cost + 10

        # Block all ports except the one connected to the root if we where the root
        if (root_port == -1):
            for i in interfaces:
                if (get_interface_name(i) in switchs):
                    ports[i] = "BLOCKED"

        root_port = port

        # Unblock the port that is connected to the root
        if (ports[port] == "BLOCKED"):
            ports[port] = "DESIGNATED"

        # Send a BPDU to all switchs that we found a better path to the root
        for i in interfaces:
            if (get_interface_name(i) in switchs):
                data = create_bpdu_packet()
                send_to_link(i, len(data), data)
    # If the root bridge ID of the BPDU is the same as our root bridge ID
    elif (BPDU_root_bridge_ID == root_bridge_ID):
        # We check if the path cost of the BPDU is less than our path cost
        if (root_port == port and BPDU_sender_bridge_ID + 10 < root_path_cost):
            root_path_cost = BPDU_root_path_cost + 10
        # If the path cost of the BPDU is greater than our path cost
        elif (port != root_port):
            if (BPDU_root_path_cost > root_path_cost):
                # We set the port to LISTEN for hosts and DESIGNATED for switches
                if (get_interface_name(port) not in switchs):
                    ports[port] = "LISTEN"
                else:
                    ports[port] = "DESIGNATED"

    # If the sneder bridge ID of the BPDU is the same as our bridge ID we block the port
    elif (BPDU_sender_bridge_ID == own_bridge_ID):
        ports[port] = "BLOCKED"

def handle_frame(port, MAC_Table, data):
    if (ports[port] == "BLOCKED"):
        return
    # Learn the source MAC address
    MAC_Table[src_mac] = port

    # Check if the destination MAC address is a unicast or broadcast address
    if int(dest_mac[0:2], 16) & 0x01 == 0:
        # Check if the destination MAC address is in the table
        if dest_mac in MAC_Table:
            dest_port = MAC_Table[dest_mac]
            send_data(data, port, dest_port)
        else:
            # Flood the frame to all ports except the source port
            print(f"Flooding unicast from {src_mac} to {dest_mac} via all ports except port {port}")
            for i in interfaces:
                if i != port:
                    send_data(data, port, i)
    else:
        # Check if the frame is a BPDU frame
        if (str(dest_mac) != "01:80:c2:00:00:00"):
            # Flood the frame to all ports except the source port
            print(f"Flooding frame from {src_mac} to {dest_mac} via all ports except port {port}")
            for i in interfaces:
                if i != port:
                    send_data(data, port, i)
        else:
            handle_stp_frame(port, data)

    # Check if we are the root bridge
    if (own_bridge_ID == root_bridge_ID):
        # Set all ports to DESIGNATED except the root port
        for i in ports:
            if (i != root_port and get_interface_name(i) in switchs):
                ports[i] = "DESIGNATED"
        # Update the variable for the root
        global is_root
        is_root = True
        global root_path_cost
        root_path_cost = 0
    else:
        # Set is_root to False cuz we are not the root bridge and should not send BPDU
        is_root = False


def main():
    # init returns the max interface number. Our interfaces
    # are 0, 1, 2, ..., init_ret value + 1
    switch_id = sys.argv[1]

    num_interfaces = wrapper.init(sys.argv[2:])
    global is_root
    is_root = True
    vlan_configs = read_data_vlan_config(switch_id)
    # Extract own_bridge_ID from the VLAN configuration
    global own_bridge_ID
    own_bridge_ID = int(vlan_configs[0].split()[0])

    # Exctract hosts from the VLAN configuration
    global hosts 
    hosts = vlan_configs[1].split() + vlan_configs[2].split()

    # Extract switchs from the VLAN configuration
    global switchs
    switchs = [item for item in vlan_configs[3].split() + vlan_configs[4].split() if item != 'T']

    # Spanning Tree Protocol variables
    global root_bridge_ID
    root_bridge_ID = own_bridge_ID
    global root_path_cost
    root_path_cost = 0
    global root_port
    root_port = -1
    global ports
    ports = {}

    global interfaces
    interfaces = range(0, num_interfaces)
    # Printing interface names
    for i in interfaces:
        print(get_interface_name(i))

    MAC_Table = {}

    # Set all ports to DESIGNATED
    for i in interfaces:
        ports[i] = "DESIGNATED"

    print("# Starting switch with id {}".format(switch_id), flush=True)
    print("[INFO] Switch MAC", ':'.join(f'{b:02x}' for b in get_switch_mac()))

    # Create and start a new thread that deals with sending BDPU
    t = threading.Thread(target=send_bdpu_every_sec)
    t.start()

    while True:
        # Note that data is of type bytes([...]).
        # b1 = bytes([72, 101, 108, 108, 111])  # "Hello"
        # b2 = bytes([32, 87, 111, 114, 108, 100])  # " World"
        # b3 = b1[0:2] + b[3:4].
        interface, data, length = recv_from_any_link()

        global dest_mac, src_mac, vlan_id

        dest_mac, src_mac, ethertype, vlan_id = parse_ethernet_header(data)

        # Print the MAC src and MAC dst in human readable format
        dest_mac = ':'.join(f'{b:02x}' for b in dest_mac)
        src_mac = ':'.join(f'{b:02x}' for b in src_mac)

        # Note. Adding a VLAN tag can be as easy as
        # tagged_frame = data[0:12] + create_vlan_tag(10) + data[12:]

        print(f'Destination MAC: {dest_mac}')
        print(f'Source MAC: {src_mac}')
        print(f'EtherType: {ethertype}')

        print("Received frame of size {} on interface {}".format(length, interface), flush=True)

        # Implement forwarding with learning / Implement VLAN support / handle_frame STP support
        handle_frame(interface, MAC_Table, data)

if __name__ == "__main__":
    main()