# Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""This module supports creating Exynos bootprom images."""

import hashlib
import os
import struct

from tools import CmdError

HASH_ALGO_LEN = 32
HASH_LEN = 32
HASH_SIGNATURE = 0xdead4eef
HASH_VERSION = 1
HASH_FLAGS = 0
HASH_HEADER_LEN = 16


class ExynosBl2(object):
  """Class for processing Exynos SPL blob.

  Second phase loader (SPL) is also called boot loader 2 (BL2), these terms
  mean the same and are used in this class interchangeably.

  SPL is a binary blob which is in fact a short program running from internal
  SRAM. It initializes main DRAM and loads the actual boot loader after that.
  The program is encapsulated using one of two methods - fixed or variable
  size. Both methods provide rudimentary checksum protection.

  SPL is supposed to know some details about the hardware it runs on. This
  information is stored in the so called machine parameters structure in the
  blob. Some of it is available at compile time, but most of it comes form the
  platform specific flat device tree. SPL generated by u-boot make file
  includes machine parameter structure with default configuration values, not
  suitable to run on actual hardware.

  This class provides the following services:

  - check integrity of the passed in SPL blob, and determining its
    the encapsulation method along the way.

  - parse the passed in device tree and pack the retrieved information into
    the machine parameters structure. The structure location in the blob is
    identified by a 4 byte structure header signature.

    Note that this method of finding the structure in the blob is quite
    brittle: it would silently produce catastrophically wrong result if for
    some reason, this same pattern is present anywhere in the blob above the
    structure.

  - update the checksum as appropriate for the detected format, and save the
    modified SPL in a file.

  Attributes:
    _tools: A tools.Tools object to use for external tools, provided by the
            caller
    _out: A cros_output.Output object to use for console output, provided by
            the caller
    _spl_type: an enum defined below, describing type (fixed of variable size)
            of the SPL being handled
  """

  VAR_SPL = 1
  FIXED_SPL = 2

  def __init__(self, tools, output):
    """Set up a new object."""
    self._tools = tools
    self._out = output
    self._spl_type = None
    self.spl_source = 'straps'  # SPL boot according to board settings

  def _BootingUsingEFS(self, fdt, use_efs_memory):
    """Check if we are booting using early-firmware-selection.

    This is just a helper function to avoid using the same logic in many
    places.

    Args:
      fdt: Device tree file to use.
      use_efs_memory: True to use early-firmware-selection memory (i.e. IRAM),
          False, to ignore it.

    Returns:
      True if EFS is enabled and we are configured to use EFS memory, else
          False.
    """
    return (use_efs_memory and
            fdt.GetInt('/chromeos-config', 'early-firmware-selection', 0))

  def _GetAddress(self, fdt, use_efs_memory, name, config_node='/config',
                  allow_none=False):
    """Work out the correct address for a region of memory.

    This deals with EFS and the memory map automatically.

    Args:
      fdt: Device tree file containing memory map.
      use_efs_memory: True to return the address in EFS memory (i.e. SRAM),
          False to use SDRAM
      name: Name of the region to look up, e.g. 'u-boot'
      config_node: Node containing configuration information
      allow_none: True if it is OK to find nothing.

    Returns:
      Address to load that region, or None if none.
    """
    efs_suffix = ''
    if self._BootingUsingEFS(fdt, use_efs_memory):
      efs_suffix = ',efs'

    # Use the correct memory section, and then find the offset in that.
    default = 'none' if allow_none else None
    memory = fdt.GetString(config_node, '%s-memory%s' % (name, efs_suffix),
                           default)
    if memory == 'none':
      return None
    base = fdt.GetIntList(memory, 'reg')[0]
    offset = fdt.GetIntList(config_node, '%s-offset%s' % (name, efs_suffix))[0]
    addr = base + offset
    return addr

  def GetUBootAddress(self, fdt, use_efs_memory):
    """Work out the correct address for loading U-Boot.

    This deals with EFS and the memory map automatically.

    Args:
      fdt: Device tree file containing memory map.
      use_efs_memory: True to return the address in EFS memory (i.e. SRAM),
          False to use SDRAM

    Returns:
      Address to load U-Boot
    """
    addr = self._GetAddress(fdt, use_efs_memory, 'u-boot')
    self._out.Notice('EFS: Loading U-Boot to %x' % addr)
    return addr

  def _GetRWSPLDetails(self, fdt, use_efs_memory):
    if not self._BootingUsingEFS(fdt, use_efs_memory):
      return 0, 0, 0

    memory = fdt.GetString('/chromeos-config', 'rw-spl-memory,efs')
    base = fdt.GetIntList(memory, 'reg')[0]
    offset, size = fdt.GetIntList('/chromeos-config', 'rw-spl-offset,efs')
    addr = base + offset
    return 1, addr, size

  def _UpdateParameters(self, fdt, spl_load_offset, spl_load_size, data, pos,
                        use_efs_memory, skip_sdram_init):
    """Update the parameters in a BL2 blob.

    We look at the list in the parameter block, extract the value of each
    from the device tree, and write that value to the parameter block.

    Args:
      fdt: Device tree containing the parameter values.
      spl_load_offset: Offset in boot media that SPL must start loading (bytes)
      spl_load_size: Size of U-Boot image that SPL must load
      data: The BL2 data.
      pos: The position of the start of the parameter block.
      use_efs_memory: True to return the address in EFS memory (i.e. SRAM),
          False to use SDRAM
      skip_sdram_init: True to skip SDRAM initialization.

    Returns:
      The new contents of the parameter block, after updating.
    """
    version, size = struct.unpack('<2L', data[pos + 4:pos + 12])
    if version != 1:
      raise CmdError("Cannot update machine parameter block version '%d'" %
                     version)
    if size < 0 or pos + size > len(data):
      raise CmdError('Machine parameter block size %d is invalid: '
                     'pos=%d, size=%d, space=%d, len=%d' %
                     (size, pos, size, len(data) - pos, len(data)))

    # Move past the header and read the parameter list, which is terminated
    # with \0.
    pos += 12
    param_list = struct.unpack('<%ds' % (len(data) - pos), data[pos:])[0]
    param_len = param_list.find('\0')
    param_list = param_list[:param_len]
    pos += (param_len + 4) & ~3

    # Work through the parameters one at a time, adding each value
    new_data = ''
    upto = 0
    for param in param_list:
      value = struct.unpack('<1L', data[pos + upto:pos + upto + 4])[0]

      # Use this to detect a missing value from the fdt.
      not_given = 'not-given-invalid-value'
      if param == 'm':
        mem_type = fdt.GetString('/dmc', 'mem-type', not_given)
        if mem_type == not_given:
          mem_type = 'ddr3'
          self._out.Warning("No value for memory type: using '%s'" % mem_type)
        mem_types = ['ddr2', 'ddr3', 'lpddr2', 'lpddr3']
        if mem_type not in mem_types:
          raise CmdError("Unknown memory type '%s'" % mem_type)
        value = mem_types.index(mem_type)
        self._out.Info('  Memory type: %s (%d)' % (mem_type, value))
      elif param == 'M':
        mem_manuf = fdt.GetString('/dmc', 'mem-manuf', not_given)
        if mem_manuf == not_given:
          mem_manuf = 'samsung'
          self._out.Warning("No value for memory manufacturer: using '%s'" %
                            mem_manuf)
        mem_manufs = ['autodetect', 'elpida', 'samsung']
        if mem_manuf not in mem_manufs:
          raise CmdError("Unknown memory manufacturer: '%s'" % mem_manuf)
        value = mem_manufs.index(mem_manuf)
        self._out.Info('  Memory manufacturer: %s (%d)' % (mem_manuf, value))
      elif param == 'f':
        mem_freq = fdt.GetInt('/dmc', 'clock-frequency', -1)
        if mem_freq == -1:
          mem_freq = 800000000
          self._out.Warning("No value for memory frequency: using '%s'" %
                            mem_freq)
        mem_freq /= 1000000
        if mem_freq not in [533, 667, 800]:
          self._out.Warning("Unexpected memory speed '%s'" % mem_freq)
        value = mem_freq
        self._out.Info('  Memory speed: %d' % mem_freq)
      elif param == 'a':
        arm_freq = fdt.GetInt('/dmc', 'arm-frequency', -1)
        if arm_freq == -1:
          arm_freq = 1700000000
          self._out.Warning("No value for ARM frequency: using '%s'" %
                            arm_freq)
        arm_freq /= 1000000
        value = arm_freq
        self._out.Info('  ARM speed: %d' % arm_freq)
      elif param == 'i':
        i2c_addr = -1
        lookup = fdt.GetString('/aliases', 'pmic', '')
        if lookup:
          i2c_addr, size = fdt.GetIntList(lookup, 'reg', 2)
        if i2c_addr == -1:
          self._out.Warning('No value for PMIC I2C address: using %#08x' %
                            value)
        else:
          value = i2c_addr
        self._out.Info('  PMIC I2C Address: %#08x' % value)
      elif param == 's':
        serial_addr = -1
        lookup = fdt.GetString('/aliases', 'console', '')
        if lookup:
          serial_addr, size = fdt.GetIntList(lookup, 'reg', 2)
        if serial_addr == -1:
          self._out.Warning('No value for Console address: using %#08x' %
                            value)
        else:
          value = serial_addr
        self._out.Info('  Console Address: %#08x' % value)
      elif param == 'v':
        value = 31
        self._out.Info('  Memory interleave: %#0x' % value)
      elif param == 'u':
        value = spl_load_size
        self._out.Info('  U-Boot size: %#0x' % value)
      elif param == 'S':
        value = self.GetUBootAddress(fdt, use_efs_memory)
        self._out.Info('  U-Boot start: %#0x' % value)
      elif param == 'o':
        value = spl_load_offset
        self._out.Info('  U-Boot offset: %#0x' % value)
      elif param == 'l':
        load_addr = fdt.GetInt('/config', 'u-boot-load-addr', -1)
        if load_addr == -1:
          self._out.Warning("No value for U-Boot load address: using '%08x'" %
                            value)
        else:
          value = load_addr
        self._out.Info('  U-Boot load address: %#0x' % value)
      elif param == 'b':
        # These values come from enum boot_mode in U-Boot's cpu.h
        # For EFS we select SPI as the boot source always. We could support
        # eMMC if we want to add EFS support for eMMC.
        if (self.spl_source == 'spi' or
            self._BootingUsingEFS(fdt, use_efs_memory)):
          value = 20
        elif self.spl_source == 'straps':
          value = 32
        elif self.spl_source == 'emmc':
          value = 4
        elif self.spl_source == 'usb':
          value = 33
        else:
          raise CmdError("Invalid boot source '%s'" % self.spl_source)
        self._out.Info('  Boot source: %#0x' % value)
      elif param in ['r', 'R']:
        records = fdt.GetIntList('/board-rev', 'google,board-rev-gpios',
                                 None, '0 0')
        gpios = []
        for i in range(1, len(records), 3):
          gpios.append(records[i])
        gpios.extend([0, 0, 0, 0])
        if param == 'r':
          value = gpios[0] + (gpios[1] << 16)
          self._out.Info('  Board ID GPIOs: tit0=%d, tit1=%d' % (gpios[0],
                                                                 gpios[1]))
        else:
          value = gpios[2] + (gpios[3] << 16)
          self._out.Info('  Board ID GPIOs: tit2=%d, tit3=%d' % (gpios[2],
                                                                 gpios[3]))
      elif param == 'w':
        records = fdt.GetIntList('/config', 'google,bad-wake-gpios',
                                 3, '0 0xffffffff 0')
        value = records[1]
        self._out.Info('  Bad Wake GPIO: %#x' % value)
      elif param == 'z':
        compress = fdt.GetString('/flash/ro-boot', 'compress', 'none')
        compress_types = ['none', 'lzo']
        if compress not in compress_types:
          raise CmdError("Unknown compression type '%s'" % compress)
        value = compress_types.index(compress)
        self._out.Info('  Compression type: %#0x' % value)
      elif param == 'c':
        rtc_type = 0
        try:
          rtc_alias = fdt.GetString('/aliases/', 'rtc')
          rtc_compat = fdt.GetString(rtc_alias, 'compatible')
          if rtc_compat == 'samsung,s5m8767-pmic':
            rtc_type = 1
          elif rtc_compat == 'maxim,max77802-pmic':
            rtc_type = 2
        except CmdError:
          self._out.Warning('Failed to find rtc')
        value = rtc_type
      elif param == 'W':
        try:
          records = fdt.GetIntList('/chromeos-config/vboot-flag-write-protect',
                                   'gpio', 3)
          value = records[1]
          self._out.Info('  Write Protect GPIO: %#x' % value)
        except CmdError:
          self._out.Warning('No value for write protect GPIO: using %#x' %
                            value)
      elif param in ['j', 'A', 'U']:
        jump, addr, size = self._GetRWSPLDetails(fdt, use_efs_memory)
        if param == 'j':
          value = jump
          self._out.Info('  Jump to RW SPL: %d' % value)
        elif param == 'A':
          value = addr
          self._out.Info('  RW SPL addr: %#x' % value)
        elif param == 'U':
          value = size
          self._out.Info('  RW SPL size: %#x' % value)
      elif param == 'd':
        value = 1 if skip_sdram_init else 0
        self._out.Info('  Skip SDRAM init: %d' % value)
      elif param == 'p':
        addr = self._GetAddress(fdt, use_efs_memory, 'vboot-persist',
                                '/chromeos-config', True)
        if addr is None:
          value = 0
        else:
          value = addr
        self._out.Info('  Vboot persist addr: %x' % value)
      elif param == 'D':
        value = fdt.GetBool('/config', 'spl-debug')
        self._out.Info('  SPL debug: %d' % value)
      else:
        self._out.Warning("Unknown machine parameter type '%s'" % param)
        self._out.Info('  Unknown value: %#0x' % value)
      new_data += struct.pack('<L', value)
      upto += 4

    # Put the data into our block.
    data = data[:pos] + new_data + data[pos + len(new_data):]
    return data

  def _UpdateChecksum(self, data):
    """Update the BL2 checksum.

    The checksum is a 4 byte sum of all the bytes in the image before the
    last 4 bytes (which hold the checksum).

    Args:
      data: The BL2 data to update.

    Returns:
      The new contents of the BL2 data, after updating the checksum.

    Raises:
      CmdError if spl type is not set properly.
    """

    if self._spl_type == self.FIXED_SPL:
      checksum = sum(ord(x) for x in data[:-4])
      checksum_offset = len(data) - 4
    elif self._spl_type == self.VAR_SPL:
      checksum = sum(ord(x) for x in data[8:])
      checksum_offset = 4
    else:
      raise CmdError('SPL type not set')

    return data[:checksum_offset]+ struct.pack(
      '<L', checksum) + data[checksum_offset + 4:]

  def _UpdateHash(self, data, digest):
    """Update the BL2 hash.

    The BL2 header may have a pointer to the hash block, but if not, then we
    add it (at the end of SPL).

    Args:
      data: The data to update.
      digest: The hash digest to write.

    Returns:
      The new contents of the BL2 data, after updating the hash.

    Raises:
      CmdError if spl type is not variable size. We don't support this
          function with fixed-sized SPL.
    """
    if self._spl_type != self.VAR_SPL:
      raise CmdError('Hash is only supported for variable-size SPL')

    # See if there is already a hash there.
    hash_offset = struct.unpack('<L', data[8:12])[0]
    if not hash_offset:
      hash_offset = len(data)
    algo = 'sha256'.ljust(HASH_ALGO_LEN, '\x00')
    hash_block = algo + digest
    hash_block_len = len(hash_block) + HASH_HEADER_LEN
    hash_hdr = struct.pack('<4L', HASH_SIGNATURE, HASH_VERSION, hash_block_len,
                           HASH_FLAGS)
    data = (data[:hash_offset] + hash_hdr + hash_block +
            data[hash_offset + hash_block_len:])

    # Update the size and hash_offset.
    data = struct.pack('<LLL', len(data), 0, hash_offset) + data[12:]
    self._out.Info('  Added hash: %s' % ''.
                   join(['%02x' % ord(d) for d in digest]))
    return data

  def _VerifyBl2(self, data, loose_check):
    """Verify BL2 integrity.

    Fixed size and variable size SPL have different formats. Determine format,
    verify SPL integrity and save its type (fixed or variable size) for future
    reference.

    Args:
      data: The BL2 data to update.
      loose_check: a Boolean, if true - the variable size SPL blob could be
                   larger than the size value in the header
    Raises:
      CmdError if SPL blob is of unrecognizable format.
    """

    # Variable size format is more sophisticated, check it first.
    try:
      size = struct.unpack('<I', data[:4])[0]
      if size == len(data) or (loose_check and (size < len(data))):
        check_sum = sum(ord(x) for x in data[8:size])
        # Compare with header checksum
        if check_sum == struct.unpack('<I', data[4:8])[0]:
          # this is a variable size SPL
          self._out.Progress('Variable size BL2 detected')
          self._spl_type = self.VAR_SPL
          return

      # This is not var size spl, let's see if it's the fixed size one.
      # Checksum is placed at a fixed location in the blob, as defined in
      # tools/mkexynosspl.c in the u--boot tree. There are two possibilities
      # for blob sizes - 14K or 30K. The checksum is put in the last 4 bytes
      # of the blob.
      #
      # To complicate things further the blob here could have come not from
      # mkexynosspl directly, it could have been pulled out of a previously
      # bundled image. I that case it the blob will be in a chunk aligned to
      # the closest 16K boundary.
      blob_size = ((len(data) + 0x3fff) & ~0x3fff) - 2 * 1024
      if blob_size == len(data) or (loose_check and (blob_size < len(data))):
        check_sum = sum(ord(x) for x in data[:blob_size - 4])
        if check_sum == struct.unpack('<I', data[blob_size - 4:blob_size])[0]:
          self._spl_type = self.FIXED_SPL
          self._out.Progress('Fixed size BL2 detected')
          return
    except IndexError:
      # This will be thrown if bl2 is too small
      pass
    raise CmdError('Unrecognizable bl2 format')

  def Configure(self, fdt, spl_load_offset, spl_load_size, orig_bl2, name='',
                loose_check=False, digest=None, use_efs_memory=True,
                skip_sdram_init=False):
    """Configure an Exynos BL2 binary for our needs.

    We create a new modified BL2 and return its file name.

    Args:
      fdt: Device tree containing the parameter values.
      spl_load_offset: Offset in boot media that SPL must start loading (bytes)
      spl_load_size: Size of U-Boot image that SPL must load
      orig_bl2: Filename of original BL2 file to modify.
      name: a string, suffix to add to the generated file name
      loose_check: if True - allow var size SPL blob to be larger, then the
                   size value in the header. This is necessary for cases when
                   SPL is pulled out of an image (and is padded).
      digest: If not None, hash digest to add to this BL2 (a string of bytes).
      use_efs_memory: True to return the address in EFS memory (i.e. SRAM),
          False to use SDRAM
      skip_sdram_init: True to skip SDRAM initialization.

    Returns:
      Filename of configured bl2.

    Raises:
      CmdError if machine parameter block could not be found.
    """
    bl2 = os.path.join(self._tools.outdir, 'updated-spl%s.bin' % name)
    self._out.Info('Configuring BL2 %s' % bl2)
    data = self._tools.ReadFile(orig_bl2)
    self._VerifyBl2(data, loose_check)

    # Locate the parameter block
    marker = struct.pack('<L', 0xdeadbeef)
    pos = data.rfind(marker)
    if not pos:
      raise CmdError("Could not find machine parameter block in '%s'" %
                     orig_bl2)
    data = self._UpdateParameters(fdt, spl_load_offset, spl_load_size, data,
                                  pos, use_efs_memory, skip_sdram_init)
    if digest:
      data = self._UpdateHash(data, digest)
    data = self._UpdateChecksum(data)

    self._tools.WriteFile(bl2, data)
    return bl2

  def MakeSpl(self, pack, fdt, blob, vanilla_bl2):
    """Create a configured SPL based on the supplied vanilla one.

    This handles the process of working out what to configure in a SPL
    and also doing it. The settings of what to configure are in the
    flash map node which requested this SPL to be included. For example
    a 'compress' property sets the type of compresion to use for the
    payload that SPL loads.

    Args:
      pack: The PackFirmware object, providing access to the contents.
      fdt: The device tree containing the flash map.
      blob: A namedtuple with these members:
          node   - Full path to device tree node
          key    - Key name (e.g. exynos-bl2')
          params - List of parameters to the node - the first element is the
                   list of files within the blob, for example 'boot,dtb'
      vanilla_bl2: The original SPL that needs configuring.

    Returns:
      Filename of the configured SPL.

    Raises:
      CmdError if there are no parameters provided (and therefore no payload).
    """
    spl_payload = blob.params

    if not spl_payload:
      raise CmdError('No parameters provided for Exynos SPL/BL2')
    prop_list = spl_payload[0].split(',')
    name = blob.key.split('.')

    # At this stage name may be plain 'exynos-bl2', but it is possible to have
    # different versions, named 'exynos-bl2.rw' for RW SPL, and
    # 'exynos-bl2.rec' for recovery SPL (in fact anything else can be used).
    # This logic selects the string to append to the standard name, i.e. we
    # want '.rw' or '.rec', or an empty string if there is no suffix.
    if len(name) > 1:
      name = '.' + name[1]
    else:
      name = ''
    compress = fdt.GetString(blob.node, 'compress', 'none')
    if compress == 'none':
      compress = None
    data = pack.ConcatPropContents(prop_list, compress, False)[0]
    spl_load_size = len(data)

    # Figure out what flash region SPL needs to load.
    payload = fdt.GetString(blob.node, 'payload', 'none')
    if payload == 'none':
      payload = '/flash/ro-boot'
      self._out.Warning("No payload boot media for '%s' - using %s" %
                        (blob.node, payload))
    spl_load_offset = fdt.GetIntList(payload, 'reg', 2)[0]

    # Tell this SPL to use EFS memory (i.e. SRAM, if available) if we are
    # loading ro-boot. Otherwise we are loading normal U-Boot, so will use
    # SDRAM.
    use_efs_memory = 'ro-boot' in prop_list

    self._out.Info("BL2/SPL contains '%s', size is %d / %#x" %
                   (', '.join(prop_list), spl_load_size, spl_load_size))
    if fdt.GetBool(blob.node, 'hash-target'):
      hasher = hashlib.sha256()
      hasher.update(data)
      digest = hasher.digest()
    else:
      digest = None
    skip_sdram_init = fdt.GetBool(blob.node, 'skip-sdram-init')
    bl2 = self.Configure(fdt, spl_load_offset, spl_load_size, vanilla_bl2,
                         name=name, digest=digest,
                         use_efs_memory=use_efs_memory,
                         skip_sdram_init=skip_sdram_init)
    return bl2