flashbios

#!/usr/bin/env python3
"""Protectli device BIOS flasher.

This tool will flash new BIOS onto the machine that runs this script.
"""

from ast import Constant
from asyncio import constants
import os
import subprocess  # noqa:S404
import sys
import textwrap
import time

from flashli import configurations, hardware

# Set a debug hardware here.
DEBUGMODE = ''

CONFIGURATIONS = configurations.CONFIGURATIONS

displaySize = 89


VERSION = '1.1.50'


if os.geteuid() != 0 and not DEBUGMODE:
    print('Need to be run as root user')
    print('Please run: sudo ./flashbios')
    sys.exit()


def get_version() -> str:
    """Functional way to get global VERSION.

    Returns:
        str: Value of global VERSION
    """
    global VERSION
    return VERSION


def get_terminal_width() -> int:
    """Get the width of the current terminal in characters.

    Returns:
        int: Width in characters
    """
    return os.get_terminal_size().columns


def get_image_path(model: str, requested_bios: str) -> str:
    """Get path to BIOS image.

    Args:
        model: Protectli device model name
        requested_bios: The BIOS to be located

    Returns:
        str: Path to the BIOS file requested
    """
    global CONFIGURATIONS
    for bios in CONFIGURATIONS[model]['bios']:
        if bios['vendor'] == requested_bios:
            return 'images/{0}'.format(bios['file'])
        

def check_vp46_flash_config (bios_select: str) -> str:
    """checks the current BIOS to user selected flash image.

    Args:
        bios_select: selected flash BIOS image

    Returns:
        str: Path to the BIOS file requested
    """
    full_flash_exceptions = ('Dasharo (coreboot+UEFI) v1.0.1', 'Dasharo (coreboot+UEFI) v1.0.19')
    full_flash_unit_exceptions = {'vp2420', 'vp4670'}
    full_flash = {'v1210', 'v1211', 'v1410', 'v1610',
                  'vp2410r', 'vp2420', 'vp4630', 'vp4650', 'vp4670', 'vp4670[s1]',
                  'vp6630', 'vp6650', 'vp6670'}
    
    vp_devices = {'vp2410r', 'vp2420', 'vp4630', 'vp4650', 'vp4670', 'vp4670[s1]',
                  'vp6630', 'vp6650', 'vp6670'}


    #Checks if the device is a VP
    if hardware.get_protectli_device(DEBUGMODE, check_mac()) in vp_devices :

        #Checks if current BIOS/UEFI is AMI
        if hardware.check_for_ami(DEBUGMODE) :

            #Checks if the BIOS lock is enabled
            if hardware.check_bios_lock(DEBUGMODE) :
                display_bios_lock_ami()

            else:
                return 'command'
            
        #Checks if the current BIOS is coreboot
        elif hardware.check_for_coreboot(DEBUGMODE) :


            #Checks if ami image is select to flash
            if 'ami' in bios_select:

                #Checks if the BIOS lock is on
                if hardware.check_bios_lock(DEBUGMODE) :
                    display_bios_lock_coreboot()

                else :
                    return 'command'
            #Checks if there is an exception    
            elif hardware.has_param(DEBUGMODE, full_flash_exceptions):

                #Checks if the BIOS lock is on
                if hardware.check_bios_lock(DEBUGMODE) :
                    display_bios_lock_coreboot()

                elif hardware.get_protectli_device(DEBUGMODE, check_mac()) in full_flash_unit_exceptions :
                    print ('Alternative Upgrade')
                    return 'alt_upgrade'

                else :
                    print ('Regular Flashing')
                    return 'command'
                
            elif hardware.get_protectli_device(DEBUGMODE, check_mac()) in full_flash:

                #Checks if the BIOS lock is on
                if hardware.check_bios_lock(DEBUGMODE) :
                    display_bios_lock_coreboot()
                else:
                    print('Overwrite')
                    return 'command'

            
            else:
                print ('Upgrading')
                return 'upgrade'
        
        else:
            sys.exit('Error has occured\nVP46xx device not found')

def do_flash(model: str, bios: str) -> int:
    """Perform the BIOS flash.

    Args:
        model: Protectli device model name
        bios: The BIOS to be flashed

    Returns:
        Returns exit status from flashrom binary:
            0: Success
            1: General Failure
            2: /dev/mem cannot be opened
            3: mmap() failed
    """
    global DEBUGMODE
    global CONFIGURATIONS

    vp_device = {'vp2410r', 'vp2420', 'VP46xx', 'vp4630', 'vp4650', 'vp4670'}

    if DEBUGMODE:
        print('Not actually flashing, script is in debug mode.')
        return 0

    vp46_command = check_vp46_flash_config(bios)
    file_path = get_image_path(model, bios)

    if hardware.has_param(DEBUGMODE, 'FW6D') or hardware.has_param(DEBUGMODE, 'FW6E') and hardware.has_param(DEBUGMODE, 'coreboot'):

        completed_process = subprocess.run(CONFIGURATIONS[model]['override'].format(file_path), shell=True)  # noqa:S602,S603

    elif model in vp_device : 

        completed_process = subprocess.run(CONFIGURATIONS[model][vp46_command].format(file_path), shell=True)  # noqa:S602,S603

    else:

        completed_process = subprocess.run(CONFIGURATIONS[model]['command'].format(file_path), shell=True)  # noqa:S602,S603

    return completed_process.returncode

def check_mac() -> bool:

    global DEBUGMODE

    str_mac = int("0x646266210000", base=16)
    end_mac = int("0x646266210314", base=16)
    

    device_mac = hardware.get_mac(DEBUGMODE)

    if device_mac :
        device_mac = int("0x" + str(device_mac).replace(':', ''), base=16)
        mac_match = ''

    else :
        print("Issue obtaining MAC address")
        print("Exiting program")
        sys.exit(1)

    if  device_mac >= str_mac and device_mac <= end_mac:
        mac_match = 'vp_vr1'
    
    elif "J4125" in hardware.get_cpu(DEBUGMODE):
        mac_match = 'vp_vr2'
    
    return mac_match

def print_supported_products():
    """Get list of devices from Configurations and prints to STDOUT."""
    global CONFIGURATIONS
    devices = list(map(lambda dev: dev.upper(), list(CONFIGURATIONS)))
    devices.sort()
    print(*devices, sep='\n')


def get_user_selection(device: str) -> str:
    """Get BIOS selection from user based on available images for device.

    Args:
        device: Which device to ask the user about

    Returns:
        str: Vendor name
    """
    global CONFIGURATIONS
    available_options = CONFIGURATIONS[device]['bios']
    while True:
        number = 1
        for available_option in available_options:
            vendor = available_option['vendor']
            filename = available_option['file']
            print('[{0}]: {1} ({2})'.format(str(number), vendor, filename))
            number += 1
        print('\nEnter the [#] of an image file, or [0] to quit. Flashing will not begin yet')
        print('> ', end='')
        try:
            user_input = int(input())
        except Exception:
            user_input = -1
        if (user_input == 0):
            sys.exit('Exiting now.')
        elif (0 < user_input <= len(available_options)):
            return available_options[user_input - 1]['vendor']
        else:
            print('Invalid choice.')
            print('Available BIOS:\n')


def show_debug_info():
    """Collect debug information and tell user how to submit an issue."""
    print('TODO: Collect info and display instructions on how to submit a Github issue.')

def display_logo () :

    logo_str =  ' ______   ______     ______     ______   ______     ______     ______   __         __    \n'
    logo_str += '/\  == \ /\  == \   /\  __ \   /\__  _\ /\  ___\   /\  ___\   /\__  _\ /\ \       /\ \   \n'
    logo_str += '\ \  _-/ \ \  __<   \ \ \/\ \  \/_/\ \/ \ \  __\   \ \ \____  \/_/\ \/ \ \ \____  \ \ \  \n'
    logo_str += ' \ \_\    \ \_\ \_\  \ \_____\    \ \_\  \ \_____\  \ \_____\    \ \_\  \ \_____\  \ \_\ \n'
    logo_str += '  \/_/     \/_/ /_/   \/_____/     \/_/   \/_____/   \/_____/     \/_/   \/_____/   \/_/ \n'
    logo_str += '_________________________________________________________________________________________\n\n'

    return logo_str


def general_notice(device) :

    revision = ['v1410']

    if device in revision:
        print('Revised PKfail BIOS for {0} is available'.format(device))
        print('please see https://kb.protectli.com/kb/bios-versions-for-the-vault\nfor changes')


def display_warning () :

    display_logo()

    general_warning = '\nFlashing new firmware onto any hardware is potentially dangerous\nin that if the procedure is interrupted or otherwise not able to complete, your hardware may be rendered useless.\nPlease proceed with caution.\nIf there are any questions, please contact Protectli support BEFORE proceeding.\n\n'

    ram_warning = '\nUnless there is a compelling reason to update the BIOS, we recommend to stay with your current known working BIOS version'

    print('\n' + '!'.center(displaySize, '*')) 
    print (textwrap.fill('Are you sure you would like to flash this device?', displaySize) + '\n')
    print (textwrap.fill(general_warning, displaySize) + '\n')
    print (textwrap.fill(ram_warning, displaySize) + '\n')
    print('!'.center(displaySize, '*'))

    user_agrement = str(input('\nAcknowledgement Yes [Y]: ')).lower()

    if user_agrement == 'y':
        print ('')
    
    else :
        sys.exit('**Acknowledgement rejected**\n')

def display_bios_lock_ami ():
        
    print('\nBIOS lock option in the BIOS is enabled.\nFlashli will now exit to prevent harm.')
    print('Please disable the BIOS lock in the BIOS\n')

    print('\tUSING AMI\n')
    print('-Please reboot the device\n-Press DEL when the splash screen appears')
    print('-Navigate to Chipset\n-Navigate to PCH-10 Configuration')
    print('\n-You will see BIOS Lock at the bottom')
    print('-Disable BIOS lock\n-Press F4 to save and reboot')
    print('-After reboot please try flashli again\n\n')

    sys.exit()

def display_bios_lock_coreboot ():

    print('\nBIOS lock option in the BIOS is enabled.\nFlashli will now exit to prevent harm.')
    print('Please disable the BIOS lock in the BIOS\n')

    print('\tUSING COREBOOT\n')
    print('-Please reboot the device\n-Press DEL when the splash screen appears')
    print('-From the main menu navigate to \'Dasharo System Features\'\n-Select \'Dasharo Security Options\'')
    print('-Disable Lock the BIOS boot medium and SMM BIOS write protection by pressing enter')
    print('-Press ESC until you see the main menu\n-Select Reset to save and reboot.')
    print('-After reboot please try flashli again\n\n')

    sys.exit()

def display_secureboot_enabled_ami ():
        
    print('\nSecure Boot option in the BIOS is enabled.\nFlashli will now exit to prevent harm.')
    print('Please disable Secure Boot in the BIOS\n')

    print('\tUSING AMI\n')
    print('-Please reboot the device\n-Press DEL when the splash screen appears')
    print('-Navigate to Security\n-Navigate Secure Boot')
    print('\n-You will see Secure Boot as the first option')
    print('-Disable Secure Boot\n-Press F4 to save and reboot')
    print('-After reboot please try flashli again\n\n')

    sys.exit()

def display_secureboot_enabled_coreboot ():

    print('\nSecure Boot option in the BIOS is enabled.\nFlashli will now exit to prevent harm.')
    print('Please disable Secure Boot in the BIOS\n')

    print('\tUSING COREBOOT\n')
    print('-Please reboot the device\n-Press DEL when the splash screen appears')
    print('-From the main menu navigate to \'Device Manager\'\n-Select \'Secure Boot Configureation\'')
    print('-Disable Secure Boot by pressing enter')
    print('-Press ESC until you see the main menu\n-Select Reset to save and reboot.')
    print('-After reboot please try flashli again\n\n')

    sys.exit()

def display_long_boot_time (device):

    long_boot = ['vp6630', 'vp6650', 'vp6670']

    if device in long_boot:

        print('\ncoreboot on the VP6600 has a notably longer memory initialization process')

        print('After the unit has completed flashing the coreboot BIOS, the device may take up to 2 minutes to boot.')

        print('DO NOT unplug the device during this time.')

        print('Doing so will reset the initialization process and may cause a significant delay in device initialization.')


def display_remove_power_warning (device):

    poweroff_devices = ['v1210', 'v1211', 'v1410', 'v1610', 'vp2420', 'vp6630', 'vp6650', 'vp6670']

    user_input = '0'

    if device in poweroff_devices:

        while (user_input != 'Y'):

            print('When flashing new firmware onto the {0}, you MUST disconnect and reconnect the power cable\nfrom the unit in order to fully complete the update process.'.format(device))

            print('Failure to perform these steps will cause your {0} to fail to power on or behave erratic until rectified'.format(device))

            print('\n\nAfter you have completed flashing the firmware:')

            print('\n1. Shut down the {0}.'.format(device))
            print('2. Disconnect the power cable from the unit.')
            print('3. Wait five seconds.')
            print('4. Reconnect the power cable to the unit.')

            print('\nYOUR {0} WILL BE IN AN UNUSABLE STATE UNTIL THE ABOVE STEPS ARE PERFORMED.'.format(device))

            display_long_boot_time(device)

            print('\nPress [Y] to acknowledge these steps, or press [N] to stop the firmware update process.')

            user_input = input(':').upper()

            if (user_input == 'Y'):
                print('\nAcknowledgement confirmed')
                print('proceeding')

            elif (user_input == 'N'):
                print('\nFlashing firmware process has been stopped')
                print('Exiting program')
                sys.exit(0)
            else:
                print('\nPlease enter Y for yes or N for no')


def display_pkfail_support (device):

    not_supported = ['fw1', 'fw2', 'fw4a', 'fw6m']

    if device in not_supported:

        print('\nThe {0} is out of the support window and will not be getting further updates.'.format(device))
        print('{0} will not be receiving the PKfail fix.\n'.format(device))


def display_serial_uuid_warning (device):

    default_serial_uuid = ['v1210', 'v1211', 'v1410', 'v1610']

    if device in default_serial_uuid:

        print('\nWith the current implementation of firmware flashing, {0}\'s serial and UUID will be set to default.'.format(device))
        print('A solution is currently in development.\n')


def main():  # noqa:WPS213
    """Main program."""

    os.system('/bin/clear')  # noqa:S607,S605

    print (display_logo())

    global DEBUGMODE
    device = hardware.get_protectli_device(DEBUGMODE, check_mac())

    print('NOTICE:\n')
    general_notice(device)
    display_pkfail_support(device)
    display_serial_uuid_warning(device)
    display_remove_power_warning(device)

    if hardware.check_secureboot(DEBUGMODE):

        if hardware.check_for_ami(DEBUGMODE):
            display_secureboot_enabled_ami()

        elif hardware.check_for_coreboot(DEBUGMODE):
            display_secureboot_enabled_coreboot()

    if device == 'Unknown' :
        print ('device is returning unknown')
        return -1
    cpu = hardware.get_cpu(DEBUGMODE)
    print('FlashLi'.center(displaySize, '='))
    print('--Version {0}--\n'.format(get_version()).center(displaySize, ' '))

    print('Device:    Protectli {0}'.format(device))
    print('CPU:       {0}'.format(cpu))
    if not hardware.is_protectli_device(DEBUGMODE) or 'Unknown' == hardware.is_protectli_device(DEBUGMODE):
        print('Sorry, this is an unsupported device.')
        print('This tool is used to flash BIOS onto the following Protectli products:')
        print_supported_products()
        sys.exit()

    bios_mode = hardware.get_bios_mode(DEBUGMODE)
    print('BIOS Mode: {0}'.format(bios_mode))

    efi_devices = {'fw4c', 'v1210', 'v1211', 'v1410', 'v1610', 'vp2410r', 'vp2420', 'vp2410r', 'vp4630', 'vp4650', 'vp4670', 'vp4670[s1]',
                   'vp6630', 'vp6650', 'vp6670'}

    if bios_mode == 'EFI' and not str(device) in efi_devices :

        print(textwrap.fill(
            '\n\nThis tool must be run in Legacy BIOS mode, not EFI.\n'
            'If you are using this tool to update an existing EFI system,\n'
            'you may experience issues booting into your operating system after flashing a new BIOS.\n'
            'If you are aware of the risks and wish to proceed with flashing a new BIOS image,\n'
            ' please reboot your device and configure your current BIOS to boot into Legacy Mode.',
            displaySize,
        ))
        sys.exit()

    print()
    print('Available BIOS:')
    print()

    selection = get_user_selection(device)
    
    display_warning()

    returncode = do_flash(device, selection)
    if returncode == 0:
        print('Flash completed and successful.')

        if device in ['v1210', 'v1211', 'v1410', 'v1610', 'vp2420', 'vp6630', 'vp6650', 'vp6670']:
            print('Please REMOVE the POWER SUPPLY from the device and reinsert it')
        
        else:
            print('Please restart your device.')

    else:
        print('BIOS Flash failed, is this script running with root permissions?')
        print('Please try again, but if problems persist, please contact Protectli.')
        show_debug_info()


try:
    main()
except KeyboardInterrupt:
    print('\n')
    sys.exit()