#!/usr/libexec/platform-python

import argparse
import random
import string
import subprocess
import signal
import sys

# Colors
RED = "\033[91m"
GREEN = "\033[92m"
YELLOW = "\033[93m"
BLUE = "\033[94m"
RESET = "\033[0m"
INFO = GREEN

# Global variable to hold the dummy net device name
GLOBAL_DUMMY_NDEV = None

def generate_random_string(length=4):
    """Generate a random string of uppercase letters."""
    return ''.join(random.choices(string.ascii_letters, k=length))

def create_dummy_netdevice():
    """Create a dummy net device with a random name."""
    global GLOBAL_DUMMY_NDEV
    random_suffix = generate_random_string()
    dummy_ndev = f"smc-dummy{random_suffix}"

    try:
        subprocess.run(f"ip link add dev {dummy_ndev} type dummy", shell=True, check=True)
        subprocess.run(f"ip link set dev {dummy_ndev} up", shell=True, check=True)
    except subprocess.CalledProcessError as e:
        print(f"{RED}Failed to create dummy netdevice {dummy_ndev}: {e}{RESET}")
        sys.exit(0)

    try:
        subprocess.run(f"smcr dumpdev set {dummy_ndev}", shell=True, check=True)
    except subprocess.CalledProcessError as e:
        print(f"{RED}Failed to set smcr dumpdev for {dummy_ndev}: {e}{RESET}")
        subprocess.run(f"ip link delete dev {dummy_ndev}", shell=True, check=True)
        sys.exit(0)

    GLOBAL_DUMMY_NDEV = dummy_ndev
    print(f"{INFO}Create dummy netdevice {dummy_ndev} for SMC dump function{RESET}")

    return dummy_ndev

def destroy_dummy_netdevice(dummy_ndev):
    subprocess.run("smcr dumpdev reset", shell=True, check=True)
    subprocess.run(f"ip link delete dev {dummy_ndev}", shell=True, check=True)
    print(f"{INFO}Remove dummy netdevice {dummy_ndev} used by SMC dump function{RESET}")

def signal_handler(sig, _frame):
    """Signal handler for cleanup."""
    global GLOBAL_DUMMY_NDEV
    print(f"{INFO}Received signal: {sig}{RESET}")
    if GLOBAL_DUMMY_NDEV:
        destroy_dummy_netdevice(GLOBAL_DUMMY_NDEV)
    sys.exit(0)

def get_dump_netdevice(mode, dummy_ndev):
    """Get the dump net device based on mode."""
    if mode in ('smc', 'smcr', 'smcd'):
        return dummy_ndev
    return 'any'

def get_base_filter(mode, ptype, legacy):
    """Get the base filter for SMC dump."""
    filters = []

    if legacy:
        filters.append("udp and (udp[8:1] = 0xcf) and (udp[9:1] & 0xF0 = 0x10)")

        if mode in ('smcr', 'smcd'):
            print(f"{RED}WARNING: "
                  f"legacy format does not support 'smcr' or 'smcd' filters, "
                  f"it turns to be 'smc'.{RESET}")

        if ptype == 'raw':
            filters.append("and (udp[9:1] & 0x0F = 0x01)")
        elif ptype == 'cdc':
            filters.append("and (udp[9:1] & 0x0F = 0x02)")
    else:
        filters.append("udp and (udp[8:4] = 0xcfd3e7a5)")

        if mode == 'smcr':
            filters.append("and (udp[12:1] & 0x0F = 0x01)")
        elif mode == 'smcd':
            filters.append("and (udp[12:1] & 0x0F = 0x02)")

        if ptype == 'raw':
            filters.append("and (udp[13:1] & 0x0F = 0x01)")
        elif ptype == 'cdc':
            filters.append("and (udp[13:1] & 0x0F = 0x02)")

    if mode == 'all':
        base_filter = f"tcp or ({' '.join(filters)})"
    else:
        base_filter = ' '.join(filters)

    return base_filter


def run_tcpdump(dump_ndev, base_filter, other_params):
    """Run tcpdump with the specified parameters."""
    command = [f"tcpdump -i {dump_ndev}"]

    if base_filter:
        command.append(f"'{base_filter}'")

    if other_params:
        command.append(other_params)

    final_command = ' '.join(command)
    print(f"{INFO}Command: {final_command}{RESET}")
    try:
        subprocess.run(final_command, shell=True, check=True)
    except subprocess.CalledProcessError as e:
        print(f"Error running tcpdump: {e}")

def main():
    # Set up signal handling
    signal.signal(signal.SIGINT, signal_handler)  # Handle Ctrl+C
    signal.signal(signal.SIGTERM, signal_handler)  # Handle kill signal

    parser = argparse.ArgumentParser(description="SMC Dump - SMC Traffic Capture (Experimental)")
    parser.add_argument('-m', '--mode', choices=['all', 'smc', 'smcd', 'smcr'],
            default='smc', help="Select the mode (default: smc)")
    parser.add_argument('-t', '--type', choices=['all', 'raw', 'cdc'], default='all',
            help="Select the packet type (default: all)")
    parser.add_argument('--param', type=str, default='',
            help="Additional parameters for tcpdump. e.g. --param '-w packets.pcap'")
    parser.add_argument('--filter', type=str, default='',
            help="Additional filter expressions for tcpdump. "
                "e.g. --filter 'host 192.168.0.1 and port 8080'")
    parser.add_argument('--legacy', action='store_true',
            help="Use the legacy SMC dump header format")  # Compatible with legacy format

    args = parser.parse_args()

    # Create smc dummy interface
    dummy_ndev = create_dummy_netdevice()

    # Get dump net device(s)
    dump_ndev = get_dump_netdevice(args.mode, dummy_ndev)

    # Get basic filters
    base_filter = get_base_filter(args.mode, args.type, args.legacy)

    # Combine filters
    user_filter = args.filter.strip()
    if user_filter:
        final_filter = f"({base_filter}) and ({user_filter})"
    else:
        final_filter = base_filter

    # Combine parameters
    params = args.param.strip()
    if params:
        params = params.split()

    # Run tcpdump
    run_tcpdump(dump_ndev, final_filter, " ".join(params))

    # Clean up the dummy net device
    destroy_dummy_netdevice(dummy_ndev)

if __name__ == "__main__":
    main()
