parallel_reset.py

"""
Performs parallel factory resets on multiple Cisco IR800 routers simultaneously.
Takes the number of routers as a command-line argument, finds active USB ports,
and runs a factory reset on each router in its own thread. Handles login,
flash and nvram cleanup, reboot, and ROMMON-2 mode automatically.

Usage:
    python3 parallel_reset.py <number_of_routers>
"""

import serial
import time
import os
import glob
import subprocess
import re
import threading
import sys
from dotenv import load_dotenv

# Load the router login credentials from the .env file
load_dotenv()

username = os.getenv("ROUTER_USERNAME")
password = os.getenv("ROUTER_PASSWORD")
enable_password = os.getenv("ROUTER_ENABLE_PASSWORD")

# Keeps track of completed routers
completed_routers = []
completed_lock = threading.Lock()


def close_screen_sessions():
    """
    Closes all active screen sessions to free the serial ports.
    Runs `screen -list` to check for active sessions.
    If any sessions are found, they are closed using `screen -X quit`.
    """

    result = subprocess.run(["screen", "-list"], capture_output=True, text=True)

    if "No Sockets found" not in result.stdout:
        subprocess.run(["screen", "-X", "quit"])
        time.sleep(1)

        print("Screen sessions closed")

    else:
        print("No screen sessions found")


def find_all_ports(antal):
    """
    Searches for and returns the requested number of active USB serial ports.
    Loops until enough active ports are found by sending a newline
    and checking if the router responds. Retries every 5 seconds
    if not enough active ports are found.

    Arguments:
        antal (int): Number of active ports required before returning.

    Returns:
        list: A list of active port paths, for example
        ['/dev/ttyUSB0', '/dev/ttyUSB1'].
    """

    print(f"Searching for {antal} active USB ports...")

    while True:
        active_ports = []
        ports = glob.glob("/dev/ttyUSB*")

        for port in ports:
            try:
                test = serial.Serial(port=port, baudrate=9600, timeout=2)
                test.write(b"\r\n")
                time.sleep(3)

                output = test.read(test.in_waiting).decode("utf-8", errors="ignore")

                test.close()

                time.sleep(1)

                if output:
                    print(f"Found active port: {port}")

                    active_ports.append(port)

            except:
                continue

        if len(active_ports) >= antal:
            print(f"Found {len(active_ports)} ports, starting cleanup...")
            return active_ports[:antal]

        print(f"Found {len(active_ports)} of {antal} ports, retrying in 5 seconds...")
        time.sleep(5)


def send_and_wait(ser, command, wait=3):
    """
    Sends a command through the serial connection and waits for a response
    from the router. Reads the output until a known prompt or keyword
    is found, or until the 15 second timeout is reached.

    Arguments:
        ser (serial.Serial): The active serial connection.
        command (str): The command to send to the router.
        wait (float): Seconds to wait before reading output. Default is 3.

    Returns:
        str: The decoded output from the router.
    """

    ser.write(f"{command}\r\n".encode())
    time.sleep(wait)

    output = ""
    deadline = time.time() + 15

    while time.time() < deadline:
        if ser.in_waiting:
            output += ser.read(ser.in_waiting).decode("utf-8", errors="ignore")

        time.sleep(0.2)

        if any(p in output for p in ["#", ">", "confirm", "save", "[OK]", "dialog", "complete", "erase", "yes/no", "filename", "rommon", "deleted", "No such", "Error"]):
            break

    return output


def login_and_enable(ser, port):
    """
    Logs in to the router and activates enable mode.
    Attempts up to 10 times to log in using the username and password,
    and then activate enable mode using the enable password.

    Arguments:
        ser (serial.Serial): The active serial connection.
        port (str): The port path used for logging.

    Returns:
        bool: True if enable mode was activated, otherwise False.
    """

    MAX_ATTEMPTS = 10

    for attempt in range(MAX_ATTEMPTS):
        print(f"[{port}] Login attempt {attempt + 1} of {MAX_ATTEMPTS}...")

        ser.write(b"\r\n")
        time.sleep(6)

        output = ser.read(ser.in_waiting).decode("utf-8", errors="ignore")

        if "(config)" in output:
            send_and_wait(ser, "end")

        if "Username" in output:
            output = send_and_wait(ser, username)

            if "Password" in output:
                output = send_and_wait(ser, password)

        if ">" in output:
            output = send_and_wait(ser, "enable")

            if "Password" in output:
                output = send_and_wait(ser, enable_password)

        if "(config)" in output:
            send_and_wait(ser, "end")

        if "#" in output:
            print(f"[{port}] Enable mode activated!")
            return True

        print(f"[{port}] Could not reach enable mode, retrying in 5 seconds...")
        time.sleep(5)

    print(f"[{port}] Maximum number of attempts reached, exiting")
    return False
  
  def clean_flash(ser, port):
    """
    Cleans the flash memory and removes all files except protected system files.
    Lists the flash contents and removes all files that do not match
    the protected file list. Short numeric filenames are also skipped.

    Arguments:
        ser (serial.Serial): The active serial connection.
        port (str): The port path used for logging.
    """

    print(f"[{port}] Cleaning flash...")
    ser.write(b"dir flash:\r\n")

    output = ""
    deadline = time.time() + 20

    while time.time() < deadline:
        if ser.in_waiting:
            output += ser.read(ser.in_waiting).decode("utf-8", errors="ignore")

        if "#" in output:
            break

        time.sleep(0.3)

    never_delete = [
        "eem",
        "managed",
        "ir800-universalk9",
        "ir800-hv",
        "ir800-uk9",
        "gemboa",
        "MC7304",
        "no_usb_emul"
    ]

    all_files = re.findall(r"(?:^|\s)([\w\-\.]+)(?:\s|$)", output, re.MULTILINE)

    for fil in all_files:
        if not fil or any(skip in fil for skip in never_delete):
            continue

        if re.match(r"^\d+$", fil) and len(fil) < 6:
            continue

        send_and_wait(ser, f"delete /force /recursive flash:{fil}", wait=3)

    print(f"[{port}] Flash cleaned!")


def clean_nvram(ser, port):
    """
    Cleans nvram and removes all files.
    Lists the nvram contents and removes all files that are not
    short numeric filenames.

    Arguments:
        ser (serial.Serial): The active serial connection.
        port (str): The port path used for logging.
    """

    print(f"[{port}] Cleaning nvram...")

    ser.write(b"dir nvram:\r\n")
    output = ""

    deadline = time.time() + 20

    while time.time() < deadline:
        if ser.in_waiting:
            output += ser.read(ser.in_waiting).decode("utf-8", errors="ignore")

        if "#" in output:
            break

        time.sleep(0.3)

    all_files = re.findall(r"(?:^|\s)([\w\-\.]+)(?:\s|$)", output, re.MULTILINE)

    for fil in all_files:
        if not fil:
            continue

        if re.match(r"^\d+$", fil) and len(fil) < 6:
            continue

        send_and_wait(ser, f"delete /force nvram:{fil}", wait=3)

    print(f"[{port}] Nvram cleaned!")


def find_ios_file_rommon(ser, port):
    """
    Searches for the IOS file in flash while the router is in ROMMON mode.
    Sends `dir flash:` and reads the response until the ROMMON prompt appears again.
    Loops until a matching IOS file is found.

    Arguments:
        ser (serial.Serial): The active serial connection.
        port (str): The port path used for logging.

    Returns:
        str: The IOS image filename, for example `ir800-universalk9-mz.xxx`.
    """

    print(f"[{port}] Searching for IOS file...")

    while True:
        ser.write(b"dir flash:\r\n")
        output = ""

        deadline = time.time() + 20

        while time.time() < deadline:
            if ser.in_waiting:
                output += ser.read(ser.in_waiting).decode("utf-8", errors="ignore")

            if "rommon" in output.lower() and output.count("rommon") >= 2:
                break

            time.sleep(0.3)

        match = re.search(r"ir800-universalk9-mz\S+", output)

        if match:
            ios_file = match.group(0).strip().rstrip(".")

            print(f"[{port}] Found IOS file: {ios_file}")

            return ios_file

        print(f"[{port}] Could not find IOS file, retrying in 5 seconds...")
        time.sleep(5)


def save_boot_variable(ser, port, ios_file):
    """
    Saves the boot variable permanently in the router startup configuration.
    Enters enable and configuration mode, removes existing
    boot directives, sets a new boot variable, and saves using `write memory`.

    Arguments:
        ser (serial.Serial): The active serial connection.
        port (str): The port path used for logging.
        ios_file (str): The IOS image filename used during boot.

    Returns:
        bool: True if the save was successful, otherwise False.
    """

    print(f"[{port}] Saving boot variable permanently...")
    time.sleep(3)

    ser.write(b"enable\r\n")
    time.sleep(3)

    ser.write(b"conf t\r\n")
    time.sleep(3)

    ser.write(b"no boot system\r\n")
    time.sleep(3)

    ser.write(f"boot system flash:{ios_file}\r\n".encode())
    time.sleep(3)

    ser.write(b"end\r\n")
    time.sleep(3)

    ser.write(b"write memory\r\n")
    output = ""

    deadline = time.time() + 30

    while time.time() < deadline:
        if ser.in_waiting:
            output += ser.read(ser.in_waiting).decode("utf-8", errors="ignore")

        if "[OK]" in output or "Building" in output:
            time.sleep(3)
            break

        time.sleep(0.3)

    if "[OK]" in output or "Building" in output:
        print(f"[{port}] Boot variable saved permanently!")
        return True

    else:
        print(f"[{port}] Could not save boot variable!")
        return False

  def wait_for_boot(ser, port):
    """
    Waits for the router and handles ROMMON-2 mode automatically.
    Monitors serial output and reacts to different states: ROMMON-2,
    'Press RETURN' prompt, initial configuration dialog, and normal IOS prompt.

    If the router gets stuck in ROMMON-2, it is booted manually from flash.
    When IOS is up, the boot variable is saved permanently.

    Arguments:
        ser (serial.Serial): The active serial connection.
        port (str): The port path used for logging.

    Returns:
        bool: True if the router successfully booted into IOS,
        otherwise False on timeout.
    """

    print(f"[{port}] Waiting for the router...")
    output = ""

    rommon_handled = False
    ios_booted = False
    ios_file = None

    last_enter = time.time()

    deadline = time.time() + 400

    while time.time() < deadline:
        if not ser.in_waiting and time.time() - last_enter > 5:
            ser.write(b"\r\n")

            last_enter = time.time()

        if ser.in_waiting:
            chunk = ser.read(ser.in_waiting).decode("utf-8", errors="ignore")
            output += chunk

            # Handle ROMMON-2
            if "rommon" in output.lower() and not rommon_handled:
                print(f"[{port}] Router in ROMMON-2, booting up...")
                time.sleep(3)

                ios_file = find_ios_file_rommon(ser, port)

                ser.write(f"set BOOT=flash:{ios_file}\r\n".encode())
                time.sleep(2)

                ser.write(b"set BOOT_IOS_SEQUENCE=20\r\n")
                time.sleep(2)

                ser.write(f"boot flash:{ios_file}\r\n".encode())
                time.sleep(2)

                rommon_handled = True

                output = ""

            # Handle Press RETURN
            if "press ret" in output.lower() or "press return" in output.lower():
                time.sleep(2)

                ser.write(b"\r\n")
                time.sleep(2)

                output = ""

            # Handle initial configuration dialog
            if "yes/no" in output.lower() and "configuration dialog" in output.lower():
                print(f"[{port}] Answering no to configuration dialog...")
                time.sleep(2)

                ser.write(b"no\r\n")
                time.sleep(5)

                output = ""

            # Router booted into IOS
            if "IR800>" in output and not ios_booted and ios_file:
                print(f"[{port}] Router booted into IOS!")
                ios_booted = True

                time.sleep(8)

                save_boot_variable(ser, port, ios_file)

                return True

        time.sleep(0.2)

    print(f"[{port}] Timeout!")
    return False


def reset_router(port):
    """
    Performs a full factory reset on a router through a serial connection.
    Logs in, runs write erase, cleans flash and nvram, reloads the router,
    and waits for it to boot correctly. Registers the router as completed
    when the process is finished.

    Arguments:
        port (str): The path to the serial port connected to the router.
    """

    print(f"[{port}] Starting factory reset...")

    try:
        ser = serial.Serial(port=port, baudrate=9600, timeout=5)

        if not login_and_enable(ser, port):
            ser.close()
            return

        print(f"[{port}] Running write erase...")

        output = send_and_wait(ser, "write erase", wait=5)

        if "confirm" in output.lower() or "erase" in output.lower():
            send_and_wait(ser, "", wait=8)

        print(f"[{port}] Write erase completed!")

        clean_flash(ser, port)
        clean_nvram(ser, port)

        print(f"[{port}] Reloading the router...")

        output = send_and_wait(ser, "reload", wait=5)

        if "yes/no" in output.lower() or "save" in output.lower() or "modified" in output.lower():
            send_and_wait(ser, "no", wait=5)

        if "confirm" in output.lower() or "proceed" in output.lower():
            send_and_wait(ser, "", wait=3)

        if "confirm" in output.lower() or "proceed" in output.lower():
            send_and_wait(ser, "", wait=3)

        time.sleep(5)

        wait_for_boot(ser, port)

        # Register this router as completed
        with completed_lock:
            completed_routers.append(port)

            print(f"[{port}] Factory reset completed! ({len(completed_routers)} of {antal_routrar} routers completed)")

        ser.close()

    except Exception as e:
        print(f"[{port}] Error: {e}")


# Main program - check that the number of routers was provided as an argument
if len(sys.argv) < 2:
    print("Usage: python3 parallel_reset.py <number_of_routers>")
    print("Example: python3 parallel_reset.py 3")
    exit()

antal_routrar = int(sys.argv[1])

print(f"Preparing cleanup of {antal_routrar} router(s)...")

close_screen_sessions()
time.sleep(2)

active_ports = find_all_ports(antal_routrar)

print(f"Starting parallel factory reset on {len(active_ports)} router(s)...")

# Start one thread per router
threads = []

for port in active_ports:
    t = threading.Thread(target=reset_router, args=(port,))
    threads.append(t)
    t.start()

# Wait until all threads are finished
for t in threads:
    t.join()

print(f"\nAll {len(completed_routers)} of {antal_routrar} routers have been cleaned!")