fix_switch_backup_firmware.py

"""
-------------------------------------------------------------------------------

    Written by Thomas Munzer (tmunzer@juniper.net)
    Github repository: https://github.com/tmunzer/Mist_library/

    This script is licensed under the MIT License.

-------------------------------------------------------------------------------
Python script to automate firmware snapshot/backup operations on Juniper EX. 
This script processes CSV reports from the report_switch_firmware.py script to
identify and remediate switches with missing snapshot or backup firmware
partitions.

Key features:
- Automated identification of switches needing firmware backup
- CSV-based input from firmware reporting tools
- Site-specific filtering for targeted operations
- Bulk processing with progress tracking
- Safety confirmation prompts (unless auto-approved)
- Support for both standalone and clustered switches

How it works:
1. Reads CSV report identifying switches needing firmware backup
2. Filters switches based on site criteria (if specified)
3. Displays list of affected switches for user review
4. Triggers firmware snapshot/backup API calls for each switch
5. Monitors and reports operation success/failure

Use cases:
- Preventive maintenance to ensure backup firmware availability
- Remediation after firmware upgrade issues
- Compliance with backup firmware policies
- Disaster recovery preparation
- Mass deployment firmware backup operations

IMPORTANT: This script triggers firmware backup operations that may impact
switch performance during execution. Use during maintenance windows when possible.

-------
Requirements:
mistapi: https://pypi.org/project/mistapi/

-------
Usage:
This script can be run as is (without parameters), or with the options below.
If no options are defined, or if options are missing, the missing options will
be asked by the script or the default values will be used.

It is recommended to use an environment file to store the required information
to request the Mist Cloud (see https://pypi.org/project/mistapi/ for more
information about the available parameters).

-------
Options:
-h, --help          display this help

-s, --site_id=      Set the site_id if only devices from a specific site must be
                    processed. If not set, the script will process all the devices
                    from the report that need a firmware backup
-f, --in_file=      path to the report generated by the report_switch_firmware script
                    default is "./report_switch_firmware.csv"
--auto-approve      Does not ask for user confirmation before triggering the
                    firmware backup API Calls

-l, --log_file=     define the filepath/filename where to write the logs
                    default is "./script.log"
-e, --env=          define the env file to use (see mistapi env file documentation
                    here: https://pypi.org/project/mistapi/)
                    default is "~/.mist_env"

-------
Examples:
python3 ./fix_switch_backup_firmware.py
python3 ./fix_switch_backup_firmware.py --site_id=203d3d02-xxxx-xxxx-xxxx-76896a3330f4

"""

#### IMPORTS #####
import sys
import csv
import argparse
import logging

MISTAPI_MIN_VERSION = "0.53.0"

try:
    import mistapi
    from mistapi.__logger import console as CONSOLE
except ImportError:
    print("""
        Critical: 
        \"mistapi\" package is missing. Please use the pip command to install it.

        # Linux/macOS
        python3 -m pip install mistapi

        # Windows
        py -m pip install mistapi
        """)
    sys.exit(2)


#### LOGS ####
LOGGER = logging.getLogger(__name__)

#### PARAMETERS #####
CSV_FILE = "./report_switch_firmware.csv"
LOG_FILE = "./script.log"
ENV_FILE = "~/.mist_env"


#####################################################################
# PROGRESS BAR AND DISPLAY
class ProgressBar:
    """Progress bar for long-running operations."""
    def __init__(self):
        self.steps_total = 0
        self.steps_count = 0

    def _pb_update(self, size: int = 80):
        if self.steps_count > self.steps_total:
            self.steps_count = self.steps_total

        percent = self.steps_count / self.steps_total
        delta = 17
        x = int((size - delta) * percent)
        print("Progress: ", end="")
        print(f"[{'█' * x}{'.' * (size - delta - x)}]", end="")
        print(f"{int(percent * 100)}%".rjust(5), end="")

    def _pb_new_step(
        self,
        message: str,
        result: str,
        inc: bool = False,
        size: int = 80,
        display_pbar: bool = True,
    ):
        if inc:
            self.steps_count += 1
        text = f"\033[A\033[F{message}"
        print(f"{text} ".ljust(size + 4, "."), result)
        print("".ljust(80))
        if display_pbar:
            self._pb_update(size)

    def _pb_title(
        self, text: str, size: int = 80, end: bool = False, display_pbar: bool = True
    ):
        print("\033[A")
        print(f" {text} ".center(size, "-"), "\n\n")
        if not end and display_pbar:
            print("".ljust(80))

    def set_steps_total(self, steps_total: int):
        """Set the total number of steps for the progress bar."""
        self.steps_count = 0
        self.steps_total = steps_total

    def log_message(self, message, display_pbar: bool = True):
        """Log a message."""
        self._pb_new_step(message, " ", display_pbar=display_pbar)

    def log_success(self, message, inc: bool = False, display_pbar: bool = True):
        """Log a success message."""
        LOGGER.info("%s: Success", message)
        self._pb_new_step(
            message, "\033[92m\u2714\033[0m\n", inc=inc, display_pbar=display_pbar
        )

    def log_warning(self, message, inc: bool = False, display_pbar: bool = True):
        """Log a warning message."""
        LOGGER.warning("%s: Warning", message)
        self._pb_new_step(
            message, "\033[93m\u2b58\033[0m\n", inc=inc, display_pbar=display_pbar
        )

    def log_failure(self, message, inc: bool = False, display_pbar: bool = True):
        """Log a failure message."""
        LOGGER.error("%s: Failure", message)
        self._pb_new_step(
            message, "\033[31m\u2716\033[0m\n", inc=inc, display_pbar=display_pbar
        )

    def log_title(self, message, end: bool = False, display_pbar: bool = True):
        """Log a title message."""
        LOGGER.info("%s", message)
        self._pb_title(message, end=end, display_pbar=display_pbar)


PB = ProgressBar()


###############################################################################
#### FUNCTIONS ####
def _process_switches(apisession: mistapi.APISession, switches: list) -> None:
    PB.set_steps_total(len(switches))
    for switch in switches:
        LOGGER.debug("Processing device: %s", switch)
        if switch.get("cluster_site_id"):
            site_id = switch.get("cluster_site_id")
        elif switch.get("vc_site_id"):
            site_id = switch.get("vc_site_id")
        else:
            site_id = switch.get("site_id")

        if switch.get("cluster_device_id"):
            device_id = switch.get("cluster_device_id")
        elif switch.get("vc_device_id"):
            device_id = switch.get("vc_device_id")
        else:
            device_id = switch.get("device_id")

        device_mac = switch.get("module_mac")
        message = f"Processing device {device_id}"
        PB.log_message(message)
        if not site_id:
            PB.log_failure(message, inc=True)
            CONSOLE.error(f"Missing site_id for device {device_mac}")
            continue
        if not device_id:
            PB.log_failure(message, inc=True)
            CONSOLE.error(f"Missing device_id for device {device_mac}")
            continue
        try:
            resp = mistapi.api.v1.sites.devices.createSiteDeviceSnapshot(
                apisession, site_id, device_id
            )
            if resp.status_code == 200:
                PB.log_success(message, inc=True)
            else:
                PB.log_failure(message, inc=True)
        except Exception:
            PB.log_failure(message, inc=True)
            LOGGER.error("Exception occurred", exc_info=True)


### READ REPORT
def _read_csv(csv_file: str, site_id: str) -> list:
    data = []
    device_ids = []
    message = "Reading CSV Report"
    PB.log_message(message, display_pbar=False)
    try:
        with open(csv_file, newline="") as f:
            reader = csv.DictReader(filter(lambda row: row[0] != "#", f))
            for row in reader:
                if site_id and row.get("vc_site_id") != site_id:
                    LOGGER.debug("Skipping row due to site_id filter: %s", row)
                    continue
                if row.get("fpc_need_snapshot") != "True":
                    LOGGER.debug("Skipping row due to fpc_need_snapshot filter: %s", row)
                    continue
                if row.get("vc_device_id") in device_ids:
                    LOGGER.debug("Skipping row due to duplicate vc_device_id: %s", row)
                    continue
                LOGGER.debug("Adding row to process: %s", row)
                data.append(row)
                device_ids.append(row.get("cluster_device_id"))
        PB.log_success(message, inc=False, display_pbar=False)
    except Exception:
        PB.log_failure(message, inc=False, display_pbar=False)
        LOGGER.error("Exception occurred", exc_info=True)
        sys.exit(1)
    return data


def _request_approval(data: list):
    print("".center(80, "-"))
    print("List of switches to process:")
    print()
    headers = []
    for entry in data:
        for key in entry:
            if key not in headers:
                headers.append(key)
    mistapi.cli.display_list_of_json_as_table(data, headers)
    r = input("Do you want to continue (y/N)? ")
    if r.lower() == "y":
        print("".center(80, "-"))
        return
    CONSOLE.info("process stopped by the user. Exiting...")
    sys.exit(0)


###############################################################################
### START
def _start(
    apisession: mistapi.APISession,
    site_id: str,
    csv_file: str = CSV_FILE,
    auto_approve: bool = False,
) -> None:
    data = _read_csv(csv_file, site_id)
    if not data:
        print("All the switches are compliant... Exiting...")
        sys.exit(0)
    if auto_approve:
        CONSOLE.info(
            "auto-approve parameter has been set to True. Starting the process"
        )
    elif isinstance(data, list):
        _request_approval(data)
    _process_switches(apisession, data)


###############################################################################
### USAGE
def usage(error_message: str|None = None):
    """
    Display usage information for the script.
    """
    print("""
-------------------------------------------------------------------------------

    Written by Thomas Munzer (tmunzer@juniper.net)
    Github repository: https://github.com/tmunzer/Mist_library/

    This script is licensed under the MIT License.

-------------------------------------------------------------------------------
Python script to automate firmware snapshot/backup operations on Juniper EX. 
This script processes CSV reports from the report_switch_firmware.py script to
identify and remediate switches with missing snapshot or backup firmware
partitions.

Key features:
- Automated identification of switches needing firmware backup
- CSV-based input from firmware reporting tools
- Site-specific filtering for targeted operations
- Bulk processing with progress tracking
- Safety confirmation prompts (unless auto-approved)
- Support for both standalone and clustered switches

How it works:
1. Reads CSV report identifying switches needing firmware backup
2. Filters switches based on site criteria (if specified)
3. Displays list of affected switches for user review
4. Triggers firmware snapshot/backup API calls for each switch
5. Monitors and reports operation success/failure

Use cases:
- Preventive maintenance to ensure backup firmware availability
- Remediation after firmware upgrade issues
- Compliance with backup firmware policies
- Disaster recovery preparation
- Mass deployment firmware backup operations

IMPORTANT: This script triggers firmware backup operations that may impact
switch performance during execution. Use during maintenance windows when possible.

-------
Requirements:
mistapi: https://pypi.org/project/mistapi/

-------
Usage:
This script can be run as is (without parameters), or with the options below.
If no options are defined, or if options are missing, the missing options will
be asked by the script or the default values will be used.

It is recommended to use an environment file to store the required information
to request the Mist Cloud (see https://pypi.org/project/mistapi/ for more 
information about the available parameters).

-------
Options:
-h, --help          display this help

-s, --site_id=      Set the site_id if only devices from a specific site must be 
                    processed. If not set, the script will process all the devices
                    from the report that need a firmware backup
-f, --in_file=      path to the report generated by the report_switch_firmware script
                    default is "./report_switch_firmware.csv"
--auto-approve      Does not ask for user confirmation before triggering the 
                    firmware backup API Calls

-l, --log_file=     define the filepath/filename where to write the logs
                    default is "./script.log"
-e, --env=          define the env file to use (see mistapi env file documentation 
                    here: https://pypi.org/project/mistapi/)
                    default is "~/.mist_env"

-------
Examples:
python3 ./fix_switch_backup_firmware.py                  
python3 ./fix_switch_backup_firmware.py --site_id=203d3d02-xxxx-xxxx-xxxx-76896a3330f4 

""")
    if error_message:
        CONSOLE.critical(error_message)
    sys.exit(0)


def check_mistapi_version():
    """Check if the installed mistapi version meets the minimum requirement."""

    current_version = mistapi.__version__.split(".")
    required_version = MISTAPI_MIN_VERSION.split(".")

    try:
        for i, req in enumerate(required_version):
            if current_version[int(i)] > req:
                break
            if current_version[int(i)] < req:
                raise ImportError(
                    f'"mistapi" package version {MISTAPI_MIN_VERSION} is required '
                    f"but version {mistapi.__version__} is installed."
                )
    except ImportError as e:
        LOGGER.critical(str(e))
        LOGGER.critical("Please use the pip command to update it.")
        LOGGER.critical("")
        LOGGER.critical("    # Linux/macOS")
        LOGGER.critical("    python3 -m pip install --upgrade mistapi")
        LOGGER.critical("")
        LOGGER.critical("    # Windows")
        LOGGER.critical("    py -m pip install --upgrade mistapi")
        print(
            f"""
Critical:\r\n
{e}\r\n
Please use the pip command to update it.
# Linux/macOS
python3 -m pip install --upgrade mistapi
# Windows
py -m pip install --upgrade mistapi
            """
        )
        sys.exit(2)
    finally:
        LOGGER.info(
            '"mistapi" package version %s is required, '
            "you are currently using version %s.",
            MISTAPI_MIN_VERSION,
            mistapi.__version__,
        )


###############################################################################
### ENTRY POINT
if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Trigger a snapshot/firmware backup on EX devices.",
        add_help=False,
    )
    # Add help manually to maintain control over usage function
    parser.add_argument("-h", "--help", action="store_true", help="display this help")

    parser.add_argument(
        "-s",
        "--site_id",
        help="Set the site_id if only devices from a specific site must be processed",
    )
    parser.add_argument(
        "-f",
        "--in_file",
        default=CSV_FILE,
        help="path to the report generated by the report_switch_firmware script (default: ./report_switch_firmware.csv)",
    )
    parser.add_argument(
        "--auto-approve",
        action="store_true",
        help="Does not ask for user confirmation before triggering the firmware backup API Calls",
    )
    parser.add_argument(
        "-e",
        "--env",
        default=ENV_FILE,
        help="define the env file to use (default: ~/.mist_env)",
    )
    parser.add_argument(
        "-l",
        "--log_file",
        default=LOG_FILE,
        help="define the filepath/filename where to write the logs (default: ./script.log)",
    )

    args = parser.parse_args()

    if args.help:
        usage()

    SITE_ID = args.site_id
    CSV_FILE = args.in_file
    AUTO_APPROVE = args.auto_approve
    ENV_FILE = args.env
    LOG_FILE = args.log_file

    #### LOGS ####
    logging.basicConfig(filename=LOG_FILE, filemode="w")
    LOGGER.setLevel(logging.DEBUG)
    check_mistapi_version()
    ### MIST SESSION ###
    APISESSION = mistapi.APISession(env_file=ENV_FILE)
    APISESSION.login()
    ### START ###
    _start(APISESSION, SITE_ID, CSV_FILE, AUTO_APPROVE)