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
Updated 2 months ago