#!/usr/bin/python

# dell staging fw updater script

import os
import sys
import subprocess
import argparse


onie_boot_folder  = '/mnt/onie-boot/onie/tools/bin/onie-fwpkg'
onie_fwpkg_tool  = '/mnt/onie-boot/onie/tools/bin/onie-fwpkg'
ONIE_BOOT_MODE_CMD = '/mnt/onie-boot/onie/tools/bin/onie-boot-mode'
HOST_GRUB_DIR = '/host'
HOST_GRUB_CFG = HOST_GRUB_DIR + '/grub/grub.cfg'
HOST_GRUB_ENV = HOST_GRUB_DIR + '/grub/grubenv'
HOST_GRUB_BOOT_DIR = '--boot-directory=' + HOST_GRUB_DIR
HOST_PLATFORM_INFO = HOST_GRUB_DIR + '/platform'
dell_reload_tool = '/usr/bin/reboot'




def set_onie_mode(option):
    """Select the ONIE boot mode, and set the next_entry to point to ONIE"""
    _set_env_option('next_entry', 'ONIE')
    subprocess.check_call([ONIE_BOOT_MODE_CMD, '-o', option])

def set_onie_fw_update_env():
    """Select the ONIE boot mode, and set the next_entry to point to ONIE"""

    if not os.path.exists(onie_boot_folder):
            os.makedirs(onie_boot_folder)

    try:
        subprocess.check_call(['mount','/dev/disk/by-label/ONIE-BOOT','/mnt/onie-boot'])
    except:
        print "onie-boot not able to mount"

def _set_env_option(option, value):
    """Set an option in the GRUB environment block. Pass None to value to
    unset the option"""
    if value is None:
        action = 'unset'
        key_value = option
    else:
        action = 'set'
        key_value = '%s=%s' % (option, value)

    subprocess.check_call(['grub-editenv', HOST_GRUB_ENV, action, key_value])


def dell_firmware_update_staging(image_name):

    try:
        p = subprocess.Popen([onie_fwpkg_tool,"purge"],stdout=subprocess.PIPE,stdin=subprocess.PIPE)
        p.communicate("y")
    except:
        print "onie-fwpkg command not found for purging old fw updates"

    try:
        subprocess.check_call([onie_fwpkg_tool,"add", str(image_name)])
    except:
        print "onie-fwpkg is not found to stage fw updates"

    try:
        set_onie_mode("update")
    except:
        print "dell-image command not found"

    try:
        subprocess.check_call([dell_reload_tool])
    except:
        print "reload command not found"


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Dell HOST Firmware updates')
    opts = parser.add_mutually_exclusive_group(required=True)
    opts.add_argument('-u', '--update', nargs=1, metavar='IMAGE',
                       help='update specified image')

    args = parser.parse_args()

    if os.getuid() != 0:
        parser.exit(127, 'ERROR: Must be root\n')

    if args.update:
        set_onie_fw_update_env()
        dell_firmware_update_staging(args.update[0])