HomeGuidesAPI Reference
GuidesAPI ReferenceGitHubAirheads Developer CommunityLog In

firmware_site_distribution.py

Switch Firmware Script

Setup

To setup and run this script first copy both code blocks into separate python files in the same directory. Please make sure to name the files the same as shown. Then, install the requirements listed below. For full documentation and usage of this script see here.

firmware_site_distribution.py

# MIT License
#
# Copyright (c) 2024 Aruba, a Hewlett Packard Enterprise company
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import threading
import urllib3
import sys
import csv
import os
from concurrent.futures import ThreadPoolExecutor
from argparse import ArgumentParser
from collections import deque
from termcolor import cprint
from getpass import getpass
from switch import Switch
from time import sleep
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

root_host = None

def define_arguments():
    """
    Defines command line arguments.

    :return: Argparse namespace with central auth filepath and ssid config
        filepath.
    :type return: argparse.Namespace
    """

    description = "Reimage switches from other host switches. "\
                  "Optionally reboot all reimaged switches. Input"\
                  " csv uses the first switch as a host to start."\
                  " csv format: external ip, internal ip, user, pass."

    parser = ArgumentParser(description=description)

    parser.add_argument('input', type=validate_path,
                         help=('switch csv filepath'))
    parser.add_argument('threads', type=validate_threads,
                        help=('number of concurrent threads reimaging'))
    parser.add_argument('-vrf', help="set vrf param", default="default")
    parser.add_argument('-r', help="reboot all reimaged switches",
                        action='store_true')
    parser.add_argument('-cred', help="prompt for login credentials", action='store_true')
    parser.add_argument('-d', help="client image destination",
                        default="secondary", type=validate_location)
    parser.add_argument('-s', help="source image location",
                        default="primary", type=validate_location)
    return parser.parse_args()

def validate_path(path):
    """Validates input filepath."""
    if not os.path.exists(path):
        sys.exit("Invalid filepath for input argument. Exiting...")
    return path

def validate_threads(threads):
    """Validates thread input."""
    t = int(threads)
    if t < 1:
        sys.exit("Invalid number of threads. Must be greater than zero. Exiting...")
    return t

def validate_location(location):
    if location != 'secondary' and location != 'primary':
        sys.exit("Invalid destination or source location. Valid options are"\
                 " primary or secondary. Exiting...")
    return location

def reimage(client, host, pool, sccs, err):
    """Download firmware image from host on new thread.
    Setup client for hosting and add to host pool.
    
    Keyword arguments:
    client -- client Switch object
    host -- host Switch object
    pool -- host pool
    sccs -- list of successful Switch updates
    err -- list of failed Switch updates
    """
    thread = threading.get_ident()
    print(f"    streaming from host {host.ip}")
    print(f"    operating on switch {client.ip} on thread {thread}.")

    if host == root_host:
        res = client.update_firm(host, ARGS.d, ARGS.s, ARGS.vrf)
    else:
        res = client.update_firm(host, ARGS.d, ARGS.d, ARGS.vrf)
    if res.status_code != 200:
        print(f"Update failed on switch {client.ip} with {res.status_code}")
        image_fail(client, host, pool, err)
        return
    
    # Wait for reimage to complete.
    updated = False
    while not updated:
        res = client.status()
        body = res.json()
        if body["status"] == "in_progress":
            sleep(5)
            print(f"    {client.ip} in progress sleeping on thread {thread}")
        elif body["status"] == "failure":
            image_fail(client, host, pool, err) 
            return
        else:
            updated = True
            sccs.append(client)

    host.seed -= 1
    if host not in pool:
        pool.append(host)
    # Enable client to become a host.
    client.setup_remote()
    pool.append(client)
    cprint(f"    Reimage complete on switch {client.ip} in thread {thread}", "green")

def image_fail(client, host, pool, err):
    """On failure to image, add client to error list and reset host."""
    err.append(client)
    host.seed -= 1
    if host not in pool:
        pool.append(host)

def main():
    threads = ThreadPoolExecutor(ARGS.threads)
    if ARGS.cred:
        print("Enter credential for switches")
        user = input("user: ")
        pw = getpass(prompt="pass: ")
    
    # get list of switch info from input
    inputs = []
    with open(ARGS.input, mode='r') as file:
        f = csv.reader(file)
        for line in f:
            inputs.append(line)

    host_pool = deque([])
    clients = []
    updated = []
    fails = []

    # create switch objects, setup host and client pools
    print("Setting up switch objects...")
    for switch in (inputs):
        if ARGS.cred:
            new_switch = Switch(switch[0], switch[1], user, pw)
        else:
            new_switch = Switch(switch[0], switch[1], switch[2], switch[3])
        res = new_switch.login()
        if res.status_code != 200:
            if len(host_pool) == 0:
                sys.exit("Failed to login on root host. Check credentials. Exiting...")
            cprint(f"   login failed for {new_switch.ip}. Removing from inputs.", "red")
            continue
        if len(host_pool) == 0:
            # Designate root switch
            global root_host
            root_host = new_switch
            new_switch.setup_remote()
            host_pool.append(new_switch)
        else:    
            clients.append(new_switch)

    try:
        print("\nBegin reimaging...")
        # Loop until clients is empty
        while clients:  
            # Wait for host to be availble
            while not host_pool:
                sleep(1)
            client = clients.pop()
            host = host_pool.pop()
            host.seed += 1
            if host.seed == 1:      
                host_pool.appendleft(host)
            threads.submit(reimage, client, host, host_pool, updated, fails)
            
        threads.shutdown()
        
        if not fails:
            cprint("Reimaging successful for all switches!", "green")
        else:
            cprint("Reimaging process complete with errors.", "red")

        print("\nCleaning up remote config...")
        for switch in host_pool:
            print(f"    Turning off remote endpoint for {switch.ip}")
            switch.setup_remote(False)

    finally:
        if fails:
            print("\nFailed to update on switches..")
            for switch in fails:
                cprint(f"    {switch.ip}", "red")

        if ARGS.r:
            print("\nFinishing up, rebooting switches...")
            for switch in updated:
                switch.boot()
            print("Reboots in progress, script complete.")
        else:       
            print("\nFinishing up, logging out of switches...")
            for switch in host_pool:
                switch.logout()
            for switch in fails:
                switch.logout()
            print("\nScript complete!")


if __name__ == "__main__":
    ARGS = define_arguments()
    main()

switch.py

# MIT License
#
# Copyright (c) 2024 Aruba, a Hewlett Packard Enterprise company
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import requests

class Switch:
    def __init__(self, ip, internal, user, pword):
        self.ip = ip
        self.internal = internal
        self.user = user
        self.pword = pword
        self.seed = 0
        self.session = requests.session()
        self.proxies = {'https': None, 'http': None}

    def login(self):
        """Login and establish session for switch authentication."""
        url = f"https://{self.ip}/rest/latest/login"
        payload = {"username": self.user, "password": self.pword}
        resp = self.session.post(url, params=payload, verify=False, proxies=self.proxies)
        print(f"    Login code {resp} for switch {self.ip}")
        return resp

    def update_firm(self, host, dest="primary", source="primary", vrf="default"):
        """Update self firmware image from remote host.

        Keyword arguments:
        host -- switch object to serve as remote host
        dest -- image location (primary/secondary)
        source -- remote image location (primary/secondary)
        vrf -- vrf name
        """
        path = f"https://{host.internal}/fwimages/{source}.swi"
        url = f"https://{self.ip}/rest/latest/firmware"
        payload = {
            "image": dest,
            "from": path,
            "vrf": vrf
        }
        resp = self.session.put(url, params=payload, verify=False, proxies=self.proxies)
        return resp
    
    def status(self):
        """Get status of firmware update."""
        url = f"https://{self.ip}/rest/latest/firmware/status"
        resp = self.session.get(url, verify=False, proxies=self.proxies)
        return resp
    
    def setup_remote(self, option=True):
        """Toggle firmware site distribution."""
        url = f"https://{self.ip}/rest/latest/system/rest_config"
        payload = {"firmware_site_distribution_enabled": option}
        resp = self.session.put(url, json=payload, verify=False, proxies=self.proxies)
        return resp
    
    def boot(self, image="primary"):
        """Reboot switch from image"""
        url = f"https://{self.ip}/rest/latest/boot"
        payload = {"image": image}
        resp = self.session.post(url, params=payload, verify=False, proxies=self.proxies)
        print(f"    Switch {self.ip} rebooted to {image}")
        return resp

    def logout(self):
        """Logout of switch."""
        url = f"https://{self.ip}/rest/latest/logout"
        resp = self.session.post(url, verify=False, proxies=self.proxies)
        print(f"    Switch {self.ip} logged out")
        return resp

Requirements

The following python modules need to be installed:
Requests
Termcolor


What’s Next

Find full documentation by following the link!