generate_factory_images

#!/usr/bin/env python
"""

Generate archive for fastboot factory installations.

Archive will have 3 scripts for Linux and Windows each:

    flash-base: Provision the non-AOSP parts of the device, including
    the partition table, bootloader, and Fastboot implementation

    flash-all: Same as 'flash-base' but also run 'fastboot update' to
    provision the AOSP pieces

    flash-ci: Special script for continuous integration, runs
    'fastboot boot' into the included Fastboot image, and then
    runs fastboot update from there. Does not disturb any of the
    installed non-AOSP components.

"""

from sys import exit, stderr
from argparse import ArgumentParser
from tempfile import gettempdir
from tarfile import DIRTYPE, TarInfo, open as TarOpen
from hashlib import sha1
from os import path, chmod, rename, remove
from time import time
from subprocess import check_call
import re
import tempfile
import os

_FLASHALL_FILENAME = "flash-all.sh"
# chmod (octal) -rwxr-x--x
_PERMS = 0o751
_FLASH_HEADER = """#!/bin/bash

# Copyright 2012 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

set -e

"""
_FLASH_FOOTER = "\n"
_FASTBOOT_REBOOT_COMMAND = "reboot-bootloader"
_WIN_FLASHALL_FILENAME = "flash-all.bat"
_WIN_FLASH_HEADER = """@ECHO OFF
:: Copyright 2012 The Android Open Source Project
::
:: Licensed under the Apache License, Version 2.0 (the "License");
:: you may not use this file except in compliance with the License.
:: You may obtain a copy of the License at
::
::      http://www.apache.org/licenses/LICENSE-2.0
::
:: Unless required by applicable law or agreed to in writing, software
:: distributed under the License is distributed on an "AS IS" BASIS,
:: WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
:: See the License for the specific language governing permissions and
:: limitations under the License.

PATH=%PATH%;\"%SYSTEMROOT%\System32\"
"""
_WIN_FLASH_FOOTER = """
echo Press any key to exit...
pause >nul
exit

"""
_LINUX_PREAMBLE   = 'fastboot "$@" '
_WINDOWS_PREAMBLE = 'fastboot %* '


def Checksum(filename):
    """Find SHA-1 checksum"""
    with open(filename) as fid:

        # Create sha-1 checksum
        h = sha1()
        h.update(fid.read())

        fid.close()

        return h.hexdigest()

def ConvertToDOSFormat(filename):
    """Convert to DOS file format"""
    try:
        check_call(["unix2dos", filename])
    except OSError:
        exit("Please install dos2unix package on your workstation")


def AddFlashScript(filename, tar, commands, windows):
    print("Archiving", filename)
    with tempfile.NamedTemporaryFile(delete=False, mode='w+') as tf:
        if (windows):
            tf.write(_WIN_FLASH_HEADER)
        else:
            tf.write(_FLASH_HEADER)

        for c in commands:
            if windows:
                tf.write(c.get_windows_command())
            else:
                tf.write(c.get_linux_command())

        if (windows):
            tf.write(_WIN_FLASH_FOOTER)
        else:
            tf.write(_FLASH_FOOTER)

        tf_name = tf.name
    if (windows):
        ConvertToDOSFormat(tf_name)
    chmod(tf_name, _PERMS)
    tar.add(tf_name, arcname=path.basename(filename))
    os.unlink(tf_name)

def RequireFile(filename):
    """Ensure file exists"""
    if not path.exists(filename):
        raise Usage("Cannot find " + filename)


class Usage(Exception):
    def __init__(self, msg):
        self.msg = msg


class CommandlineParser(ArgumentParser):
    """Enhanced argument parser for command line"""
    def __init__(self):
        super(CommandlineParser, self).__init__()
        self.description = __doc__

    def error(self, message):
        print("ERROR: {}".format(message), file=stderr)
        print("\n------\n", file=stderr)
        self.print_help()
        exit(2)

class Command:
    def get_linux_command(self):
        return _LINUX_PREAMBLE + self.get_command()

    def get_windows_command(self):
        return _WINDOWS_PREAMBLE + self.get_command()

class RebootBootloaderCommand(Command):
    def __init__(self, delay):
        self.delay = delay

    def get_linux_command(self):
        return '{} reboot-bootloader\nsleep {}\n'.format(_LINUX_PREAMBLE, self.delay)

    def get_windows_command(self):
        return '{} reboot-bootloader\nping -n {} 127.0.0.1 >nul\n'.format(_WINDOWS_PREAMBLE,
                self.delay)

class FlashCommand(Command):
    def __init__(self, target, image):
        self.target = target
        self.image = path.basename(image)

    def get_command(self):
        return "flash {} {}\n".format(self.target, self.image)

class UpdateCommand(Command):
    def __init__(self, image, erase):
        self.image = path.basename(image)
        self.args = ""
        if erase:
            self.args += " -w"

    def get_command(self):
        return "{} update {}\n".format(self.args, self.image)


def main():
    """Main"""
    parser = CommandlineParser()
    parser.add_argument('--product', required=True,
                        help="Product name, e.g. hammerhead")
    parser.add_argument('--release', required=True,
                        help="Release name, e.g. krt16m")
    parser.add_argument('-b', '--bootloader', required=True,
                        help="Bootloader image for fastboot flash")
    parser.add_argument('-u', '--update-archive', required=True,
                        help="Zipped images for fastboot update")
    parser.add_argument('-s', '--sleeptime',
                        type=int, default=5,
                        help="Sleep in seconds for fastboot reboot, default=5")
    parser.add_argument('-o', '--output',
                        help="Output file for archived factory scripts")
    parser.add_argument('-r', '--radio',
                        help="Radio image")
    args = parser.parse_args()

    commands = []
    files = []

    buildprefix = args.product + "-" + args.release + "-"
    update_fn = buildprefix + "img.zip"

    fn = buildprefix + "bootloader"
    files.append((args.bootloader, fn))
    commands.append(FlashCommand("bootloader", fn))

    commands.append(RebootBootloaderCommand(args.sleeptime))

    if args.radio:
        fn = buildprefix + "radio.img"
        files.append((args.radio, fn))
        commands.append(FlashCommand("radio", fn))
        commands.append(RebootBootloaderCommand(args.sleeptime))

    # Archive Name
    archive_name = args.output

    # Create Archive
    print("Creating archive: " + archive_name)
    tar = TarOpen(archive_name, "w:gz")

    for src_path, dst_path in files:
        print("Archiving " + src_path)
        RequireFile(src_path)
        tar.add(src_path, arcname=dst_path)

    # 'fastboot update' covers the additional AOSP pieces, add this to the
    # command list now
    commands.append(UpdateCommand(update_fn, True))
    print("Archiving " + args.update_archive)
    RequireFile(args.update_archive)
    tar.add(args.update_archive, update_fn)
    AddFlashScript(_FLASHALL_FILENAME, tar, commands, windows=False)
    AddFlashScript(_WIN_FLASHALL_FILENAME, tar, commands, windows=True)

    tar.close()

    print("Done.")

if __name__ == "__main__":
    try:
        exit(main())
    except Usage as err:
        print("ERROR: {}".format(err.msg), file=stderr)
        print("       for help use --help", file=stderr)
        exit(2)