fix_rommon.py

"""
Handles a Cisco IR809 router stuck in ROMMON-2 mode.
Automatically finds the correct USB port, detects the router state,
and boots the router from flash if it is in ROMMON-2 mode.

The script then saves the boot variable permanently in the startup configuration
so the router can boot normally in the future.
"""

import serial
import time
import glob
import subprocess
import re


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_port():
    """
    Searches for and returns an active USB serial port.

    The function loops until an active port is found by sending a newline
    and checking if the router responds. If no active port is found, it
    retries every 5 seconds.

    Returns:
        str: The path to the active port, for example '/dev/ttyUSB0'.
    """

    print("Searching for USB port...")

    while True:
        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}")
                    return port

            except:
                continue

        print("No active port found, retrying in 5 seconds...")
        time.sleep(5)


def find_ios_file_rommon():
    """
    Searches for the IOS file in flash while the router is in ROMMON mode.

    Sends the `dir flash:` command and reads the response until the ROMMON
    prompt appears again. The function loops until a matching IOS file is found.

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

    print("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"Found IOS file: {ios_file}")
            return ios_file

        print("Could not find IOS file, retrying in 5 seconds...")
        time.sleep(5)


def find_ios_file_ios():
    """
    Searches for the IOS file in flash while the router is running IOS.
    Sends the `show flash:` command and reads the response until an IOS
    or enable prompt appears.

    Returns:
        str or None: The IOS image filename if found,
        otherwise None.
    """

    print("Searching for IOS file...")

    ser.write(b"show 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 or "IR800>" in output:
            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"Found IOS file: {ios_file}")
        return ios_file

    print("Could not find IOS file!")
    return None


def save_boot_variable(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 it using `write memory`.

    Arguments:
        ios_file (str): The IOS image filename to use during boot.

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

    print("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("Boot variable saved permanently!")
        print("The router will now boot normally without ROMMON-2!")
        return True

    else:
        print("Could not save boot variable!")
        return False


def fix_rommon():
    """
    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.

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

    print("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("Router in ROMMON-2, booting up...")

                time.sleep(3)
                ios_file = find_ios_file_rommon()

                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("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("Router booted into IOS!")

                ios_booted = True

                time.sleep(8)
                save_boot_variable(ios_file)

                return True

        time.sleep(0.2)

    print("Timeout!")
    return False


# Main program - closes screen sessions and finds active port
close_screen_sessions()
time.sleep(2)

port = find_port()

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

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

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


# Check the router state and act accordingly
if "rommon" in output.lower():
    print("Router is in ROMMON-2, fixing...")
    fix_rommon()

elif "IR800>" in output or "IR800#" in output:
    print("Router is already running IOS...")

    ios_file = find_ios_file_ios()

    if ios_file:
        save_boot_variable(ios_file)

    else:
        print("Could not find IOS file!")

else:
    print("Waiting for the router...")
    fix_rommon()


print("Done!")

ser.close()