build-uboot-legacy-image

#!/usr/bin/env python

import sys
from struct      import Struct
from argparse    import ArgumentParser, FileType, Action
from collections import namedtuple
from zlib        import crc32
from time        import time
from shutil      import copyfileobj

IH_MAGIC = 0x27051956

class HexAction(Action):
    def __call__(self, parser, namespace, v, option_string = None):
        setattr(namespace,
                self.dest,
                int(v, 16 if v.startswith('0x') else 10))

parser = ArgumentParser(description='Make a U-Boot legacy image')
parser.add_argument('image',
                    type = FileType('rb'),
                    help = 'the image to wrap up')
parser.add_argument('-o', '--output',
                    type = FileType('wb'),
                    help = 'output file - stdout is used if not present')
parser.add_argument('-t', '--template',
                    type = FileType('rb'),
                    help = 'file name of a U-Boot legacy image (or just ' \
                           'header) to use as a template for header values')
parser.add_argument('--header-only',
                    action = 'store_true',
                    help = 'only print header, not image data')
parser.add_argument('--preserve-magic',
                    action = 'store_true',
                    help = 'preserve magic value from template')
parser.add_argument('--preserve-time',
                    action = 'store_true',
                    help = 'preserve timestamp from template')
parser.add_argument('--load-address',
                    action = HexAction,
                    help = 'address can be decimal or prefixed hexadecimal')
parser.add_argument('--entry-point', action = HexAction)
parser.add_argument('--os')
parser.add_argument('--arch')
parser.add_argument('--image-type')
parser.add_argument('--compression')
parser.add_argument('--name')

args = parser.parse_args()

Header = namedtuple('Header',
                    'magic header_crc time size load_address entry_point ' \
                    'data_crc os arch image_type compression name')
header_format = Struct('!IIIIIIIBBBB32s')

h = None
if(args.template):
    ih = args.template.read(header_format.size)
    args.template.close()

    if len(ih) < header_format.size:
        sys.exit('Error: Template file is too small')

    h = Header._make(header_format.unpack(ih))._asdict()
    if not args.preserve_magic:
        if(h['magic'] != IH_MAGIC):
            sys.exit('Error: Incorrect magic number in template file')
else:
    h = Header(*map(lambda _: None, Header._fields)).asdict()

# Set values from arguments.
for f in Header._fields:
    av = vars(args).get(f)
    if av is not None:
        h[f] = av

# Magic!
if not args.template:
    h['magic'] = IH_MAGIC

# Current time...
if not args.template or (args.template and not args.preserve_time):
    h['time'] = int(time())

# Read and calculate CRC32 for image.
h['data_crc'] = crc32('')
while True:
    a = args.image.read(1024 * 16)
    if len(a) == 0:
        break
    h['data_crc'] = crc32(a, h['data_crc']) & 0xffffffff

# Grab the file size while we're at the end of the image.
h['size'] = args.image.tell()

# Pack header and calculate CRC32. The field should be blank when the value
# is calculated, so we zero it first.
h['header_crc'] = 0
h['header_crc'] = crc32(header_format.pack(*h.values())) & 0xffffffff

output = args.output if args.output else sys.stdout

# Print header.
output.write(header_format.pack(*h.values()))

# Print image data.
if not args.header_only:
    args.image.seek(0)
    copyfileobj(args.image, output)