#!/usr/bin/env python

"""
rpi-eeprom-config
"""

import argparse
import atexit
import os
import subprocess
import struct
import sys
import tempfile
import time

IMAGE_SIZE = 512 * 1024

MAX_BOOTCONF_SIZE = 2024

# Each section starts with a magic number followed by a 32 bit offset to the
# next section (big-endian).
# The number, order and size of the sections depends on the bootloader version
# but the following mask can be used to test for section headers and skip
# unknown data.
MAGIC = 0x55aaf00f
MAGIC_MASK = 0xfffff00f
FILE_MAGIC = 0x55aaf11f # id for modifiable file, currently only bootconf.txt
FILE_HDR_LEN = 20
FILENAME_LEN = 12
TEMP_DIR = None

def exit_handler():
    """
    Delete any temporary files.
    """
    if TEMP_DIR is not None and os.path.exists(TEMP_DIR):
        tmp_image = os.path.join(TEMP_DIR, 'pieeprom.upd')
        if os.path.exists(tmp_image):
            os.remove(tmp_image)
        tmp_conf = os.path.join(TEMP_DIR, 'boot.conf')
        if os.path.exists(tmp_conf):
            os.remove(tmp_conf)
        os.rmdir(TEMP_DIR)

def create_tempdir():
    global TEMP_DIR
    if TEMP_DIR is None:
        TEMP_DIR = tempfile.mkdtemp()

def exit_error(msg):
    """
    Trapped a fatal arror, output message to stderr and exit with non-zero
    return code.
    """
    sys.stderr.write("ERROR: %s\n" % msg)
    sys.exit(1)

def shell_cmd(args):
    """
    Executes a shell command waits for completion returning STDOUT. If an
    error occurs then exit and output the subprocess stdout, stderr messages
    for debug.
    """
    start = time.time()
    arg_str = ' '.join(args)
    result = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    while time.time() - start < 5:
        if result.poll() is not None:
            break

    if result.poll() is None:
        exit_error("%s timeout" % arg_str)

    if result.returncode != 0:
        exit_error("%s failed: %d\n %s\n %s\n" %
                   (arg_str, result.returncode, result.stdout.read(), result.stderr.read()))
    else:
        return result.stdout.read()

def get_latest_eeprom():
    """
    Returns the path of the latest EEPROM image file if it exists.
    """
    latest = shell_cmd(['rpi-eeprom-update', '-l']).rstrip()
    if not os.path.exists(latest):
        exit_error("EEPROM image '%s' not found" % latest)
    return latest

def apply_update(config, eeprom=None):
    """
    Applies the config file to the latest available EEPROM image and spawns
    rpi-eeprom-update to schedule the update at the next reboot.
    """
    if eeprom is not None:
       eeprom_image = eeprom
    else:
       eeprom_image = get_latest_eeprom()
    create_tempdir()
    tmp_update = os.path.join(TEMP_DIR, 'pieeprom.upd')
    image = BootloaderImage(eeprom_image, tmp_update)
    image.write(config)
    config_str = open(config).read()

    sys.stdout.write("Updating bootloader EEPROM\n image: %s\nconfig: %s\n%s\n" %
                     (eeprom_image, config, config_str))
    args = ['sudo', 'rpi-eeprom-update', '-d', '-f', tmp_update]
    resp = shell_cmd(args)
    sys.stdout.write(resp)

def edit_config(eeprom=None):
    """
    Implements something like visudo for editing EEPROM configs.
    """
    if 'EDITOR' not in os.environ:
        exit_error('EDITOR environment variable not defined')
    create_tempdir()
    current_config = read_current_config()
    tmp_conf = os.path.join(TEMP_DIR, 'boot.conf')
    out = open(tmp_conf, 'w')
    out.write(current_config)
    out.close()
    cmd = "\'%s\' \'%s\'" % (os.environ['EDITOR'], tmp_conf)
    result = os.system(cmd)
    if result != 0:
        exit_error("Aborting update because \'%s\' exited with code %d." % (cmd, result))

    new_config = open(tmp_conf, 'r').read()
    if len(new_config.splitlines()) < 2:
        exit_error("Aborting update because \'%s\' appears to be empty." % tmp_conf)
    apply_update(tmp_conf, eeprom)

def read_current_config():
    """
    Reads the configuration used by the current bootloader.
    """
    result = shell_cmd(['vcgencmd', 'bootloader_config'])
    if result is None:
        exit_error('Failed to read the current bootloader configuration')
    return result

class BootloaderImage(object):
    def __init__(self, filename, output):
        """
        Instantiates a Bootloader image writer with a source eeprom (filename)
        and optionally an output filename.
        """
        self._filename = filename
        try:
            self._bytes = bytearray(open(filename, 'rb').read())
        except IOError as err:
            exit_error("Failed to read \'%s\'\n%s\n" % (filename, str(err)))
        self._out = None
        if output is not None:
            self._out = open(output, 'wb')

        if len(self._bytes) != IMAGE_SIZE:
            exit_error("%s: Expected size %d bytes actual size %d bytes" %
                       (filename, IMAGE_SIZE, len(self._bytes)))

    def find_config(self):
        offset = 0
        magic = 0
        while offset < IMAGE_SIZE:
            magic, length = struct.unpack_from('>LL', self._bytes, offset)
            if (magic & MAGIC_MASK) != MAGIC:
                raise Exception('EEPROM is corrupted')

            if magic == FILE_MAGIC: # Found a file
                name = self._bytes[offset + 8: offset + FILE_HDR_LEN]
                if name.decode('utf-8') == 'bootconf.txt':
                    return (offset, length)

            offset += 8 + length # length + type
            offset = (offset + 7) & ~7

        raise Exception('EEPROM parse error: Bootloader config not found')

    def write(self, new_config):
        hdr_offset, length = self.find_config()
        new_config_bytes = open(new_config, 'rb').read()
        new_len = len(new_config_bytes) + FILENAME_LEN + 4
        if len(new_config_bytes) > MAX_BOOTCONF_SIZE:
            raise Exception("Config is too large (%d bytes). The maximum size is %d bytes."
                            % (len(new_config_bytes), MAX_BOOTCONF_SIZE))
        if hdr_offset + len(new_config_bytes) + FILE_HDR_LEN > IMAGE_SIZE:
            raise Exception('EEPROM image size exceeded')

        struct.pack_into('>L', self._bytes, hdr_offset + 4, new_len)
        struct.pack_into(("%ds" % len(new_config_bytes)), self._bytes,
                         hdr_offset + 4 + FILE_HDR_LEN, new_config_bytes)

        # If the new config is smaller than the old config then set any old
        # data which is now unused to all ones (erase value)
        pad_start = hdr_offset + 4 + FILE_HDR_LEN + len(new_config_bytes)
        pad = 0
        while pad < (length - len(new_config_bytes)):
            struct.pack_into('B', self._bytes, pad_start + pad, 0xff)
            pad = pad + 1

        if self._out is not None:
            self._out.write(self._bytes)
            self._out.close()
        else:
            if hasattr(sys.stdout, 'buffer'):
                sys.stdout.buffer.write(self._bytes)
            else:
                sys.stdout.write(self._bytes)

    def read(self):
        hdr_offset, length = self.find_config()
        offset = hdr_offset + 4 + FILE_HDR_LEN
        config_bytes = self._bytes[offset:offset+length-FILENAME_LEN-4]
        if self._out is not None:
            self._out.write(config_bytes)
            self._out.close()
        else:
            if hasattr(sys.stdout, 'buffer'):
                sys.stdout.buffer.write(config_bytes)
            else:
                sys.stdout.write(config_bytes)

def main():
    """
    Utility for reading and writing the configuration file in the
    Raspberry Pi 4 bootloader EEPROM image.
    """
    description = """\
Bootloader EEPROM configuration tool for the Raspberry Pi 4.
Operating modes:

1. Outputs the current bootloader configuration to STDOUT if no arguments are
   specified OR the given output file if --out is specified.

   rpi-eeprom-config [--out boot.conf]

2. Extracts the configuration file from the given \'eeprom\' file and outputs
   the result to STDOUT or the output file if --output is specified.

   rpi-eeprom-config pieeprom.bin [--out boot.conf]

3. Writes a new EEPROM image replacing the configuration file with the contents
   of the file specified by --config.

   rpi-eeprom-config --config boot.conf --out newimage.bin pieeprom.bin

   The new image file can be installed via rpi-eeprom-update
   rpi-eeprom-update -d -f newimage.bin

4. Applies a given config file an EEPROM image and invokes rpi-eeprom-update
   to schedule an update of the bootloader when the system is rebooted.

   rpi-eeprom-config --apply boot.conf [pieeprom.bin]

   If the \'eeprom\' argument is not specified then the latest available image
   is selected by calling \'rpi-eeprom-update -l\'.

5. The `--edit` parameter behaves the same as `--apply` except that instead of
   applying a predefined configuration file the default text editor ($EDITOR) is
   launched with the contents of the current EEPROM configuration.

   N.B. The 'current' EEPROM configuration is read via 'vcgencmd bootloader_config'
   and is the configuration that the system was booted with. It does not reflect any
   pending updates so invoking this a second time without rebooting would discard
   any previous edits.

   rpi-eeprom-config --edit [pieeprom.bin]

See 'rpi-eeprom-update -h' for more information about the available EEPROM
images.
"""
    parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter,
                                     description=description)

    parser.add_argument('-a', '--apply', required=False,
                        help='Updates the bootloader to the given config plus latest available EEPROM release.')
    parser.add_argument('-c', '--config', help='Name of bootloader configuration file', required=False)
    parser.add_argument('-e', '--edit', action='store_true', default=False, help='Edit the current EEPROM config')
    parser.add_argument('-o', '--out', help='Name of output file', required=False)
    parser.add_argument('eeprom', nargs='?', help='Name of EEPROM file to use as input')
    args = parser.parse_args()

    if args.edit:
        edit_config(args.eeprom)
    elif args.apply is not None:
        if not os.path.exists(args.apply):
            exit_error("config file '%s' not found" % args.apply)
        apply_update(args.apply, args.eeprom)
    elif args.eeprom is not None:
        image = BootloaderImage(args.eeprom, args.out)
        if args.config is not None:
            if not os.path.exists(args.config):
                exit_error("config file '%s' not found" % args.config)
            image.write(args.config)
        else:
            image.read()
    elif args.config is None and args.eeprom is None:
        current_config = read_current_config()
        if args.out is not None:
            open(args.out, 'w').write(current_config)
        else:
            sys.stdout.write(current_config)

if __name__ == '__main__':
    atexit.register(exit_handler)
    main()