#!/usr/bin/python
#
# Copyright (c) 2019, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#

import argparse
import copy
import os
import struct
import subprocess
import sys
import textwrap

if sys.hexversion < 0x02070000:
  print >> sys.stderr, "Python 2.7 or newer is required."
  sys.exit(1)

def parse_args():
    parser = argparse.ArgumentParser(description=textwrap.dedent('''\
                                                 Perform a bootloader and/or kernel update on a target with OTA blobs
                                                 '''), epilog=__doc__, formatter_class=argparse.RawTextHelpFormatter
                                    )
    parser.add_argument('-p', '--print-all', dest='print_all', action='store_true', default=False,
                        help=textwrap.dedent('''\
                             print all entries in blob file even if
                             entry is not updatable on current system
                             ''')
                       )
    parser.add_argument('blob', type=argparse.FileType(mode='rb'), nargs='?',
                        help=textwrap.dedent('''\
                            file path to the blob file used for installation
                            (default "/opt/ota_package/bl_update_payload")
                            ''')
                       )
    parser.add_argument('-v', '--ver', dest='print_ver', action='store_true', default=False,
                        help=textwrap.dedent('''\
                             print current bootloader version
                             (no need argument for valid blob file)
                             ''')
                       )
    params = parser.parse_args()
    return params

def print_VER(arg):
    system_info_obj = system_info(arg)
    system_info_obj.print_current_ver()

def install_BUP(arg):
    if arg.blob is None:
        try:
            arg.blob=open("/opt/ota_package/bl_update_payload")
        except:
            sys.stderr.write("Error. Missing default blob /opt/ota_package/bl_update_payload\r\n")
            sys.stderr.write("Or last argument must be a path to a valid blob file.\r\n")
            sys.stderr.write("Exiting...\r\n")
            sys.exit(1)

    print 'BLOB PATH:'
    print os.path.realpath(arg.blob.name)
    print

    payload_obj = install_update_payload(arg)

    payload_obj.print_blob_header()
    print
    payload_obj.print_entry_table()
    print
    print "Beginning update."
    print
    payload_obj.install_binaries()
    print "Update successful."

    arg.blob.close()

class payload(object):
    def __init__(self, arg):
        super(payload, self).__init__(arg)
        self.magic = 'NVIDIA__BLOB__V2'
        self.version = 0x00020000
        self.header_packing = '=16sIIIIII'

class update_type_payload(payload):
    def __init__(self, arg):
        super(update_type_payload, self).__init__(arg)
        self.blob_type = 0
        self.entry_packing = '=40sIIII64s'
        self.header_accessory_packing = '=Q'
        self.header_name_tuple = ("magic", "version", "blob_size", "header_size", "entry_count", "type", "uncomp_blob_size", "accessory")
        self.entry_name_tuple = ("part_name", "offset", "part_size", "version", "op_mode", "tnspec", "updatable", "intra_part_offset")
        self.accessory_present = False

        self.blob_file = arg.blob

        self.blob_header_tuple = struct.unpack(self.header_packing, self.blob_file.read(struct.calcsize(self.header_packing)))
        self.blob_header_dict = dict(zip(self.header_name_tuple, self.blob_header_tuple))

        # Detect if optional accessory field (8 bytes) is present
        if self.blob_header_dict['header_size'] > struct.calcsize(self.header_packing):
            self.blob_header_dict['accessory'] = struct.unpack(self.header_accessory_packing, self.blob_file.read(struct.calcsize(self.header_accessory_packing)))[0]
            self.accessory_present = True

        if not self._valid():
            sys.stderr.write("Error. Invalid input blob file.\r\n" \
                             "      Input magic: " + self.blob_header_dict['magic'] + "\r\n" \
                             "   Expected magic: " + self.magic + "\r\n" \
                             "    Input version: " + format(self.blob_header_dict['version'], "#010x") + "\r\n" \
                             " Expected version: " + format(self.version, "#010x") + "\r\n" \
                             "       Input type: " + str(self.blob_header_dict['type']) + "\r\n" \
                             "    Expected type: " + str(self.blob_type) + " (type 1 BMP blobs are not currently supported)\r\n" \
                             "No changes have been made. Exiting...\r\n" \
                            )
            sys.exit(1)

    def _valid(self):
        expected_header_tuple = (self.magic, self.version, self.blob_type)
        input_header_tuple = self.blob_header_tuple[0:2] + (self.blob_header_tuple[5],)

        return input_header_tuple == expected_header_tuple

class system_info(object):
    # part_name : (offset, size)

    # jetson-nano-qspi*
    _spi_main_boot_partition_table = {
        'BCT'  : (0, 262144),
        'NVC'  : (262144, 196608),
        'PT'   : (458752, 65536),
        'NVC_R': (524288, 196608), # Used for optional redundant bootloader upgrades
        'PAD'  : (720896, 3407872),
        'VER'  : (4128768, 65536)
    }

    # jetson-tx1*
    _emmc_main_boot_partition_table = {
        'BCT' : (0, 4194304)
    }
    _emmc_sec_boot_partition_table = {
        'NVC' : (0, 4194304)
    }

    def __init__(self, arg):
        super(system_info, self).__init__()
        self.system_spec = self._get_system_prop("TNSPEC")
        self.system_op_mode = self._get_system_op_mode()
        self.system_boot_dev_info = self._get_system_boot_dev_info()
        self.system_boot_partition_info = self._get_system_boot_partition_info()

        self.system_gpt_partition_dir = "/dev/disk/by-partlabel/"
        self.system_gpt_partition_list = os.listdir(self.system_gpt_partition_dir)

    def _get_system_prop(self, sys_property):
        sys_property_value = ""

        try:
            nv_boot_control_conf_dir = "/etc/nv_boot_control.conf"
            with open(nv_boot_control_conf_dir, 'r') as nv_boot_control_conf:
                target = [line for line in nv_boot_control_conf if sys_property in line]
                try:
                    sys_property_value = target[0].split()[1]
                except IndexError:
                    sys.stderr.write("Warning. Cannot find " + sys_property + " in " + nv_boot_control_conf_dir + "\r\n")
        except IOError:
            sys.stderr.write("Error. Cannot open " + nv_boot_control_conf_dir + "\r\n" \
                             "No changes have been made. Exiting...\r\n" \
                            )
            sys.exit(1)

        return str(sys_property_value)

    def _get_bct_prop(self, bct_device, bct_base_offset, bct_property):
        bct_prop_packing = ''
        bct_property_value = 0

        if bct_property == "boot_block_size_log_2":
            bct_offset = 1332
            bct_prop_packing = 'I'
        elif bct_property == "boot_page_size_log_2":
            bct_offset = 1336
            bct_prop_packing = 'I'
        else:
            raise LookupError("Cannot retrieve unsupported BCT property: " + bct_property)

        try:
            bct_device.seek(bct_base_offset + bct_offset, os.SEEK_SET)
            bct_property_value = struct.unpack(bct_prop_packing, bct_device.read(struct.calcsize(bct_prop_packing)))[0]
        except:
            sys.stderr.write("Warning. Could not retrieve " + bct_property + " from BCT.\r\n" \
                             "Skipping boot partition updates.\r\n\r\n" \
                            )
            bct_property_value = 0

        return bct_property_value

    def _get_system_op_mode(self):
        try:
            tegra_prod_mode_dir = "/sys/module/tegra_fuse/parameters/tegra_prod_mode"
            with open(tegra_prod_mode_dir, 'r') as tegra_prod_mode:
                op_mode = int(tegra_prod_mode.read())
        except IOError:
            sys.stderr.write("Error. Cannot open " + tegra_prod_mode_dir + ".\r\n" \
                             "No changes have been made. Exiting...\r\n" \
                            )
            sys.exit(1)
        return op_mode

    def _get_system_boot_dev_info(self):
        boot_dev_properties = {
            'dev_path'     : "",
            'write_en_path': "",
            'block_size'   : 0,
            'page_size'    : 0
        }
        boot_dev_system_prop = {
            'main'     : "TEGRA_OTA_BOOT_DEVICE",
            'secondary': "TEGRA_OTA_GPT_DEVICE"
        }
        boot_dev_info = {
            'main'     : copy.copy(boot_dev_properties),
            'secondary': copy.copy(boot_dev_properties)
        }

        for dev in boot_dev_info:
            dev_path = self._get_system_prop(boot_dev_system_prop[dev])

            if not os.path.exists(dev_path):
                sys.stderr.write("Warning. Boot device " + dev_path + " is not exposed.\r\n" \
                                 "Skipping boot partition updates.\r\n\r\n" \
                                )
                dev_path = ""

            boot_dev_info[dev]['dev_path'] = dev_path

        main_boot_dev_path = boot_dev_info['main']['dev_path']

        if main_boot_dev_path is not "":
            main_boot_dev = open(main_boot_dev_path, 'rb')

            boot_block_size_log_2 = self._get_bct_prop(bct_device=main_boot_dev, bct_base_offset=0, bct_property="boot_block_size_log_2")
            boot_page_size_log_2 = self._get_bct_prop(bct_device=main_boot_dev, bct_base_offset=0, bct_property="boot_page_size_log_2")

            for dev in boot_dev_info:
                # check if data from BCT is reasonable
                if (boot_block_size_log_2 == 0) or (boot_block_size_log_2 > 23) or (boot_page_size_log_2 == 0) or (boot_page_size_log_2 > 23):
                    sys.stderr.write("Warning. First BCT is potentially corrupt.\r\n" \
                                    "Skipping boot partition updates.\r\n\r\n" \
                                    )
                    boot_dev_info[dev]['dev_path'] = ""
                else:
                    boot_dev_info[dev]['block_size'] = 1 << boot_block_size_log_2
                    boot_dev_info[dev]['page_size'] = 1 << boot_page_size_log_2

            main_boot_dev.close()

        if "jetson-tx1" in self.system_spec:
            for dev in boot_dev_info:
                write_en_path = "/sys/block/" + boot_dev_info[dev]['dev_path'].split("/")[-1] + "/force_ro"

                if not os.path.isfile(write_en_path):
                    sys.stderr.write("Warning. Boot partition write enable " + write_en_path + " is not exposed.\r\n" \
                                     "Skipping boot partition updates.\r\n\r\n" \
                                    )
                    write_en_path = ""

                boot_dev_info[dev]['write_en_path'] = write_en_path

        return boot_dev_info

    def _get_system_boot_partition_info(self):
        boot_partition_info = {
            'name_list'  : list(), # (e.g. [BCT, NVC, ...])
            'dev_dict'   : dict(), # (e.g. {'BCT': 'main', 'NVC': 'secondary', ...})
            'offset_dict': dict(), # (e.g. {'BCT': 0, 'NVC': 262144, ...})
            'size_dict'  : dict()  # (e.g. {'BCT': 262144, 'NVC': 196608, ...})
        }
        boot_partition_table = {
            'main'     : dict(),
            'secondary': dict()
        }

        if "jetson-nano-qspi" in self.system_spec:
            boot_partition_table['main'] = system_info._spi_main_boot_partition_table # /dev/mtdblock0
            boot_partition_table['secondary'] = None # only main boot device present
        elif "jetson-tx1" in self.system_spec:
            boot_partition_table['main'] = system_info._emmc_main_boot_partition_table # /dev/mmcblk0boot0
            boot_partition_table['secondary'] = system_info._emmc_sec_boot_partition_table # /dev/mmcblk0boot1
        # elif "jetson-NEWBOARD" in self.system_spec:
        #     boot_partition_table['main'] = system_info._NEWBOARD_main_boot_partition_table # /dev/NEWBOARDDEV
        #     boot_partition_table['secondary'] = system_info._NEWBOARD_sec_boot_partition_table # /dev/NEWBOARDDEV
        else:
            sys.stderr.write("Warning. Boot device updates are not currently supported on this platform.\r\n" \
                             "Skipping boot partition updates.\r\n\r\n" \
                            )

        for dev in boot_partition_table:
            if boot_partition_table[dev] is not None:
                boot_partition_info['name_list'] += list(boot_partition_table[dev].keys())

                for part_name, offset_size_tuple in boot_partition_table[dev].iteritems():
                    boot_partition_info['dev_dict'].update({part_name: dev})
                    boot_partition_info['offset_dict'].update({part_name: offset_size_tuple[0]})
                    boot_partition_info['size_dict'].update({part_name: offset_size_tuple[1]})

        return boot_partition_info

    def print_current_ver(self):
        ver_partition_boot_dev = self.system_boot_partition_info['dev_dict']['VER']
        ver_partition_path = self.system_boot_dev_info[ver_partition_boot_dev]['dev_path']
        ver_offset = self.system_boot_partition_info['offset_dict']['VER']
        ver_size = self.system_boot_partition_info['size_dict']['VER']

        ver_partition_device = open(ver_partition_path, 'wb+')
        ver_partition_device.seek(ver_offset, os.SEEK_SET)
        ver_data = ver_partition_device.read(ver_size)
        ver_partition_device.close();

        # For NV1 or older, there are no version data in VER,
        # we print out "NV1" directly.
        # For NV2 or newer VER, there are three lines (or more) version data,
        # we print out line 2:
        # NV2
        # #R32 (release), REVISION: 2.0, GCID: , BOARD: t210ref, EABI: aarch64, DATE: xxxx
        # BOARDID=3448 BOARDSKU=0000 FAB=100
        ver = ver_data.splitlines();
        if ver[0] == 'NV1' or 'NV' not in ver[0]:
            print "NV1"
        else:
            print ver[1]

class install_update_payload(update_type_payload, system_info):
    def __init__(self, arg):
        super(install_update_payload, self).__init__(arg)
        self.blob_entry_list = range(self.blob_header_dict['entry_count'])

        self.blob_entry_max_width_list = range(len(self.entry_name_tuple))
        for i in range(len(self.blob_entry_max_width_list)):
            self.blob_entry_max_width_list[i] = len(self.entry_name_tuple[i])

        self.print_all = arg.print_all

        self._generate_entry_list()
        self._validate_entry_list()

    def _get_partition_size(self, partition_path):
        partition_file = os.open(partition_path, os.O_RDONLY)

        try:
            return os.lseek(partition_file, 0, os.SEEK_END)
        finally:
            os.close(partition_file)

    def _generate_entry_list(self):
        self.blob_file.seek(self.blob_header_dict['header_size'], os.SEEK_SET)

        for idx, boot_entry in enumerate(self.blob_entry_list):
            try:
                blob_entry_tuple = struct.unpack(self.entry_packing, self.blob_file.read(struct.calcsize(self.entry_packing)))
                blob_entry_dict = dict(zip(self.entry_name_tuple, blob_entry_tuple))
                blob_entry_dict['part_name'] = blob_entry_dict['part_name'].strip(' \t\n\0')
                blob_entry_dict['tnspec'] = blob_entry_dict['tnspec'].strip(' \t\n\0')
                blob_entry_dict['updatable'] = True
                blob_entry_dict['intra_part_offset'] = 0 # offset within a partition

                blob_entry_op_mode = int(blob_entry_dict['op_mode'])
                blob_entry_spec = blob_entry_dict['tnspec']

                # Op mode of blob binary must either be 0 or equal to the system op mode to be updatable
                # Spec of blob binary must either be empty or equal to the system spec to be updatable
                if (blob_entry_op_mode != 0) and (blob_entry_op_mode != self.system_op_mode):
                    blob_entry_dict['updatable'] = False
                elif (len(blob_entry_spec) > 0) and (blob_entry_spec != self.system_spec):
                    blob_entry_dict['updatable'] = False

                self.blob_entry_list[idx] = blob_entry_dict

                # finding max column width for blob entry table print out
                for n in range(len(self.blob_entry_max_width_list)):
                    try:
                        if len(str(blob_entry_tuple[n]).strip(' \t\n\0')) > self.blob_entry_max_width_list[n]:
                            self.blob_entry_max_width_list[n] = len(str(blob_entry_tuple[n]).strip(' \t\n\0'))
                    except:
                        pass
            except:
                sys.stderr.write("Warning. Cannot parse partition number " + str(idx) + ".\r\n" \
                                 "Payload blob corrupt.\r\n" \
                                )

        if (self.system_boot_dev_info['main']['dev_path'] != "") and (self.system_boot_dev_info['secondary']['dev_path'] != ""):
            self._prep_boot_entries()

    def _prep_boot_entries(self):
        # update order: NVC_R, block-1 BCT "BCT2" to block-n BCT "BCT(N=n+1)", NVC, block-0 BCT "BCT" (contains 2 page-aligned BCT copies in 2 slots for qspi)
        # there will be as many BCT copies as required to fill the system's entire BCT partition, up to a maximum of 64 blocks (n=63)
        # these partitions will be updated last
        boot_part_name_update_order_list = ['NVC_R', 'BCTN', 'NVC', 'BCT']

        bct_copy_start_idx = boot_part_name_update_order_list.index('BCTN')
        bct_copy_list = list()

        for N in range(2,64):
            bct_copy_list.append('BCT' + str(N))

        boot_part_name_update_order_list[bct_copy_start_idx:bct_copy_start_idx] = bct_copy_list
        boot_part_name_update_order_list.remove('BCTN')

        if 'NVC_R' not in self.system_boot_partition_info['name_list']:
            boot_part_name_update_order_list.remove('NVC_R')

        boot_entry_list_append = list()

        try:
            bct_entry = filter(lambda blob_entry: (blob_entry['part_name'] == 'BCT') and (blob_entry['updatable'] is True), self.blob_entry_list)[0]
        except:
            # quietly skip BCT and NVC updates if no updatable entries exist in the payload for those partitions
            self._skip_boot_parts()
            return

        # make a secondary entry list only containing the entries from the part names in
        # boot_part_name_update_order_list
        for part_name in boot_part_name_update_order_list:
            part_name_base = part_name

            if 'BCT' in part_name:
                part_name_base = 'BCT'

                if part_name == part_name_base:
                    bct_num = 0
                else:
                    bct_num = int(part_name.replace(part_name_base, ''))

                if bct_num == 0:
                    # for block-0 slot-0, slot-1 will be handled in install_binaries() for qspi
                    bct_offset = 0
                elif bct_num > 0:
                    # calculate the offset for subsequent BCT copies
                    # only block aligned
                    bct_offset = self.system_boot_dev_info['main']['block_size'] * (bct_num - 1)
                else:
                    raise Exception("BCT number \"" + bct_num + "\" cannot be negative.")

                if bct_offset + bct_entry['part_size'] >= self.system_boot_partition_info['size_dict']['BCT']:
                    continue
            elif 'NVC' in part_name:
                part_name_base = 'NVC'

            try:
                blob_entry = filter(lambda entry: (entry['part_name'] == part_name_base) and (entry['updatable'] is True), self.blob_entry_list)[0]
            except:
                sys.stderr.write("Warning. " + part_name_base + " entry not found in payload.\r\n" \
                                "Skipping boot partition updates.\r\n\r\n" \
                                )
                self._skip_boot_parts()
                continue

            if part_name_base is 'BCT':
                # all entries in "boot_part_name_update_order_list" containing "BCT" become
                # aliases of the one updatable BCT entry in the payload, each BCT alias will
                # simply have a different intra_part_offset so that the same payload image
                # will be copied to different locations within the system BCT partition
                blob_entry = copy.deepcopy(blob_entry)
                blob_entry['intra_part_offset'] = bct_offset
                blob_entry['part_name'] = part_name
            elif part_name is 'NVC_R':
                # similarily for NVC_R, it will be an alias to the one updatable NVC payload
                # entry and will be copied to the system under the NVC_R partition if one exists
                blob_entry = copy.deepcopy(blob_entry)
                blob_entry['part_name'] = part_name

            boot_entry_list_append.append(blob_entry)

        # remove the partitions added to boot_entry_list_append from the master blob_entry_list
        for blob_entry in boot_entry_list_append:
            try:
                self.blob_entry_list.remove(blob_entry)
            except:
                pass

        # schedule the boot partition updates at the end
        self.blob_entry_list = self.blob_entry_list + boot_entry_list_append

    def _validate_entry_list(self):
        for blob_entry in self.blob_entry_list:
            if blob_entry['updatable'] == False:
                continue

            part_name = str(blob_entry['part_name'])

            part_name_base = part_name
            if 'BCT' in part_name:
                part_name_base = 'BCT'

            if part_name_base in self.system_gpt_partition_list:
                part_path = self.system_gpt_partition_dir + part_name_base

                # Check partition size in blob does not exceed GPT partition on system
                system_partition_size = self._get_partition_size(partition_path=part_path)
            elif part_name_base in self.system_boot_partition_info['name_list']:
                if self.system_boot_dev_info['main']['dev_path'] == self.system_boot_dev_info['secondary']['dev_path'] == "":
                    sys.stderr.write("Warning. Skipping boot partition updates due to previous warning(s).\r\n\r\n")
                    self._skip_boot_parts()
                    continue

                # Check partition size in blob does not exceed boot partition on system
                system_partition_size = self.system_boot_partition_info['size_dict'][part_name_base]

                # Checking block-0 BCT
                if part_name == 'BCT':
                    # require two BCTs in block-0 for jetson-nano-qspi*
                    if "jetson-nano-qspi" in self.system_spec:
                        block0_bct_cpy = 2
                    else:
                        block0_bct_cpy = 1

                    if (blob_entry['part_size'] * block0_bct_cpy > self.system_boot_dev_info['main']['block_size']) or (blob_entry['part_size'] % self.system_boot_dev_info['main']['page_size'] != 0):
                        sys.stderr.write("Warning. One or more conditions below are true.\r\n" \
                                         "Payload BCT size * " + block0_bct_cpy + " (" + str(blob_entry['part_size'] * block0_bct_cpy) + " bytes) exceeds boot device block size (" + str(self.system_boot_dev_info['main']['block_size']) + " bytes).\r\n" \
                                         "Payload BCT size (" + str(blob_entry['part_size']) + " bytes) is not a multiple of boot device page size (" + str(self.system_boot_dev_info['main']['page_size']) + " bytes).\r\n" \
                                         "Skipping boot partition updates.\r\n\r\n" \
                                        )
                        self._skip_boot_parts()
                        continue
            else:
                blob_entry['updatable'] = False
                sys.stderr.write("Warning. Partition \"" + part_name_base + "\" in blob not found in system GPT or boot partition list.\r\n" \
                                 "Skipping this partition.\r\n"
                                )

            if blob_entry['part_size'] > system_partition_size:
                sys.stderr.write("Error. Partition \"" + part_name_base + "\" in blob exceeds system partition size on system.\r\n" \
                    "  Blob partition size:" + " {:,}".format(blob_entry['part_size']) + " bytes\r\n" \
                    "System partition size:" + " {:,}".format(system_partition_size) + " bytes\r\n" \
                    "No changes have been made. Exiting...\r\n" \
                )
                sys.exit(1)

            # Check spec of partition in blob matches spec of system
            blob_entry_spec = blob_entry['tnspec']
            if len(blob_entry_spec) > 0:
                if blob_entry_spec != self.system_spec:
                    sys.stderr.write("Error. Partition \"" + part_name_base + "\" in blob does not match system spec.\r\n" \
                                        "Blob partition spec: " + blob_entry_spec + "\r\n" \
                                        "        System spec: " + self.system_spec + "\r\n" \
                                        "No changes have been made. Exiting...\r\n" \
                                    )
                    sys.exit(1)
        return

    def _skip_boot_parts(self):
        filtered_blob_entry_list = filter(lambda blob_entry: (blob_entry['part_name'] in self.system_boot_partition_info['name_list']) and (blob_entry['updatable'] is True), self.blob_entry_list)

        for filtered_entry in filtered_blob_entry_list:
            self.blob_entry_list[self.blob_entry_list.index(filtered_entry)]['updatable'] = False

    def _write_partition(self, bin_data, block_dev, block_dev_offset=0):
        print "[W]riting" + " {:,}".format(len(bin_data)) + " bytes to " + block_dev.name + " at offset" + " {:,}".format(block_dev_offset)

        block_dev.seek(block_dev_offset, os.SEEK_SET)
        block_dev.write(bin_data)
        return

    def _read_partition(self, bin_data_size, block_dev, block_dev_offset=0):
        print "[R]eading" + " {:,}".format(bin_data_size) + " bytes of " + block_dev.name + " at offset" + " {:,}".format(block_dev_offset)

        block_dev.seek(block_dev_offset, os.SEEK_SET)
        system_bin_data = block_dev.read(bin_data_size)

        return system_bin_data

    def _set_boot_part_write_en(self, write_enable):
        if "jetson-nano-qspi" in self.system_spec:
            return # qspi device always has writes enabled

        for dev in self.system_boot_dev_info:
            try:
                boot_partition_wr_en = open(self.system_boot_dev_info[dev]['write_en_path'], 'rb+')

                if write_enable == True:
                    boot_partition_wr_en.write(str(0))
                elif write_enable == False:
                    boot_partition_wr_en.write(str(1))

                boot_partition_wr_en.close()
            except Exception:
                sys.stderr.write("Error. Cannot enable writing to boot partition. Exiting...\r\n")
                sys.exit(1)

    def print_blob_header(self):
        print "BLOB HEADER:"
        print "       Magic: " + self.blob_header_dict['magic']
        print "     Version: " + format(self.blob_header_dict['version'], "#010x")
        print "   Blob Size: " + "{:,}".format(self.blob_header_dict['blob_size']) + " bytes"
        print " Header Size: " + "{:,}".format(self.blob_header_dict['header_size']) + " bytes"
        print " Entry Count: " + str(self.blob_header_dict['entry_count']) + " partition(s)"
        print "        Type: " + str(self.blob_header_dict['type']) + " (0 for update, 1 for BMP)"
        print "Uncompressed\r\n" \
              "   Blob Size: " + "{:,}".format(self.blob_header_dict['uncomp_blob_size']) + " bytes"
        print "   Accessory:",
        if self.accessory_present == True:
            print format(self.blob_header_dict['accessory'], "#018x")
        else:
            print "Not Present"
        return

    def print_entry_table(self):
        print "ENTRY TABLE:"
        print "|",
        for idx, entry_name in enumerate(self.entry_name_tuple):
            if entry_name is "intra_part_offset":
                continue
            elif (entry_name is "updatable") and (self.print_all is False):
                continue
            else:
                print entry_name.center(self.blob_entry_max_width_list[idx]) + " |",
        print
        for blob_entry in self.blob_entry_list:
            if (blob_entry['updatable'] is True) or (self.print_all is True):
                print "|",
                try:
                    print str(blob_entry['part_name']).rjust(self.blob_entry_max_width_list[0]) + " |",
                    print str(blob_entry['offset']).rjust(self.blob_entry_max_width_list[1]) + " |",
                    print str(blob_entry['part_size']).rjust(self.blob_entry_max_width_list[2]) + " |",
                    print str(blob_entry['version']).center(self.blob_entry_max_width_list[3]) + " |",
                    print str(blob_entry['op_mode']).center(self.blob_entry_max_width_list[4]) + " |",
                    print str(blob_entry['tnspec']).ljust(self.blob_entry_max_width_list[5]) + " |",
                    if self.print_all is True:
                        print str(blob_entry['updatable']).ljust(self.blob_entry_max_width_list[6]) + " |"
                    else:
                        print ""
                except:
                    print "SKIPPED".center(sum(self.blob_entry_max_width_list) + (len(self.blob_entry_max_width_list)*2) + 3) + " |"
                    pass

        if self.print_all is False:
            print "Note: only partitions updatable on this system are shown, use \"-p\" to show all partitions in payload."

        print "Note: partitions are written to system top down in the order shown above."
        return

    def install_binaries(self):
        total_bytes_written = 0
        total_parts_written = 0
        total_bytes_updatable = 0
        total_parts_updatable = 0

        if len(self.system_boot_partition_info['name_list']) > 0:
            self._set_boot_part_write_en(write_enable=True)

        for blob_entry in self.blob_entry_list:
            if blob_entry['updatable'] == False:
                continue

            skip_write = False

            part_name = str(blob_entry['part_name'])
            dev_offset = 0

            part_name_base = part_name
            if 'BCT' in part_name:
                part_name_base = 'BCT'

                if part_name == part_name_base:
                    bct_num = 0
                else:
                    bct_num = int(part_name.replace(part_name_base, ''))

            if part_name_base in self.system_gpt_partition_list:
                system_partition_path = os.path.join(self.system_gpt_partition_dir, part_name_base)
                dev_offset = blob_entry['intra_part_offset']
            elif part_name_base in self.system_boot_partition_info['name_list']:
                partition_boot_dev = self.system_boot_partition_info['dev_dict'][part_name_base]
                system_partition_path = self.system_boot_dev_info[partition_boot_dev]['dev_path']
                dev_offset = self.system_boot_partition_info['offset_dict'][part_name_base]
                dev_offset += blob_entry['intra_part_offset']

            try:
                system_partition_device = open(system_partition_path, 'wb+')
            except Exception as e:
                sys.stderr.write("Error. " + str(e) + "\r\n" \
                                 "Cannot open system partition \"" + system_partition_path + "\". Exiting...\r\n"
                                )
                sys.exit(1)

            # Seek to the offset of the partition data within the blob
            self.blob_file.seek(int(blob_entry['offset']), os.SEEK_SET)

            # Assign the binary data of the update to a variable
            update_bin_data_size = int(blob_entry['part_size'])
            update_bin_data = self.blob_file.read(update_bin_data_size)

            # require two BCTs in block-0 for jetson-nano-qspi*
            if "jetson-nano-qspi" in self.system_spec:
                block0_bct_cpy = 2
            else:
                block0_bct_cpy = 1

            # BCT payload is checked already to be page aligned, so the second slot
            # is directly appended if required
            if (part_name_base == 'BCT') and (bct_num == 0):
                update_bin_data *= block0_bct_cpy
                update_bin_data_size *= block0_bct_cpy
                if block0_bct_cpy > 1:
                    print "Note: BCT in block-0 is " + str(block0_bct_cpy) + " times the usual size as it holds " + str(block0_bct_cpy) + " copies."

            print "Updating partition: " + part_name
            pre_update_system_bin_data = self._read_partition(bin_data_size=update_bin_data_size, block_dev=system_partition_device, block_dev_offset=dev_offset)
            total_bytes_updatable += update_bin_data_size
            total_parts_updatable += 1

            if pre_update_system_bin_data == update_bin_data:
                sys.stderr.write("Warning. " + part_name + " partition data on system is identical to payload data. Skipping update.\r\n\r\n")
                skip_write = True

            if skip_write is False:
                self._write_partition(bin_data=update_bin_data, block_dev=system_partition_device, block_dev_offset=dev_offset)
                total_bytes_written += update_bin_data_size
                total_parts_written += 1

                # Run sync twice to ensure data is written to NVM
                subprocess.call("sync")
                subprocess.call("sync")

                post_update_system_bin_data = self._read_partition(bin_data_size=update_bin_data_size, block_dev=system_partition_device, block_dev_offset=dev_offset)
                if post_update_system_bin_data == update_bin_data:
                    print "Verification successful."
                else:
                    sys.stderr.write("Error. Verification failed. Exiting...\r\n")
                    sys.exit(1)

                print

            system_partition_device.close()

        if len(self.system_boot_partition_info['name_list']) > 0:
            self._set_boot_part_write_en(write_enable=False)

        print "Total data updatable for system: " + str(total_parts_updatable) + " partitions, {:,}".format(total_bytes_updatable) + " bytes."
        print "   Total data written to system: " + str(total_parts_written) + " partitions, {:,}".format(total_bytes_written) + " bytes."
        print ""

        return

# Main function
def main(arg):
    if (arg.print_ver is True):
        print_VER(arg)
    else:
        install_BUP(arg)

if __name__ == '__main__':
    if not os.geteuid() == 0:
        sys.stderr.write("Error. Script must be run as root. Exiting...\r\n")
        sys.exit(1)

    try:
        param = parse_args()
    except Exception:
        sys.stderr.write("Error. Provided blob file does not exist. Exiting...\r\n")
        sys.exit(1)

    main(param)
    sys.exit(0)
