#!/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 | --downgrade | --swap]
#
#      --verbose    : Verbose mode
#      --check      : Do everything except the firmware upgrade         
#      --downgrade  : Downgrades the firmware to the original version (DEBUG ONLY)
#      --swap       : Swaps firmware image (DEBUG ONLY)
#
###############################################################################


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"
#

hdparm_cmd      = ['hdparm', '--yes-i-know-what-i-am-doing', '--please-destroy-my-drive', '--fwdownload']
ssd_fw_path     = '/usr/lib/cumulus/firmware/'

virtium_model_name        = 'VSF302XC016G'
virtium_old_fw_version    = '0115-000'
virtium_new_fw_version    = '1210-000'
virtium_old_firmware_file = ssd_fw_path + 'VSF302XC016G-MLX_ISP_0115.bin'
virtium_new_firmware_file = ssd_fw_path + 'VSF302XC016G-MLX_ISP_1210.bin'

virtium_drive       = 1


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

# Check to make sure that the tools we'll use are available
def tool_check():
    tools = ['/usr/sbin/smartctl', '/sbin/hdparm']

    for tool in tools:
        if os.path.exists(tool) == False:
            print("The required %s program does not exist or can't be executed." % (tool))
            return False

    return True 


# Builds a list of the physical drives in the system
def physical_drives():
    drive_glob = '/sys/block/*/device'
    drives = [basename(dirname(d)) for d in glob(drive_glob)]
    drives = sorted(drives)
    return drives


# Gets the drive model, firmware version, and switch vendor 
def query_drive(drive):
    model = ''
    fw_version = ''
    mlx = False 

    drive_path = '/dev/' + drive
    cmd = ['smartctl', '-i', '-T', 'permissive', drive_path]
    try:
        output = subprocess.check_output(cmd)
    except subprocess.CalledProcessError:
        return model, fw_version, mlx 

    for row in output.split('\n'):
        fields = row.split(':')
        if len(fields) >= 2:
            if fields[0] == 'Device Model':
                model = fields[1]
                model = model.lstrip()
                model = model.rstrip()
            elif fields[0] == 'Firmware Version':
                fw_version = fields[1]
                fw_version = fw_version.lstrip()
                fw_version = fw_version.rstrip()

    cmd = ['platform-detect']
    try:
        output = subprocess.check_output(cmd)
    except subprocess.CalledProcessError:
        return model, fw_version, mlx

    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'):
            mlx = True

    if verbose:
        print ("%s model: %s firmware version: %s mlx: %d" % (drive_path,model,fw_version,mlx))

    return model, fw_version, mlx


# Determines if a drive is made by Virtium 
def is_virtium(drive_type):
    if (drive_type == virtium_drive):
        return True
    else:
        return False


# Determine if a drive's firmware needs to be upgraded  
def does_drive_need_fw_upgrade(drive):
    model,fw_version,mlx = query_drive(drive)

    if mlx:
        if string.find(model,virtium_model_name) >= 0:
            if (fw_version == virtium_old_fw_version):
                return virtium_drive
    
    return 0


# Determine if a drive's firmware needs to be downgraded  
def does_drive_need_fw_downgrade(drive):
    model,fw_version,mlx = query_drive(drive)

    if mlx:
        if string.find(model,virtium_model_name) >= 0:
            if (fw_version == virtium_new_fw_version):
                return virtium_drive
    
    return 0


# Check to make sure that the firmware binary is available
def firmware_bin_check(drive_type):
 
    if drive_type == 0:
        return ''
 
    firmware_file_name = ''
    if upgrade:
        if drive_type == virtium_drive:
            firmware_file_name = virtium_new_firmware_file
    else:
        if drive_type == virtium_drive:
            firmware_file_name = virtium_old_firmware_file
  
    if os.path.exists(firmware_file_name) == False:
        print("The required %s firmware binary does not exist." % (firmware_file_name))
        return ''

    return firmware_file_name


# Upgrades the firmware of an SSD drive
def program_ssd_fw(drive,drive_type,firmware_file_name,action):

    drive_path = '/dev/' + drive

    fw_cmd = hdparm_cmd + [firmware_file_name, drive_path]

    if verbose:
        print ' '.join(fw_cmd)
    if check:  
        return

    result = False
    output = subprocess.check_output(fw_cmd)
    for row in output.split('\n'):
        if string.find(row,'Done') >= 0:
            result = True
        if verbose:
            print row
    
    if result:
        msg = "%s firmware %s SUCCESS" % (drive_path,action) 
    else:
        msg = "%s firmware %s FAILED!!!" % (drive_path,action) 
    print(msg)
    syslog.syslog(syslog.LOG_NOTICE,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")
parser.add_argument("--downgrade", help="downgrades firmware (DEBUG ONLY)", action="store_true")
parser.add_argument("--swap", help="swaps firmware image (DEBUG ONLY)", action="store_true")
args = parser.parse_args()
if args.verbose:   
    verbose = 1
if args.check:   
    check = 1
if args.downgrade:
    upgrade = 0
if args.swap:
    swap = 1
if verbose:
    print("verbose=%d check=%d upgrade=%d swap=%d" % (verbose,check,upgrade,swap))

# Make sure we have all the utilities needed to perform the SSD firmware upgrade
if tool_check() == False:
    sys.exit()

drives = physical_drives()
for drive_name in drives:
    
    if swap:
        if does_drive_need_fw_upgrade(drive_name):
            upgrade = 1
        elif does_drive_need_fw_downgrade(drive_name):
            upgrade = 0

    action = ''
    if upgrade:
        drive_type = does_drive_need_fw_upgrade(drive_name)
        action = 'upgrade'
    else:
        drive_type = does_drive_need_fw_downgrade(drive_name)
        action = 'downgrade'

    drive_path = '/dev/' + drive_name
    if drive_type:
        msg = "%s needs firmware %s!" % (drive_path,action)
        print(msg)

        firmware_file_name = ''
        firmware_file_name = firmware_bin_check(drive_type)
        if not firmware_file_name:
            continue

        if check == False:
            program_ssd_fw(drive_name,drive_type,firmware_file_name,action)
    else:
        if verbose or check:
             print("%s does NOT need firmware %s!" % (drive_path,action))
