#!/usr/bin/env python

# Copyright (c) 2011 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.

import optparse
import os
import re
import shutil
import struct
import subprocess
import sys
import tempfile

from tools import CmdError
from tools import Tools
from fdt import Fdt
import tools

# TODO(sjg): Once we have a stable fdt, write some unit tests for this

# TODO(clchiou): Rewrite this part after official flashmap implementation is
# pulled into Chromium OS code base.

# constants imported from lib/fmap.h
FMAP_SIGNATURE = "__FMAP__"
FMAP_VER_MAJOR = 1
FMAP_VER_MINOR = 0
FMAP_STRLEN = 32

FMAP_AREA_STATIC = 1 << 0
FMAP_AREA_COMPRESSED = 1 << 1
FMAP_AREA_RO = 1 << 2

FMAP_HEADER_FORMAT = "<8sBBQI%dsH" % (FMAP_STRLEN)
FMAP_AREA_FORMAT = "<II%dsH" % (FMAP_STRLEN)

FMAP_HEADER_NAMES = (
    'signature',
    'ver_major',
    'ver_minor',
    'base',
    'image_size',
    'name',
    'nareas',
)

FMAP_AREA_NAMES = (
    'offset',
    'size',
    'name',
    'flags',
)

class ConfigError(Exception):
  """A configuration error, normally a mistake in the fdt."""
  pass


class PackError(Exception):
  """A packing error, such as a file being too large for its area."""
  pass


class Entry(dict):
  """The base class for an entry type.

  All entry types are subclasses of this class. It is basically a dictionary
  with a few methods to check that the supplied data is value.

  Subclasses should implement the following methods:

    GetData(): To return the data for this entry as a string.
    RunTools(): To run any required tools to create the data.

  Potentially in the future we could add PutData() to write data from a packed
  image file back into an entry, to allow an image file to be updated.

  Properties which we need in our dictionary:
    pack: The PackFirmware object (essentially this is our parent)
    offset: Byte offset of area.
    size: Size of area in bytes.
    name: Name of area.
    required: True if this entry is required in the image, False if not
  """
  def __init__(self, props):
    super(Entry, self).__init__(props)
    self._CheckFields(('offset', 'size', 'name'))

  def __getattr__(self, name):
    return self[name]

  def __setattr__(self, name, value):
    self[name] = value

  def _CheckFields(self, fields):
    """Check that listed fields are present.

    Args:
      fields: List of field names to check for.

    Raises:
      ConfigError: if a field is not present.
    """
    for f in fields:
      if f not in self:
        raise ConfigError('Entry %s: missing required field: %s' %
            (self['name'], f))

  def _CheckFieldsInt(self, fields):
    """Check that listed fields are present, and convert them to ints.

    Args:
      fields: List of field names to check for.

    Raises:
      ConfigError: if a field is not present or cannot be converted to an int.
    """
    self._CheckFields(fields)
    for f in fields:
      try:
        self[f] = int(self[f])
      except ValueError as str:
        raise ConfigError("Entry %s, property %s: could not convert '%s'"
            " to integer" % (self['name'], f, self[f]))

  def GetOverlap(self, entry):
    """Returns the amount by which we overlap with the supplied entry.

    Args:
      entry: Entry to check.

    Returns:
      Amount of overlap:
        0: the entries butt up together.
        <0: there is a gap bewteen the entries.
        >0: there is overlap.
    """
    a = [self.offset, self.offset + self.size]
    b = [entry.offset, entry.offset + entry.size]
    return min(a[1], b[1]) - max(a[0], b[0])

  def GetData(self):
    """Method implemented by subclasses to return data for the entry.

    Returns:
      String containing entry data.
    """
    raise PackError('class Entry does not implement GetEntry()')

  def RunTools(self, tools, out, tmpdir):
    """Method implemented by subclasses to run required tools.

    Some entry types require running external tools to create their data.
    This method provides a convenient way of doing this. The supplied
    temporary directory can be used to store files. These should ideally
    not be deleted by this method, since the user may wish to see them
    later.

    Args:
      tools: Tools object to use to run tools.
      out: Output object to send output to
      tmpdir: Temporary directory to use to create required files.
    """
    pass


class EntryFmapArea(Entry):
  """A base class for things that actually produce data for the image.

  Properties:
    flags: The FMAP flags, which is an integer made up from the bit values in
        FMAP_AREA_...
  """
  def __init__(self, props):
    super(EntryFmapArea, self).__init__(props)
    self._CheckFields(('flags',))

  def GetData(self):
    """This section is empty"""
    return ''


class EntryFmap(EntryFmapArea):
  """An entry which contains an Fmap section

  Properties:
    base: Base offset of flash device (normally 0).
    size: Size of the flash device (2MB or 4MB typically).
    nareas: Number of areas in the FMAP = number of sections.
    entries: The list of entries to put in the FMAP.
  """

  def __init__(self, props):
    super(EntryFmap, self).__init__(props)
    self._CheckFieldsInt(('ver_major', 'ver_minor'))

  def SetEntries(self, base, image_size, entries):
    self['base'] = base
    self['image_size'] = image_size
    self['nareas'] = len(entries)
    self['entries'] = entries

  def GetData(self):
    def _FormatBlob(format, names, obj):
      params = [obj[name] for name in names]
      return struct.pack(format, *params)

    self._CheckFields(('base', 'size'))
    self['signature'] = FMAP_SIGNATURE
    blob = _FormatBlob(FMAP_HEADER_FORMAT, FMAP_HEADER_NAMES, self)
    for entry in self.entries:
      blob += _FormatBlob(FMAP_AREA_FORMAT, FMAP_AREA_NAMES, entry)

    return blob


class EntryWiped(EntryFmapArea):
  """This entry will be wiped to a particular byte value.

  Properties:
    wipe_value: byte value to set area to.
  """
  def __init__(self, props):
    super(EntryWiped, self).__init__(props)
    self._CheckFields(('wipe_value',))
    if self.wipe_value:
      self.wipe_value = chr(int(self.wipe_value))
    else:
      self.wipe_value = chr(0)
    if len(self.wipe_value) != 1:
        raise ConfigError('wipe_value out of range [00:ff]: %s' %
            repr(self.wipe_value))

  def GetData(self):
    return self.wipe_value * self.size


class EntryBlobString(EntryFmapArea):
  """This entry contains a single string.

  The string is placed in the area.

  Properties:
    self.value: The string value to store.
  """
  def __init__(self, props):
    super(EntryBlobString, self).__init__(props)

  def GetData(self):
    return self.value


class EntryIfd(EntryFmapArea):
  """This entry marks the use of an Intel Firmware Descriptor.

  The entry itself covers the 'Intel' part of the firmware, containing the
  firmware descriptor and management engine. When this entry appears in a
  flash map, we strip it off the image, and use ifdtool to place the rest
  of the image into a provided skeleton file (which contains the management
  engine and a skeleton firmware descriptor).
  """
  def __init__(self, props):
    super(EntryIfd, self).__init__(props)

  def ProduceFinalImage(self, tools, out, tmpdir, image_fname):
    """Produce the final image for an Intel ME system

    Some Intel systems require that an image contains the Management Engine
    firmware, and also a firmware descriptor.

    This function takes the existing image, removes the front part of it,
    and replaces it with these required pieces using ifdtool.

    Args:
      tools: Tools object to use to run tools.
      out: Output object to send output to
      tmpdir: Temporary directory to use to create required files.
      image_fname: Output image filename
    """
    out.Progress('Setting up Intel ME')
    data = tools.ReadFile(image_fname)

    # We can assume that the ifd section is at the start of the image.
    if self.offset != 0:
      raise ConfigError('IFD section must be at offset 0 in the image')
    data = data[self.size:]
    input_fname = os.path.join(tmpdir, 'ifd-input.bin')
    tools.WriteFile(input_fname, data)
    ifd_output = os.path.join(tmpdir, 'image.ifd')

    # This works by modifying a skeleton file.
    shutil.copyfile(tools.Filename(self.pack.props['skeleton']), ifd_output)
    args = ['-i', 'BIOS:%s' % input_fname, ifd_output]
    tools.Run('ifdtool', args)

    # ifdtool puts the output in a file with '.new' tacked on the end.
    shutil.move(ifd_output + '.new', image_fname)
    tools.OutputSize('IFD image', image_fname)

class EntryBlob(EntryFmapArea):
  """This entry contains a binary blob.

  Properties:
    self.value: The filename to read to obtain the blob data.
  """
  def __init__(self, props):
    super(EntryBlob, self).__init__(props)

  def GetData(self,):
    filename = self.value
    data = self.pack.tools.ReadFile(filename)
    return data


class EntryKeyBlock(EntryFmapArea):
  """This entry contains a vblock

  Properties:
    keyblock: The filename of the keyblock to use (within the keydir)
        (e.g. 'firmware.keyblock')
    signprivate: Private key filename (e.g. 'firmware_data_key.vbprivk')
    keynelkey: Kernel key filename (e.g. 'kernel_subkey.vbpubk')
    version: Version of vblock (generally 1)
  """
  def __init__(self, props):
    super(EntryKeyBlock, self).__init__(props)
    self._CheckFields(('keyblock', 'signprivate', 'kernelkey'))
    self._CheckFieldsInt(('version','preamble_flags'))

  def RunTools(self, tools, out, tmpdir):
    """Create a vblock for the given firmware image"""
    self.path = os.path.join(tmpdir, 'vblock.%s' % self.label)
    try:
      prefix = self.pack.props['keydir'] + '/'
      args = [
          '--vblock', self.path,
          '--keyblock', prefix + self.keyblock,
          '--signprivate', prefix + self.signprivate,
          '--version', '%d' % self.version,
          '--fv', self.value,
          '--kernelkey', prefix + self.kernelkey,
          '--flags', '%d' % self.preamble_flags,
      ]
      stdout = tools.Run('vbutil_firmware', args)
      out.Debug(stdout)

      # Update value to the actual filename to be used
      self.value = self.path
    except CmdError as err:
      raise PackError('Cannot make key block: vbutil_firmware failed\n%s' %
                      err)

  def GetData(self):
    fd = open(self.path, 'rb')
    data = fd.read()
    fd.close()
    return data


class PackFirmware:
  """Class for packing a firmware image

  A firmware image consists of a number of areas, called entries here.
  Together they cover the entire firmware device from start to finish.

  We use the fdt to define the format of the output - things like the
  name, start and size of each entry, as well as what type of thing
  to pack in there.

  We maintain a list of available data files (in self.props) and these can be
  requested by each entry. Each entry is a subclass of Entry, such as
  EntryBlob which handles the 'blob' entry type.

  One of the entry types is the fmap (flash map). It basically has a header
  and a representation of the list of sections. We create this from the fdt
  auotmatically. You can find this format documented here:

      http://code.google.com/p/flashmap/wiki/FmapSpec


  The following methods must be called in order:

    SetupFiles
    SelectFdt
    PackImage
  """
  def __init__(self, tools, output):
    """
    Args:
      output: A method to call to diplay output:
                def output(msg):
                Args:
                  msg: Message to output.
      tools: A tools object for calling external tools.
      verbose: 0 for silent, 1 for progress only, 2 for some info, 3 for
          full info, 4 for debug info including sub-tool output.
    """
    self.props = {}             # Properties / files we know about.
    self.entries = []           # The entries in the flash image.
    self._out = output
    self.tools = tools

  def _GetFlags(self, props):
    """Create the fmap flags value from the given properties.

    Args:
      props: Dictionary containing properties.

    Returns
      Integer containing flags from _FMAP_AREA_...
    """
    flags = FMAP_AREA_STATIC
    if 'read-only' in props:
      flags |= FMAP_AREA_RO
    if 'compressed' in props:
      flags |= FMAP_AREA_COMPRESSED
    return flags

  def _CreateEntry(self, node, props):
    """Create an entry based on a node in the fdt.

    For example this fdt line:

      type = "blob signed";

    means it is a 'blob' type and will create an EntryBlob, and the blob
    used will be self.props['signed'].

    Args:
      node: The node to use, e.g. /flash/ro-stub
      props: Doctionary containing all properties from the node.

    Returns:
      An entry set up and ready for packing.

    Raises:
      ValueError: If fdt has an unknown entry type.
    """
    entry_list = self.fdt.GetString(node, 'type', 'empty').split()
    ftype = entry_list[0]
    if ftype == 'empty':
      ftype = ''
    key = None
    if len(entry_list) > 1:
      key = entry_list[1]

    # Create an entry of the correct type.
    entry = None
    if not ftype:
      entry = EntryFmapArea(props)
    elif ftype == 'blob':
      entry = EntryBlob(props)
      pass
    elif ftype == 'wiped':
      entry = EntryWiped(props)
      pass
    elif ftype == 'keyblock':
      entry = EntryKeyBlock(props)
      pass
    elif ftype == 'blobstring':
      entry = EntryBlobString(props)
    elif ftype == 'fmap':
      entry = EntryFmap(props)
    elif ftype == 'ifd':
      entry = EntryIfd(props)
    else:
      raise ValueError('%s: unknown entry type' % ftype)

    entry.pack = self
    entry.node = node
    entry.key = key
    entry.required = 'required' in props
    entry.ftype = ftype
    return entry

  def _IsRequired(self, entry):
    """Check if an entry is required in the final image.

    Args:
      entry: Entry to check

    Returns:
      True if this entry must be included, False if not.
    """
    return entry.required

  def SetupFiles(self, boot, signed, gbb, fwid, keydir):
    """Set up files required for packing the firmware.

    Args:
      boot: Filename of bootloader image.
      signed: Filename of signed bootloader image.
      gbb: Filename of Google Binary Block (GBB).
      fwid: String containing the firmware ID.
      keydir: Key directory to use (containing signing keys).
    """
    self.props['boot'] = boot
    self.props['signed'] = signed
    self.props['gbb'] = gbb
    self.props['fwid'] = fwid
    self.keydir = keydir

  def _CheckOverlap(self):
    """Check that no entries overlap each other.

    We only allow section areas to overlap - anything with actual data in it
    must not overlap.

    Returns:
      True if all entries are required in the image, else False
    """
    required = filter(self._IsRequired, self.entries)
    entries = sorted(required, key=lambda e: e.offset)
    all_entries = len(self.entries) == len(required)
    for e1, e2 in zip(entries, entries[1:]):
      # Allow overlap between "pure" fmap areas, but not any of its subclasses
      # Here we exploit the fact that Entry is a new-style class
      if type(e1) is EntryFmapArea or type(e2) is EntryFmapArea:
        continue

      overlap = e1.GetOverlap(e2)
      if overlap > 0:
        raise ValueError('Flash map entries overlap by %d bytes: '
            '%s: %08x-%08x, %s: %08x-%08x' %
            (overlap, e1.label, e1.offset, e1.offset + e1.size,
             e2.label, e2.offset, e2.offset + e2.size))
      elif overlap is not 0:
        self._out.Warning('Warning: Flash map has a gap of %d bytes: '
            '%s: %08x-%08x, %s: %08x-%08x' %
            (-overlap, e1.label, e1.offset, e1.offset + e1.size,
             e2.label, e2.offset, e2.offset + e2.size))

    return all_entries

  def SelectFdt(self, fdt):
    """Scan FDT and build entry objects.

    This creates a list of entry objects which we can later use to generate
    a firmware image. Each entry's properties are set up correcty but at this
    stage no data is processed. The result is a list of entries in
    self.entries.

    Args:
      fdt: fdt object containing the device tree.

    Raises:
      ConfigError if an error is detected in the fdt configuration.
    """
    self.fdt = fdt
    root = '/flash'
    self.image_size = int(fdt.GetIntList(root, 'reg', 2)[1])

    # Scan the flash map in the fdt, creating a list of Entry objects.
    re_label = re.compile('(.*)-(\w*)')
    children = fdt.GetChildren(root)

    # Current offset to use for entries with only a 'size' property.
    upto_offset = 0
    required_count = 0
    first_blob_entry = None
    for child in children:
      node = root + '/' + child
      props = fdt.GetProps(node, True)

      # Read the two cells from the node's /reg property to get entry extent.
      reg = props.get('reg', None)
      if reg:
        offset, size = fdt.DecodeIntList(node, 'reg', reg, 2)
      else:
        size = props.get('size', None)
        if not size:
          raise ValueError("Must specify either 'reg' or 'size' in flash node")
        size = int(size)
        offset = upto_offset

      props['offset'] = offset
      props['size'] = size

      # The section names must be upper case with underscores, for other tools
      props['name'] = re.sub('-', '_', props['label']).upper()
      props['flags'] = self._GetFlags(props)
      try:
        entry = self._CreateEntry(node, props)
        self.entries.append(entry)

      except ConfigError as err:
        raise ValueError('Config error: %s' % err)

      if entry.ftype:
        upto_offset = offset + size
      if entry.required:
        required_count += 1
      if entry.key == 'signed':
        first_blob_entry = entry

    # HACK: Since Tegra FDT files are not in our tree yet, but we still want
    # to use the old ones, we emulate the old behavior by marking the signed
    # entry as required, if none of the entries were marked required.
    if not required_count:
      first_blob_entry.required = True

  def GetBlobList(self):
    """Generate a list of blob types that we are going to need.

    Returns:
      List of blob type strings
    """
    blob_list = set()
    for entry in self.entries:
      if isinstance(entry, EntryBlob):
        blob_list.add(entry.key)

    return list(blob_list)

  def AddProperty(self, name, value):
    """Add a new property which can be used by the fdt.

    Args:
      name: Name of property
      value: Value of property (typically a filename)
    """
    if not value:
      raise CmdError("Cannot find value for entry property '%s'" % name)
    self.props[name] = value

  def GetProperty(self, name):
    """Get the value of a property required by the fdt.

    Args:
      name: Name of property

    Returns:
      Value of property, normally a filename string
    """
    return self.props.get(name, None)

  def CheckProperties(self):
    """Check that each entry has the properties that it needs.

    Entries with a 'key' use that to look up properties. We need to make
    sure that there is a property for that key. If not, then the entry will
    not be able to access the data it needs during the packing stage.

    Raises:
      ConfigError: The property for a required key is missing
    """
    for entry in self.entries:
      if entry.required and entry.key:
        if not self.props.get(entry.key):
          raise ConfigError("%s: Requests property '%s' but we only have %s"
              % (entry.node, entry.key, self.props))
        entry.value = self.props[entry.key]

  def RequireAllEntries(self):
    """Mark all entries as required, to produce a full image."""
    for entry in self.entries:
      entry.required = True

  def PackImage(self, tmpdir, output_path):
    """Pack the various components into a firmware image,

    You must call SetupFiles() and SelectFdt() first, to set things up. Then
    this function will create a firmware image for you.

    Args:
      output_path: Full path of file to contain the resulting image.

    Raises:
      PackError: If unable to fit something in the space available, or a
          required file or setup piece is missing.
    """
    self.tmpdir = tmpdir

    all_entries = self._CheckOverlap()

    # Set up a zeroed file of the correct size.
    with open(output_path, 'wb') as image:
      if all_entries:
        image.write('\0' * self.image_size)

      # Pack all the entriess.
      ifd = None
      for entry in self.entries:
        if not entry.required:
          continue

        # Add in the info for the fmap.
        if type(entry) == EntryFmap:
          entry.SetEntries(base=0, image_size=self.image_size,
              entries=self.entries)
        elif type(entry) == EntryIfd:
          ifd = entry

        try:
          # First run any required tools.
          entry.RunTools(self.tools, self._out, self.tmpdir)
          if 'value' in entry:
            self._out.Notice("Pack '%s' into %s" % (entry.value, entry.name))

          # Now read out the data
          data = entry.GetData()
          self._out.Debug('Entry: %s' % entry.name)
          self._out.Debug('Entry data: %s' % entry)
          self._out.Debug('Data size: %s bytes, at %#x' %
              (len(data), entry.offset))
          if len(data) > entry.size:
            raise PackError("Data for '%s' too large for area: %d/%#x >"
                " %d/%#x" % (entry.name, len(data), len(data), entry.size,
                entry.size))

          image.seek(entry.offset)
          image.write(data)

        except PackError as err:
          raise ValueError('Packing error: %s' % err)

    # If the image contain an IFD section, process it
    if ifd:
      ifd.ProduceFinalImage(self.tools, self._out, self.tmpdir, output_path)

  def _OutEntry(self, status, offset, size, name):
    """Display a flash map entry.

    Args:
      status: Status character.
      offset: Byte offset of entry.
      size: Size of entry in bytes.
      name: Name of entry.
    """
    indent = '' if status == '-' else ' '
    self._out.UserOutput('%s %08x  %08x  %s%-20s' % (status, offset, size,
        indent, name))

  def ShowMap(self):
    """Show a map of the final image."""
    self._out.UserOutput('Final Flash Map:')
    self._out.UserOutput('%s %8s  %8s  %-20s' % ('S', 'Start', 'Size', 'Name'))
    offset = 0
    for entry in self.entries:
      if not entry.ftype:
        status = '-'
      elif entry.required:
        status = 'P'
      else:
        status = '.'
      if entry.ftype and offset != entry.offset:
        self._OutEntry('!', offset, entry.offset - offset, '<gap>')
      self._OutEntry(status, entry.offset, entry.size, entry.name)
      if entry.ftype:
        offset = entry.offset + entry.size

def _Test():
  """Run any built-in tests."""
  import doctest
  doctest.testmod()

def _PackOutput(msg):
  """Helper function to write output from PackFirmware (verbose level 2).

  This is passed to PackFirmware for it to use to write output.

  Args:
    msg: Message to display.
  """
  print msg

def main():
  """Main function for pack_firmware.

  We provide a way of packing firmware from the command line using this module
  directly.
  """
  parser = optparse.OptionParser()
  parser.add_option('-v', '--verbosity', dest='verbosity', default=1,
      type='int', help='Control verbosity: 0=silent, 1=progress, 3=full, '
      '4=debug')
  parser.add_option('-k', '--key', dest='key', type='string', action='store',
      help='Path to signing key directory (default to dev key)',
      default='##/usr/share/vboot/devkeys')
  parser.add_option('-d', '--dt', dest='fdt', type='string', action='store',
      help='Path to fdt file to use (binary ,dtb)', default='u-boot.dtb')
  parser.add_option('-u', '--uboot', dest='uboot', type='string',
      action='store', help='Executable bootloader file (U-Boot)',
      default='u-boot.bin')
  parser.add_option('-S', '--signed', dest='signed', type='string',
      action='store', help='Path to signed boot binary (U-Boot + BCT + FDT)',
      default='u-boot-fdt-signed.bin')
  parser.add_option('-g', '--gbb', dest='gbb', type='string',
      action='store', help='Path to Google Binary Block file',
      default='gbb.bin')
  parser.add_option('-o', '--outdir', dest='outdir', type='string',
      action='store', help='Path to directory to use for intermediate and '
      'output files', default='out')

  (options, args) = parser.parse_args(sys.argv)

  # Set up the output directory.
  if not os.path.isdir(options.outdir):
    os.makedirs(options.outdir)

  # Get tools and fdt.
  tools = Tools(options.verbosity)
  fdt = Fdt(tools, options.fdt)

  # Pack the firmware.
  pack = PackFirmware(_PackOutput, tools, options.verbosity)
  pack.SetupFiles(boot=options.uboot, signed=options.signed, gbb=options.gbb,
      fwid=tools.GetChromeosVersion(), keydir=options.key)
  pack.SelectFdt(fdt)
  out_fname = os.path.join(options.outdir, 'image.bin')
  pack.PackImage(options.outdir, out_fname)
  print 'Output binary is %s' % out_fname

if __name__ == '__main__':
  if sys.argv[1:2] == ["--test"]:
    _Test(*sys.argv[2:])
  else:
    main()