#!/usr/bin/python

###############################################################################
#
# upgrade_ssd_fw - Updates the firmware on select SSD drives to correct 
# failures in specific cases.
#
#  Copyright (C) 2020 Cumulus Networks, Inc. All rights reserved
#
#  This software is subject to the Cumulus Networks End User License Agreement
#  available  at the following locations:
#
#  Internet: https://cumulusnetworks.com/downloads/eula/latest/view/
#  Cumulus Linux systems: /usr/share/cumulus/EULA.txt
#
# Description:
#
#
# Syntax:
#      upgrade_ssd_fw [--verbose | --check]
#
#      --verbose    : Verbose mode
#      --check      : Do everything except the firmware upgrade         
#
###############################################################################


from glob import glob
from os.path import basename, dirname
import os.path
import subprocess
import string
import sys
import argparse
import syslog

#
# Global variables and "constants"
#

ssd_health_path  = '/usr/cumulus/bin/ssd_health_cumulus/ssd_health.sh'
ssd_upgrade_path = '/usr/cumulus/bin/mlnx_ssd_fw_update.sh'
ssd_fw_path      = '/usr/lib/cumulus/firmware/mlnx_ssd_fw_package.pkg'


#
# Command line switches
#
verbose    = 1
check      = 0


def is_mlnx():
    mlnx = False

    if verbose:
        print ("Detecting platform...")
    cmd = ['platform-detect']
    try:
        output = subprocess.check_output(cmd)
    except subprocess.CalledProcessError:
        return mlnx

    switch_vendor = ''
    for row in output.split('\n'):
        fields = row.split(',')
        if len(fields) >= 1:
            switch_vendor = fields[0]
        if (switch_vendor == 'mlnx') or (switch_vendor == 'mlx'):
            mlnx = True

    if verbose:
        print ("  mlnx: %d" % (mlnx))
    return mlnx


def check_ssd_fw_upgrade():
    upgrade = False

    if verbose:
        print ("Checking if ssd fw upgrade is needed...")
    cmd = [ssd_health_path, '--mode', 'manual']
    try:
        output = subprocess.check_output(cmd)
    except subprocess.CalledProcessError:
        return upgrade

    for row in output.split('\n'):
        if row.startswith("Recommendation: Upgrade SSD FW"):
            upgrade = True
    
    if verbose:
        print ("  upgrade: %d" % (upgrade))
    return upgrade


def do_ssd_fw_upgrade():
    success = False

    # Sync the file system
    if verbose:
        print ("Syncing file system...")
    cmd = ['sync']
    subprocess.call(cmd)

    if check:
        return

    # Upgrade the SSD firmware.  System will automatically reboot when complete.
    if verbose:
        print ("Upgrading ssd firmware...")
    cmd = [ssd_upgrade_path, '-u', '-y', '-i', ssd_fw_path]
    try:
        output = subprocess.check_output(cmd)
    except subprocess.CalledProcessError:
        msg = "SSD firmware upgrade ERROR"
        syslog.syslog(syslog.LOG_ERR,msg)
        print(msg)
        return

    for row in output.split('\n'):
        if row.startswith("SSD FW update completed successfully"):
            success = True    
            msg = "SSD firmware upgrade successful"
            syslog.syslog(syslog.LOG_NOTICE,msg)
            print(msg)

    if not success:
        msg = "SSD firmware upgrade FAILED!"
        syslog.syslog(syslog.LOG_CRIT,msg)
        print(msg)


# 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 firmware", action="store_true")
args = parser.parse_args()
if args.verbose:   
    verbose = 1
if args.check:   
    check = 1
if verbose:
    print("Command line arguements: verbose=%d check=%d" % (verbose,check))

if is_mlnx():
    if check_ssd_fw_upgrade():
            do_ssd_fw_upgrade()
