blob: fc48071318389cc8955c7ef1ef3caf53e80b6ab3 [file] [log] [blame]
#!/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 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.
'tftpserverip': 'tftpserverip=%s',
# Some kernel arguments are also "first class". They are listed here.
'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.
# Format of the CRC in the output file.
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
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
def _CheckMemsize(option, opt, value):
"""Check that a 'memsize' option to optparse is good.
option: See optparse manual.
opt: Option name we're parsing; See optparse manual.
value: Value we're parsing; See optparse manual.
value: The parsed value.
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)]
# 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.
raise optparse.OptionValueError("option %s: invalid memsize value: %r" %
(opt, value))
TYPES = optparse.Option.TYPES + ('memsize',)
TYPE_CHECKER = copy.copy(optparse.Option.TYPE_CHECKER)
_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 = ""
>>> sorted(_GetSpecialEnvVars(opts))
opts: The options from the option parser.
vars: A list of 'key=value' strings for u-boot environment variables.
vars = []
for opt_name, value in OPTION_ENVIRONMENT_VARS.iteritems():
if getattr(opts, opt_name):
vars.append(OPTION_ENVIRONMENT_VARS[opt_name] % getattr(opts, opt_name))
return vars
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 = ""
>>> sorted(_GetSpecialKernelArgs(opts))
>>> opts.board = "tegra2_kaen-gobi"
>>> sorted(_GetSpecialKernelArgs(opts))
['cros_board=tegra2_kaen-gobi', 'omahaserver=']
opts: The options from the option parser.
vars: A list of strings that will eventually be joined with space to create
the kernel command line.
vars = []
for opt_name, value in OPTION_KERNEL_ARGS.iteritems():
if getattr(opts, opt_name):
vars.append(OPTION_KERNEL_ARGS[opt_name] % getattr(opts, opt_name))
return vars
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)
>>> _ParseEnvStr(s)
['merge_with_default=1', 'a=b', 'c=d', 'extra_bootargs=noinitrd x=y']
env_vars: List of env_vars to put in the environment (after the default).
env_size: Final size of the environment.
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
env_str: The string to read
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
>>> s == chr(0) * 0x1000
# No-op (just opens file); tell it to pad the size it already is.
>>> s = _MakeOutput(path1, None, 0x1000).read()
>>> len(s) == 0x1000
>>> s == chr(0) * 0x1000
# Try padding w/ a copy
>>> s = _MakeOutput(path3, path4, 0x2000).read()
>>> len(s) == 0x2000
>>> s == (chr(0xff) * 0x1000) + (chr(0) * 0x1000)
# Try padding w/out a copy
>>> s = _MakeOutput(path3, None, 0x2000).read()
>>> len(s) == 0x2000
>>> s == (chr(0xff) * 0x1000) + (chr(0) * 0x1000)
# Double-check sizes
>>> os.path.getsize(path1) == 0x1000
>>> os.path.getsize(path2) == 0x1000
>>> os.path.getsize(path3) == 0x2000
>>> os.path.getsize(path4) == 0x2000
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.
outfile: The output file. Already contains the input and has been padded.
Left seeked at 0.
if not output_path:
output_path = input_path
output_dir = os.path.dirname(output_path)
if not os.path.isdir(output_dir):
raise ArgumentError("The destination directory %s doesn't exist" %
shutil.copy(input_path, output_path)
except IOError, e:
raise ArgumentError("Problem copying input to output: %s" % str(e))
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:, os.SEEK_END)
outfile.write('\0' * padding_bytes)
return outfile
def _GetEnvVarAddr(image_file):
"""Get environment variable section address from FMAP.
image_file: The image file to look for environment variable section.
The address of RW_ENVIRONMENT section is returned.
command = ["dump_fmap", "-p",, "RW_ENVIRONMENT"]
stream = subprocess.Popen(command, stdout=subprocess.PIPE)
result = stream.communicate()[0].split()
if len(result) == 0:
raise ArgumentError("Cannot find RW_ENVIRONMENT section in FMAP.")
addr = int(result[1])
return addr
except OSError:
raise ArgumentError("Error calling dump_fmap.")
def _PutEnvInFile(outfile, 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)
>>> _PutEnvInFile(f, env_str, False)
>>> s = f.getvalue()
# File should have kept the same size.
>>> len(s) == FILE_SIZE
# File should end with env_str and start with '\0'
>>> s[:-0x100].strip(chr(0))
>>> s[-0x100:] == env_str
# Running again should get an error if clobbering not OK.
>>> _PutEnvInFile(f, 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, env_str, True)
>>> s_new = f.getvalue()
>>> s == s_new
outfile: An already opened file.
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)
addr = _GetEnvVarAddr(outfile), os.SEEK_SET)
old_env_str =
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"), os.SEEK_SET)
def _ParseOptions():
"""Parse command-line arguments.
opts: Options from optparse.OptionParser.
parser = optparse.OptionParser(description=__doc__,
parser.add_option('--env-size', dest='env_size', default=0x1000,
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,
help='Bypass ignorable errors')
# Special-case options to make things simpler for factory. Must match
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',
help='Set any arbitrary u-boot var using format arg=value')
parser.add_option('--arg', default=[], dest='args', metavar='arg',
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])
opts = _ParseOptions()
print 'Input: %s (0x%08x bytes)' % (opts.input,
except OSError:
raise ArgumentError("Error accessing input image: %s" % opts.input)
env_vars = opts.vars + _GetSpecialEnvVars(opts)
kernel_args = opts.args + _GetSpecialKernelArgs(opts)
env_str = _BuildEnvironment(env_vars, kernel_args, opts.env_size)
outfile = _MakeOutput(opts.input, opts.output, opts.fw_size)
_PutEnvInFile(outfile, env_str, opts.force)
print 'Output: %s (0x%08x bytes)' % (,
print '\n '.join(['Stored env:'] + env_vars)
except ArgumentError, e:
print >>sys.stderr, "%s: error: %s" % (prog_name, str(e))
def _Test(verbose=''):
"""Run any built-in tests."""
import doctest
assert verbose in ('', '-v')
doctest.testmod(verbose=(verbose == '-v'))
if __name__ == '__main__':
if sys.argv[1:2] == ["--test"]: