#!/usr/bin/python3

###############################################################################
#
# cl-upgrade-bios - Updates the BIOS on Mellanox switches
#
#
# Copyright (c) 2022-2023 NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# This software product is a proprietary product of Nvidia Corporation and its affiliates
# (the "Company") and all right, title, and interest in and to the software
# product, including all associated intellectual property rights, are and
# shall remain exclusively with the Company.
#
# This software product is governed by the End User License Agreement
# provided with the software product.
#
# Description:
#
# This script upgrades the BIOS on Mellanox switches.  The basic flow is:
#
# 1) Ensure that this script is being run with root level access
# 2) Ensure that all of the tools and utilities required by this script 
#    exist.
# 3) Confirm ONIE version installed on the switch is recent enough to
#    support firmware updates.
# 4) Determine what type of BIOS is used by the switch (AtomRangeley,
#    IvyBridge, Broadwell, Coffeelake)
# 5) Extract the BIOS upgrade files to /tmp/BIOS.
# 6) Stage the BIOS upgrade in ONIE.
# 7) Configure the system to enter ONIE update mode after the next reboot.
#
# The script intentionally does NOT initiate a reboot.  Rebooting
# is left to the user (or the automation framework).
#
# Syntax:
#      cl-upgrade-bios [--verbose | --check | --revert | --force]
#
#      --verbose    : Verbose mode
#      --check      : Do everything except the BIOS upgrade
#      --revert     : Undo everything (ie steps 4 and 5 above)
#      --force      : Force the BIOS upgrade (useful for automation)
#
###############################################################################

import argparse
import distutils.spawn
import os
import re
import subprocess
import string
import sys
import syslog

#
# Global variables and "constants"
#

fs_mountpoint           = '/mnt/onie-boot'
onie_path               = '/lib/onie'
min_onie_version        = '5.2.0016'
cpu_type_path           = '/run/hw-management/config/cpu_type'
atom_rangeley_cpu_type  = 0x64d
ivy_bridge_cpu_type     = 0x63a
broadwell_cpu_type      = 0x656
coffee_lake_cpu_type    = 0x69e
atom_rangeley_bios_dir  = '/usr/lib/cumulus/firmware/BIOS/AtomRangeley/'
ivy_bridge_bios_dir     = '/usr/lib/cumulus/firmware/BIOS/IvyBridge/'
broadwell_bios_dir      = '/usr/lib/cumulus/firmware/BIOS/Broadwell/'
coffee_lake_bios_dir    = '/usr/lib/cumulus/firmware/BIOS/CoffeeLake/'
tmp_bios_dir            = '/tmp/BIOS/'
onie_fw_updates_file    = tmp_bios_dir + 'onie_firmware_updates.txt'
onie_fwpkg              = 'onie/tools/bin/onie-fwpkg'

#
# Command line switches
#
verbose = 0
check   = 0
force   = 0
revert  = 0


# Checks if the script is being run with root priviledges
def is_root():
    return os.geteuid() == 0


# Ensures that all of the tools and utilities needed by this script exist
def tools_check():
    required_tools = ['awk', 'dmidecode', 'fdisk', 'grep', 'md5sum', 'mount', 'tar', 'umount']

    for tool in required_tools:
        if distutils.spawn.find_executable(tool) is None:
            if verbose:
                print(('%s not found' % (tool)))  
            return False

    return True


# Mounts the ONIE file system to give CL access to ONIE files, utilities, etc.
# I am well aware that the use of 'shell=True' when using the subprocess module
# is considered to be a security vulnerability.  That being said there does not
# seem to be any way for this code to run without it.
def mount_onie_fs(rw=False):
    if os.path.lexists(onie_path) or os.path.exists(fs_mountpoint):
        umount_onie_fs()

    cmd = "fdisk -l | grep 'ONIE boot' | awk '{print $1}'"
    fs_path = subprocess.check_output(cmd,
                                      stderr=subprocess.STDOUT,
                                      shell=True,
                                      universal_newlines=True).rstrip('\n')
    if fs_path.startswith('/dev') == False:
        if verbose:
            print(('invalid fs_path %s' % (fs_path)))  
        return None

    try:
        os.mkdir(fs_mountpoint)
    except:
        print(('Creation of directory %s failed' % (fs_mountpoint)))
        return None
    if rw:
        cmd = "mount -n -t ext4 {} {}".format(fs_path, fs_mountpoint)
    else:
        cmd = "mount -n -r -t ext4 {} {}".format(fs_path, fs_mountpoint)
    subprocess.check_call(cmd, shell=True, universal_newlines=True)

    fs_onie_path = os.path.join(fs_mountpoint, 'onie/tools/lib/onie')
    try:
        os.symlink(fs_onie_path, onie_path)
    except:
        print(('Creation of sumlink %s failed' % (onie_path)))
        return None

    return fs_mountpoint


# Unmounts the ONIE file system
# I am well aware that the use of 'shell=True' when using the subprocess module
# is considered to be a security vulnerability.  That being said there does not
# seem to be any way for this code to run without it.
def umount_onie_fs():
    if os.path.islink(onie_path):
        os.unlink(onie_path)

    if os.path.ismount(fs_mountpoint):
        cmd = "umount -rf {}".format(fs_mountpoint)
        subprocess.check_call(cmd, shell=True, universal_newlines=True)

    if os.path.exists(fs_mountpoint):
        os.rmdir(fs_mountpoint)


# Gets the verison of ONIE installed on the system
def get_onie_version():
        version = None

        try:
            fs_mountpoint = mount_onie_fs()
            if fs_mountpoint is not None:
                machine_conf_path = os.path.join(fs_mountpoint, 'onie/grub/grub-machine.cfg')

                with open(machine_conf_path, 'r') as machine_conf:
                    for line in machine_conf:
                        if line.startswith('onie_version'):
                            items = line.rstrip('\n').split('=')
                            if len(items) == 2:
                                version = items[1]
                                break
        finally:
            umount_onie_fs()

        return version


# Parse the ONIE version string into a 6-tuple
def parse_onie_version(version, is_base=False):
    onie_year = None
    onie_month = None
    onie_major = None
    onie_minor = None
    onie_release = None

    if is_base:
        # The base ONIE version string looks like: 5.2.0016   
        pattern = '([0-9]+)\.([0-9]+)\.([0-9]+)'
        m = re.search(pattern, version)
        if not m:
            return (None, None, None, None, None)

        onie_major = m.group(1)
        onie_minor = m.group(2)
        onie_release = m.group(3)
    else:
        # Structure of a typical ONIE version string: 2021.02-5.3.0006*
        pattern = '([0-9]{4})\.([0-9]{2})-([0-9]+)\.([0-9]+)\.([0-9]+).*'
        m = re.search(pattern, version)
        if not m:
            return (None, None, None, None, None)

        onie_year = m.group(1)
        onie_month = m.group(2)
        onie_major = m.group(3)
        onie_minor = m.group(4)
        onie_release = m.group(5)

    return (onie_year, onie_month, onie_major, onie_minor, onie_release)


# There is a min version of ONIE needed to upgrade the BIOS.  This function
# enforces that requirement.
def onie_version():

    # Get the version of ONIE installed on this system
    installed_onie_version = get_onie_version()
    if installed_onie_version is None:
        if verbose:
            print('Unable to determine ONIE version')
        return False, None, None
    (_, _, installed_major, installed_minor, installed_release)  = parse_onie_version(installed_onie_version)
    installed_version = int("{}{}{}".format(installed_major, installed_minor, installed_release))

    # Get the min required version of ONIE needed to do a BIOS upgrade
    (_, _, min_major, min_minor, min_release)  = parse_onie_version(min_onie_version,True)
    min_version = int("{}{}{}".format(min_major, min_minor, min_release))

    if installed_version > min_version:
        return (True, installed_version, min_version)
    else:
        return (False, installed_version, min_version)


# Identifies the type of BIOS (AtomRangeley, IvyBridge, Broadwell, CoffeeLake) used by this
# system.  Returns the full path of the corresponding BIOS binary file.
def identify_bios():

    # If the hw-mgmt node /run/hw-management/config/cpu_type does not
    # exist, we give up immediately.
    if os.path.exists(cpu_type_path) == False:
        return None

    # Read the hw-mgmt node /run/hw-management/config/cpu_type to      
    # identify the type of CPU on this system.
    bios_golden_dir = None
    cab = False
    with open(cpu_type_path) as fp:
        val = fp.read().rsplit('\n')[0]
        cpu_type = (int)(val, 16)
        if cpu_type == atom_rangeley_cpu_type:
            bios_golden_dir = atom_rangeley_bios_dir 
        elif cpu_type == ivy_bridge_cpu_type:
            bios_golden_dir = ivy_bridge_bios_dir 
        elif cpu_type == broadwell_cpu_type:
            bios_golden_dir = broadwell_bios_dir 
        elif cpu_type == coffee_lake_cpu_type:
            bios_golden_dir = coffee_lake_bios_dir 
            cab = True

    return (cab, bios_golden_dir)

# Returns the version info of the BIOS that is currently installed
def system_bios_version():
    version = None
    cmd = ['dmidecode', '-s', 'bios-version']
    version = subprocess.check_output(cmd)

    release_date = None
    cmd = ['dmidecode', '-s', 'bios-release-date']
    release_date = subprocess.check_output(cmd)

    subversion = None
    cmd = ['dmidecode', '--oem-string', '1']
    subversion = subprocess.check_output(cmd)

    return(version.rstrip(), release_date.rstrip(), subversion.rstrip())

# Extracts the BIOS TGZ files to /tmp and returns the full path
# of the BIOS ROM upgrade file.
def extract_bios(cab, bios_golden_dir):
    bios_path = None
    bios_name = None

    if cab:
        suffix = '.cab'
    else:
        suffix = '.tgz'

    # Locate the BIOS golden file(s)
    bios_golden_count = 0
    for f in os.listdir(bios_golden_dir):
        if f.endswith(suffix):
            bios_golden_file = f
            bios_golden_count += 1
 
    # If we don't find any BIOS tgz files; give up. 
    if bios_golden_count == 0:
        if verbose:
            print(('No BIOS files found in %s' % (bios_golden_dir)))
        return (None, None)

    # If we find multiple BIOS tgz files; give up.
    # (we're not going to try and figure out which file is the "right"
    # one to use).
    if bios_golden_count != 1:
        if verbose:
            print(('Multiple BIOS files found in %s' % (bios_golden_dir)))
        return (None, None)

    # Now that we've identified a single BIOS file, extract it to /tmp       
    bios_path = bios_golden_dir + bios_golden_file

    # Create a directory in /tmp if it does not already exist
    if os.path.isdir(tmp_bios_dir) == False:
        os.mkdir(tmp_bios_dir)
    if os.path.isdir(tmp_bios_dir) == False:
        if verbose:
            print(('Directory %s does not exist' % (tmp_bios_dir)))
        return (None, None)

    # Erase the contents of this directory 
    cmd = 'rm -rf' + ' ' + tmp_bios_dir + '*'
    os.system(cmd)

    # Extract the BIOS files to this directory
    if cab:
        cmd = 'cp' + ' ' + bios_path + ' ' + tmp_bios_dir
    else:
        cmd = 'tar -xvf' + ' ' + bios_path + ' ' + '-C' + ' ' + tmp_bios_dir + ' ' + '> /dev/null'
    os.system(cmd)

    # Locate the BIOS upgrade file(s)
    if cab:
        suffix = '.cab'
    else:
        suffix = '.rom'
    bios_upgrade_count = 0
    for f in os.listdir(tmp_bios_dir):
        if f.endswith(suffix):
            bios_upgrade_file = f
            bios_upgrade_count += 1

    # If we don't find any BIOS upgrade files; give up. 
    if bios_upgrade_count == 0:
        if verbose:
            print(('No BIOS files found in %s' % (tmp_bios_dir)))
        return (None, None)

    # If we find multiple BIOS upgrade files; give up.
    # (we're not going to try and figure out which file is the "right"
    # one to use).
    if bios_upgrade_count != 1:
        if verbose:
            print(('Multiple BIOS files found in %s' % (tmp_bios_dir)))
        return (None, None)

    bios_upgrade_path = tmp_bios_dir + bios_upgrade_file
    return (bios_upgrade_path, bios_upgrade_file)


# Uses an ONIE utility to display pending and past firmware updates.
# This information is also logged to a file in /tmp/BIOS.
def show_pending_upgrades():
    try:
        fs_mountpoint = mount_onie_fs()
        if fs_mountpoint is not None:
            onie_fwpkg_path = os.path.join(fs_mountpoint, onie_fwpkg)
            if os.path.exists(onie_fwpkg_path):
                cmd = onie_fwpkg_path + ' > ' + onie_fw_updates_file
                os.system(cmd)
    finally:
        umount_onie_fs()

    cmd = 'cat' + ' ' + onie_fw_updates_file
    os.system(cmd)


# Reverts a BIOS upgrade by unstaging the BIOS upgrade in 
# ONIE and cancelling the pending ONIE reboot operation.
def revert_bios_upgrade(bios_path):
    rc = True

    # Unstage the BIOS upgrade in ONIE
    try:
        fs_mountpoint = mount_onie_fs(rw=True)
        if fs_mountpoint is not None:
            onie_fwpkg_path = os.path.join(fs_mountpoint, onie_fwpkg)
            if os.path.exists(onie_fwpkg_path):
                cmd = onie_fwpkg_path + ' remove ' + bios_path
                os.system(cmd)
        else:
            rc = False
    finally:
        umount_onie_fs()

    # Cancel the pending ONIE reboot operation (update mode)
    cmd = 'onie-select -f -c'
    os.system(cmd)

    return rc


# Uses an ONIE utility to stage the BIOS upgrade.  This does not 
# upgrade the BIOS; it simply stages it for (later) installation.
def stage_bios_upgrade(bios_path):
    rc = True

    try:
        fs_mountpoint = mount_onie_fs(rw=True)
        if fs_mountpoint is not None:
            onie_fwpkg_path = os.path.join(fs_mountpoint, onie_fwpkg)
            if os.path.exists(onie_fwpkg_path):
                cmd = onie_fwpkg_path + ' add ' + bios_path
                os.system(cmd)
        else:
            rc = False
    finally:
        umount_onie_fs()

    return rc   


# Uses an ONIE utility to configure the system to enter ONIE update
# mode after the next reboot.  When the system enters ONIE update mode,
# it will locate the staged BIOS binary from above and perform the BIOS
# upgrade.
def onie_update_mode():
    cmd = 'onie-select -f -p'
    os.system(cmd)
    return True


# 'Main' part of the script starts here.

# Decipher command line arguements
parser = argparse.ArgumentParser()
parser.add_argument("--verbose", help="verbose mode", action="store_true")
parser.add_argument("--check", help="does not alter BIOS", action="store_true")
parser.add_argument("--revert", help="revert  BIOS upgrade", action="store_true")
parser.add_argument("--force", help="force BIOS upgrade", action="store_true")
args = parser.parse_args()
if args.verbose:
    verbose = 1
if args.check:
    check = 1
if args.force:
    force = 1
if args.revert:
    revert = 1
if verbose:
    print(("Command line arguements: verbose=%d check=%d revert=%d force=%d" % (verbose,check,revert,force)))

# Root access check
result_str = 'Root check: '
root = is_root()
if root:
    result_str += 'PASS'
else:
    result_str += 'FAIL'
    sys.exit(result_str)
print(result_str)

# Tools check
result_str = 'Tools check: '
tools = tools_check()
if tools:
    result_str += 'PASS'
else:
    result_str += 'FAIL'
    sys.exit(result_str)
print(result_str)

# Minimum ONIE version check
result_str = 'ONIE version check: '
version_check, installed_version, min_version = onie_version()
if version_check:
    result_str += 'PASS'
else:
    result_str += 'FAIL'
    sys.exit(result_str)
print(result_str)
if verbose:
    print(('  ONIE min version: %d' % (min_version)))
    print(('  ONIE installed version: %d' % (installed_version)))

# Get the BIOS version currently installed on the switch
result_str = 'Current BIOS version check: '
(version, release_date, subversion) = system_bios_version()
if (version is not None) and (release_date is not None) and (subversion is not None):
    result_str += 'PASS'
else:
    result_str += 'FAIL'
    sys.exit(result_str)
print(result_str)
if verbose:
    print(('  Version: %s' % (version)))
    print(('  Release Date: %s' % (release_date)))
    print(('  Subversion: %s' % (subversion)))

# Identify BIOS type (AtomRangeley, IvyBridge, Broadwell, Coffeelake, etc)
result_str = 'BIOS identification: '
(cab, bios_golden_dir) = identify_bios()
if bios_golden_dir is not None:
    result_str += 'PASS'
else:
    result_str += 'FAIL'
    sys.exit(result_str)
print(result_str)
if verbose:
    print(('  bios_golden_dir = %s' % (bios_golden_dir)))

# Extract the BIOS to /tmp/BIOS
result_str = 'BIOS extraction: '
(bios_upgrade_path, bios_upgrade_file) = extract_bios(cab, bios_golden_dir)
if bios_upgrade_file is not None:
    result_str += 'PASS'
else:
    result_str += 'FAIL'
    sys.exit(result_str)
print(result_str)
if verbose:
    print((' bios_upgrade_file = %s' % (bios_upgrade_path)))

# Revert BIOS upgrade (if desired)
if revert:
    result_str = 'BIOS revert: '
    revert = revert_bios_upgrade(bios_upgrade_file)
    if revert:
        result_str += 'PASS'
    else:
        result_str += 'FAIL'
    sys.exit(result_str)

# Stage the BIOS upgrade
if check == True:
     sys.exit('BIOS staging: SKIPPED')
else:
    # Ask the user if he/she is SURE they want to upgrade the BIOS
    if force == False:
        user_input = input("WARNING: This script upgrades the BIOS.  Are you SURE you want to do this? [yes | no]: ")
        user_input = user_input.lower()
        if user_input != 'yes':
            sys.exit('BIOS staging: SKIPPED')

    result_str = 'BIOS staging: '
    stage = stage_bios_upgrade(bios_upgrade_path)
    if stage:
        result_str += 'PASS'
    else:
        result_str += 'FAIL'
        sys.exit(result_str)
    print(result_str)

    # Configure the system to enter ONIE upgrade mode when
    # the system is rebooted.  ONIE will then locate the
    # staged BIOS upgrade, perform the upgrade, and will then
    # reboot the system again.
    onie_update_mode()
    
show_pending_upgrades()
