vswitch.py

"""
File: vswitch.py
Function: Simple Virtual Switch implementation
Author: Cyans from Chang'an University
"""
import socket
import sys


if len(sys.argv) != 2:
    print("Usage: python3 vswitch.py <port>")
    sys.exit(1)

LISTEN_PORT = int(sys.argv[1])
SERVER_ADDRESS = ("0.0.0.0", LISTEN_PORT)

vswitch_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
vswitch_socket.bind(SERVER_ADDRESS)
print(f"[VSwitch] Started. Listening on: {SERVER_ADDRESS[0]}:{SERVER_ADDRESS[1]}")

mac_address_table = {}

while True:
    ethernet_frame, sender_address = vswitch_socket.recvfrom(1518)
    ethernet_header = ethernet_frame[:14]
    destination_mac = ":".join(f"{byte:02x}" for byte in ethernet_header[0:6])
    source_mac = ":".join(f"{byte:02x}" for byte in ethernet_header[6:12])
    
    print(f"[VSwitch] From<{sender_address}> SrcMAC<{source_mac}> DstMAC<{destination_mac}> Size<{len(ethernet_frame)} bytes>")
    
    # MAC Address Learning
    if source_mac not in mac_address_table or mac_address_table[source_mac] != sender_address:
        mac_address_table[source_mac] = sender_address
        print(f"    [Table Update]: {mac_address_table}")
    
    # Forwarding Logic
    if destination_mac in mac_address_table:
        # Unicast
        target_address = mac_address_table[destination_mac]
        vswitch_socket.sendto(ethernet_frame, target_address)
        print(f"    [Forwarded] To: {destination_mac}")
    
    elif destination_mac == "ff:ff:ff:ff:ff:ff":
        # Broadcast
        all_macs = list(mac_address_table.keys())
        
        if source_mac in all_macs:
            all_macs.remove(source_mac) # Don't send back to source
        broadcast_targets = {mac_address_table[mac] for mac in all_macs}
        
        if broadcast_targets:
            print(f"    [Broadcast] Targets: {broadcast_targets}")
            for target in broadcast_targets:
                vswitch_socket.sendto(ethernet_frame, target)
        else:
            print("    [Broadcast] No other ports active")
    else:
        # Unknown Unicast -> Drop
        print(f"    [Drop] Unknown DstMAC: {destination_mac}")