#!/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.

"""Update environment variables in a legacy_image.bin."""

import copy
import fmap
import optparse
import os
import shutil
import struct
import subprocess
import sys
import zlib


# Some environment variables are "first class"--you can set them with simple
# command line options.  Those are listed here with the mapping for how to
# change the command-line option into an environment variable.
OPTION_ENVIRONMENT_VARS = {
  'tftpserverip': 'tftpserverip=%s',
}

# Some kernel arguments are also "first class".  They are listed here.
OPTION_KERNEL_ARGS = {
  'omahaserver': 'omahaserver=%s',
  'board': 'cros_board=%s',
}

# This is the environment we'll write, which tells to merge this environment
# with whatever is stored in u-boot as the "default" environment.
DEFAULT_ENV = (
  "merge_with_default=1\0"
)

# Format of the CRC in the output file.
CRC_FORMAT = '<i'


class ArgumentError(Exception):
  """We'll raise this whenever we have an error that was caused by user-error.

  This could be a bad argument or the inability to access a file on the
  filesystem.
  """
  pass


class _OptionWithMemsize(optparse.Option):
  """Option subclass that has an additional 'memsize' option.

  The memsize option allows you to add a k/m/g suffix to your integers to
  specify KiB, MiB, or GiB.


  >>> parser = optparse.OptionParser(option_class=_OptionWithMemsize)
  >>> _ = parser.add_option('--test', type='memsize')

  # Simple: just a number of bytes...
  >>> opts, args = parser.parse_args(['--test', '999'])
  >>> print opts
  {'test': 999}

  # Can do hex number of bytes...
  >>> opts, args = parser.parse_args(['--test', '0x1000'])
  >>> print opts
  {'test': 4096}

  # Kilobyte suffix...
  >>> opts, args = parser.parse_args(['--test', '1k'])
  >>> print opts
  {'test': 1024}

  # Can handle hex and upper-case suffix too..
  >>> opts, args = parser.parse_args(['--test', '0xfK'])
  >>> print opts
  {'test': 15360}

  # Check other suffixes
  >>> opts, args = parser.parse_args(['--test', '1m'])
  >>> print opts
  {'test': 1048576}
  >>> opts, args = parser.parse_args(['--test', '1g'])
  >>> print opts
  {'test': 1073741824}

  # Negative test cases; need to quiet stderr first...
  >>> import StringIO
  >>> _oldstderr = sys.stderr
  >>> sys.stderr = StringIO.StringIO()
  >>> opts, args = parser.parse_args(['--test', 'hello'])
  Traceback (most recent call last):
      ...
  SystemExit: 2
  >>> opts, args = parser.parse_args(['--test', ''])
  Traceback (most recent call last):
      ...
  SystemExit: 2
  >>> sys.stderr = _oldstderr
  """

  @staticmethod
  def _CheckMemsize(option, opt, value): # pylint: disable=W0613
    """Check that a 'memsize' option to optparse is good.

    Args:
      option: See optparse manual.
      opt: Option name we're parsing; See optparse manual.
      value: Value we're parsing; See optparse manual.
    Returns:
      value: The parsed value.
    Raises:
      OptionValueError: Upon error
    """
    # Note: purposely no 'b' suffix, since that makes 0x12b ambiguous.
    multiplier_table = [
      ('g', 1024 * 1024 * 1024),
      ('m', 1024 * 1024),
      ('k', 1024),
      ('', 1),
    ]
    for (suffix, multiplier) in multiplier_table:
      if value.lower().endswith(suffix):
        new_value = value
        if suffix:
          new_value = new_value[:-len(suffix)]
        try:
          # Convert w/ base 0 (handles hex, binary, octal, etc)
          return int(new_value, 0) * multiplier
        except ValueError:
          # Pass and try other suffixes; not useful now, but may be useful
          # later if we ever allow B vs. GB vs. GiB.
          pass
    raise optparse.OptionValueError("option %s: invalid memsize value: %r" %
                                    (opt, value))

  TYPES = optparse.Option.TYPES + ('memsize',)
  TYPE_CHECKER = copy.copy(optparse.Option.TYPE_CHECKER)

# pylint: disable=W0212
_OptionWithMemsize.TYPE_CHECKER['memsize'] = _OptionWithMemsize._CheckMemsize


def _GetSpecialEnvVars(opts):
  """Grab env vars that the user can set directly via command-line options.

  >>> class FakeOpts(object): pass
  >>> opts = FakeOpts()
  >>> opts.tftpserverip = None

  >>> _GetSpecialEnvVars(opts)
  []

  >>> opts.tftpserverip = "1.2.3.4"
  >>> sorted(_GetSpecialEnvVars(opts))
  ['tftpserverip=1.2.3.4']

  Args:
    opts: The options from the option parser.

  Returns:
    A list of 'key=value' strings for u-boot environment variables.
  """
  variables = []
  for opt_name, dummy_value in OPTION_ENVIRONMENT_VARS.iteritems():
    if getattr(opts, opt_name):
      variables.append(OPTION_ENVIRONMENT_VARS[opt_name] %
                       getattr(opts, opt_name))
  return variables


def _GetSpecialKernelArgs(opts):
  """Grab kernel args that the user can set directly via command-line options.

  >>> class FakeOpts(object): pass
  >>> opts = FakeOpts()
  >>> opts.omahaserver = None
  >>> opts.board = None

  >>> _GetSpecialKernelArgs(opts)
  []

  >>> opts.omahaserver = "3.4.5.6"
  >>> sorted(_GetSpecialKernelArgs(opts))
  ['omahaserver=3.4.5.6']

  >>> opts.board = "tegra2_kaen-gobi"
  >>> sorted(_GetSpecialKernelArgs(opts))
  ['cros_board=tegra2_kaen-gobi', 'omahaserver=3.4.5.6']

  Args:
    opts: The options from the option parser.

  Returns:
    A list of strings that will eventually be joined with space to create
      the kernel command line.
  """
  variables = []
  for opt_name, dummy_value in OPTION_KERNEL_ARGS.iteritems():
    if getattr(opts, opt_name):
      variables.append(OPTION_KERNEL_ARGS[opt_name] % getattr(opts, opt_name))
  return variables


def _BuildEnvironment(env_vars, kernel_args, env_size):
  """Build up the u-boot environment string.

  >>> s = _BuildEnvironment(['a=b', 'c=d'], ['noinitrd', 'x=y'], 0x100)
  >>> len(s)
  256

  >>> _ParseEnvStr(s)
  ['merge_with_default=1', 'a=b', 'c=d', 'extra_bootargs=noinitrd x=y']

  Args:
    env_vars: List of env_vars to put in the environment (after the default).
    env_size: Final size of the environment.

  Returns:
    env_str: The environment string for u-boot, including checksum.
  """
  assert DEFAULT_ENV.endswith('\0')

  # Add the command-line args in through extra_bootargs
  if kernel_args:
    env_vars += ['extra_bootargs=%s' % (' '.join(kernel_args))]

  env_str = DEFAULT_ENV + '\0'.join(env_vars)

  # Pad to full size...
  padding_bytes = env_size - struct.calcsize(CRC_FORMAT) - len(env_str)
  if padding_bytes < 0:
    raise ArgumentError("Environment %d bytes too big" % -padding_bytes)
  env_str += '\0' * padding_bytes

  # Add in the CRC.
  crc = zlib.crc32(env_str)
  crc_str = struct.pack(CRC_FORMAT, crc)

  return crc_str + env_str


def _ParseEnvStr(env_str):
  """Read an environment string and return variables in it.

  # Test corrupt case; valid case is tested by _BuildEnvironment()
  >>> _ParseEnvStr('hello; this is not valid') is None
  True

  Args:
    env_str: The string to read

  Returns:
    env_vars: The variables in the string, or None if the CRC was bad.
  """
  crc_bytes = struct.calcsize(CRC_FORMAT)

  stored_crc, = struct.unpack(CRC_FORMAT, env_str[:crc_bytes])
  env_str = env_str[crc_bytes:]
  calc_crc = zlib.crc32(env_str)

  if stored_crc != calc_crc:
    return None

  return env_str.rstrip('\0').split('\0')


def _MakeOutput(input_path, output_path, fw_size):
  """Make and open the output file.

  Handles copying the input file to the output file (if the user desires that)
  and also padding.

  # Doctests; for tests below...
  #   file 1: 0x1000 of zeros
  #   file 2: should be a copy of file 1
  #   file 3: 0x1000 of 0xff, eventually padded to 0x2000 with 0s.
  #   file 4: copy of file 3 padded to 0x2000 with 0s.
  >>> import tempfile
  >>> work_dir = tempfile.mkdtemp('doctest')
  >>> path1 = os.path.join(work_dir, '1.bin')
  >>> path2 = os.path.join(work_dir, '2.bin')
  >>> path3 = os.path.join(work_dir, '3.bin')
  >>> path4 = os.path.join(work_dir, '4.bin')
  >>> open(path1, 'wb').write(chr(0) * 0x1000)
  >>> open(path3, 'wb').write(chr(0xff) * 0x1000)

  # Try making a copy: from 1 to 2
  >>> s = _MakeOutput(path1, path2, 0x1000).read()
  >>> len(s) == 0x1000
  True
  >>> s == chr(0) * 0x1000
  True

  # No-op (just opens file); tell it to pad the size it already is.
  >>> s = _MakeOutput(path1, None, 0x1000).read()
  >>> len(s) == 0x1000
  True
  >>> s == chr(0) * 0x1000
  True

  # Try padding w/ a copy
  >>> s = _MakeOutput(path3, path4, 0x2000).read()
  >>> len(s) == 0x2000
  True
  >>> s == (chr(0xff) * 0x1000) + (chr(0) * 0x1000)
  True

  # Try padding w/out a copy
  >>> s = _MakeOutput(path3, None, 0x2000).read()
  >>> len(s) == 0x2000
  True
  >>> s == (chr(0xff) * 0x1000) + (chr(0) * 0x1000)
  True

  # Double-check sizes
  >>> os.path.getsize(path1) == 0x1000
  True
  >>> os.path.getsize(path2) == 0x1000
  True
  >>> os.path.getsize(path3) == 0x2000
  True
  >>> os.path.getsize(path4) == 0x2000
  True


  Args:
    input_path: Path the the input firmware file.
    output_path: Path to the output firmware file; if None/blank we'll modify
        the input in place.
    fw_size: If non-None, we'll pad to this many bytes.

  Returns:
    outfile: The output file.  Already contains the input and has been padded.
        Left seeked at 0.
  """
  if not output_path:
    output_path = input_path
  else:
    output_dir = os.path.dirname(output_path)
    if not os.path.isdir(output_dir):
      raise ArgumentError("The destination directory %s doesn't exist" %
                          output_dir)

    try:
      shutil.copy(input_path, output_path)
    except IOError, e:
      raise ArgumentError("Problem copying input to output: %s" % str(e))

  try:
    outfile = open(output_path, 'r+b')
  except IOError, e:
    raise ArgumentError("Problem opening output file: %s" % str(e))

  if fw_size:
    old_size = os.path.getsize(output_path)
    if fw_size < old_size:
      raise ArgumentError("Can't specify firmware size smaller than input")

    padding_bytes = fw_size - old_size
    if padding_bytes:
      outfile.seek(0, os.SEEK_END)
      outfile.write('\0' * padding_bytes)
      outfile.seek(0)

  return outfile


def _GetEnvVarAddrSize(image_file):
  """Get environment variable section address and size from FMAP.

  Args:
    image_file: The image file path to look for environment variable section.

  Returns:
    A tuple of the address and the size of RW_ENVIRONMENT section is returned.
  """
  try:
    with open(image_file, 'rb') as f:
      entries = fmap.fmap_decode(f.read())
  except:
    raise ArgumentError("Error decoding FMAP from image file %s." % image_file)

  rw_env = [entry for entry in entries['areas']
            if entry['name'] == 'RW_ENVIRONMENT']
  if len(rw_env) != 1:
    raise ArgumentError("Cannot find RW_ENVIRONMENT section in FMAP.")
  return (rw_env[0]['offset'], rw_env[0]['size'])


def _PutEnvInFile(outfile, addr, env_str, clobber_ok):
  """Put the given environment into the output file.

  At the moment, this just crams the env_str to the end of the file.

  # Sample call with a fake (all zero) file.
  >>> import StringIO
  >>> FILE_SIZE=0x10000
  >>> f = StringIO.StringIO(chr(0) * FILE_SIZE)
  >>> env_str = _BuildEnvironment([], [], 0x100)
  >>> addr = _GetEnvVarAddrSize[0]
  >>> _PutEnvInFile(f, addr, env_str, False)
  >>> s = f.getvalue()

  # File should have kept the same size.
  >>> len(s) == FILE_SIZE
  True

  # File should end with env_str and start with '\0'
  >>> s[:-0x100].strip(chr(0))
  ''
  >>> s[-0x100:] == env_str
  True

  # Running again should get an error if clobbering not OK.
  >>> _PutEnvInFile(f, addr, env_str, False)
  Traceback (most recent call last):
      ...
  ArgumentError: Old arguments will be clobbered; pass --force if OK

  # Should be OK if clobbering OK
  >>> _PutEnvInFile(f, addr, env_str, True)
  >>> s_new = f.getvalue()
  >>> s == s_new
  True

  Args:
    outfile: An already opened file.
    addr: The address of RW_ENVIRONMENT section.
    env_str: The str to store.
    clobber_ok: If True, it's OK to clobber the old environment; if False we'll
        raise an exception if we detect and old environment.
  """
  env_size = len(env_str)
  outfile.seek(addr, os.SEEK_SET)

  old_env_str = outfile.read(env_size)
  if old_env_str != ('\0' * env_size):
    old_env = _ParseEnvStr(old_env_str)

    # If we get here, our old env section was neither zeros nor valid.
    # that probably means that we either had an unpadded firmware image as
    # input or a completely bogus input file.
    if old_env is None:
      raise ArgumentError("Previous vars were not zero nor valid")

    if not clobber_ok:
      raise ArgumentError("Old arguments will be clobbered; pass --force if OK")

  outfile.seek(addr, os.SEEK_SET)
  outfile.write(env_str)


def _ParseOptions():
  """Parse command-line arguments.

  Returns:
    opts: Options from optparse.OptionParser.
  """
  parser = optparse.OptionParser(description=__doc__,
                                 option_class=_OptionWithMemsize)

  parser.add_option('--env-addr', dest='env_addr', default=None,
                    type='int',
                    help='If specified, overrides the default env_addr that '
                    'u-boot expects')
  parser.add_option('--env-size', dest='env_size', default=None,
                    type='memsize',
                    help='If specified, overrides the default env_size that '
                    'u-boot expects')
  parser.add_option('--fw-size', dest='fw_size', default=None, type='memsize',
                    help='If specified, firmware will be padded to be this '
                    'many bytes')
  parser.add_option('--input', '-i', dest='input', default=None,
                    help='Path to the firmware to modify; required')
  parser.add_option('--output', '-o', dest='output', default=None,
                    help='Path to store output; if not specified we will '
                    'directly modify the input file')
  parser.add_option('--force', '-f', dest='force', default=False,
                    action='store_true',
                    help='Bypass ignorable errors')

  # Special-case options to make things simpler for factory.  Must match
  # OPTION_ENVIRONMENT_VARS and OPTION_KERNEL_ARGS
  parser.add_option('--tftpserverip', default=None,
                    help='Set the TFTP server IP address')

  parser.add_option('--board', default=None,
                    help='Set the cros_board to be passed into the kernel')
  parser.add_option('--omahaserver', default=None,
                    help='Set the Omaha server IP address')

  parser.add_option('--var', default=[], dest='vars', metavar='var',
                    action='append',
                    help='Set any arbitrary u-boot var using format arg=value')
  parser.add_option('--arg', default=[], dest='args', metavar='arg',
                    action='append',
                    help='Set any arbitrary kernel command line arg')

  opts, args = parser.parse_args()
  if args:
    raise ArgumentError('Unexpected argument(s): %s' % ', '.join(args))

  if not opts.input:
    raise ArgumentError("You must specify the input image")

  return opts


def main():
  """Main function."""
  prog_name = os.path.basename(sys.argv[0])

  try:
    opts = _ParseOptions()
    try:
      print 'Input:  %s (0x%08x bytes)' % (opts.input,
                                           os.path.getsize(opts.input))
    except OSError:
      raise ArgumentError("Error accessing input image: %s" % opts.input)

    if opts.env_addr is None or opts.env_size is None:
      env_addr, env_size = _GetEnvVarAddrSize(opts.input)
    env_addr = opts.env_addr or env_addr
    env_size = opts.env_size or env_size

    env_vars = opts.vars + _GetSpecialEnvVars(opts)
    kernel_args = opts.args + _GetSpecialKernelArgs(opts)
    env_str = _BuildEnvironment(env_vars, kernel_args, env_size)
    outfile = _MakeOutput(opts.input, opts.output, opts.fw_size)
    _PutEnvInFile(outfile, env_addr, env_str, opts.force)

    print 'Output: %s (0x%08x bytes)' % (outfile.name,
                                         os.path.getsize(outfile.name))
    print '\n  '.join(['Stored env:'] + env_vars)
  except ArgumentError, e:
    print >> sys.stderr, "%s: error: %s" % (prog_name, str(e))
    sys.exit(1)


def _Test(verbose=''):
  """Run any built-in tests."""
  import doctest # pylint: disable=W0404
  assert verbose in ('', '-v')
  doctest.testmod(verbose=(verbose == '-v'))


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