#!/usr/bin/env python
#
# Copyright (c) 2012-2013 Intel, Inc.
# License: GPLv2
# Author: Artem Bityutskiy <artem.bityutskiy@linux.intel.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License, version 2,
# as published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# General Public License for more details.

"""
This is a tool to generate block map files (bmap) and to copy files using bmap.
Generally speaking, these tools are about writing large image files quickly.

The bmap file is an XML file which contains a list of mapped blocks of the
image. Mapped blocks are the blocks which have disk sectors associated with
them, as opposed to holes, which are blocks with no associated disk sectors. In
other words, the image is considered to be a sparse file, and bmap basically
contains a list of mapped blocks of this sparse file. The bmap additionally
contains some useful information like block size (usually 4KiB), image size,
mapped blocks count, etc.

The bmap is used for copying the image to a block device or to a regular file.
The idea is that we copy quickly with bmap because we copy only mapped blocks
and ignore the holes, because they are useless. And if the image is generated
properly (starting with a huge hole and writing all the data), it usually
contains only little mapped blocks, comparing to the overall image size. And
such an image compresses very well (because holes are read as all zeroes), so
it is beneficial to distribute them as compressed files along with the bmap.

Here is an example. Suppose you have a 4GiB image which contains only 100MiB of
user data and you need to flash it to a slow USB stick. With bmap you end up
copying only a little bit more than 100MiB of data from the image to the USB
stick (namely, you write only mapped blocks). This is a lot faster than copying
all 4GiB of data. We say that it is a bit more than 100MiB because things like
file-system meta-data (inode tables, superblocks, etc), partition table, etc
also contribute to the mapped blocks and are also copied.
"""

# Disable the following pylint recommendations:
#   * Too few public methods (R0903)
# pylint: disable=R0903

VERSION = "2.5"

import argparse
import sys
import os
import stat
import time
import logging
import tempfile
import traceback
from bmaptools import BmapCreate, BmapCopy, BmapHelpers, TransRead

def copy_command_open_blkdev(path, log):
    """
    Open a block device specified by 'path' in exclusive mode. Returns opened
    file object.
    """

    class NamedFile:
        """
        This simple class allows us to override the 'name' attribute of a file
        object. The problem is that 'os.fdopen()' sets the name to "<fdopen>",
        which is not very user-friendly.
        """

        def __init__(self, file_obj, name):
            self._file_obj = file_obj
            self.name = name

        def __getattr__(self, name):
            return getattr(self._file_obj, name)

    try:
        descriptor = os.open(path, os.O_WRONLY | os.O_EXCL)
    except OSError as err:
        log.error("cannot open block device '%s' in exclusive mode: %s"
                  % (path, err))
        raise SystemExit(1)

    # Turn the block device file descriptor into a file object
    try:
        file_obj = os.fdopen(descriptor, "wb")
    except OSError as err:
        os.close(descriptor)
        log.error("cannot open block device '%s': %s" % (path, err))
        raise SystemExit(1)

    return NamedFile(file_obj, path)

def find_and_open_bmap(image_path):
    """
    When the user does not specify the bmap file, we try to find it at the same
    place where the image file is located. We search for a file with the same
    path and basename, but with a ".bmap" extension. Since the image may
    contain more than one extension (e.g., image.raw.bz2), try to remove them
    one-by-one.

    This function returns a file-like object for the bmap file if it has been
    found, and 'None' otherwise.
    """

    bmap_path = None

    while True:
        bmap_path = image_path + ".bmap"

        try:
            bmap_obj = TransRead.TransRead(bmap_path)
            bmap_obj.close()
            return bmap_path
        except TransRead.Error:
            pass

        image_path, ext = os.path.splitext(image_path)
        if ext == '':
            break

    return None

def copy_command_open_all(args, log):
    """
    Open the image/bmap/destination files for the "copy" command. Returns a
    tuple of 5 elements:
        1 file-like object for the image
        2 file object for the destination file
        3 file-like object for the bmap
        4 full path to the bmap file
        5 image size in bytes
        6 'True' if the destination file is a block device and 'False' otherwise
    """

    # Open the image file using the TransRead module, which will automatically
    # recognize whether it is compressed or whether file path is an URL, etc.
    try:
        image_obj = TransRead.TransRead(args.image)
    except TransRead.Error as err:
        log.error("cannot open image: %s" % str(err))
        raise SystemExit(1)

    # Open the bmap file. Try to discover the bmap file automatically if it
    # was not specified.
    bmap_path = args.bmap
    if not bmap_path and not args.nobmap:
        bmap_path = find_and_open_bmap(args.image)
        if bmap_path:
            log.info("discovered bmap file '%s'" % bmap_path)

    bmap_obj = None
    if bmap_path:
        try:
            # The current implementation of BmapCopy requires a local file for
            # the bmap file.
            bmap_obj = TransRead.TransRead(bmap_path, local = True)
        except TransRead.Error as err:
            log.error("cannot open bmap file '%s': %s" % (bmap_path, str(err)))
            raise SystemExit(1)

    # Try to open the destination file. If it does not exist, a new regular
    # file will be created. If it exists and it is a regular file - it'll be
    # truncated. If this is a block device, it'll just be opened.
    try:
        dest_obj = open(args.dest, 'wb+')
    except IOError as err:
        log.error("cannot open destination file '%s': %s"
                  % (args.dest, err))
        raise SystemExit(1)

    # Check whether the destination file is a block device
    dest_is_blkdev = stat.S_ISBLK(os.fstat(dest_obj.fileno()).st_mode)
    if dest_is_blkdev:
        dest_obj.close()
        dest_obj = copy_command_open_blkdev(args.dest, log)

    return (image_obj, dest_obj, bmap_obj, bmap_path, image_obj.size,
            dest_is_blkdev)


def copy_command(args, log):
    """Copy an image to a block device or a regular file using bmap."""
    image_obj, dest_obj, bmap_obj, bmap_path, image_size, dest_is_blkdev = \
                                          copy_command_open_all(args, log)
    try:
        if dest_is_blkdev:
            dest_str = "block device '%s'" % args.dest
            # For block devices, use the specialized class
            writer = BmapCopy.BmapBdevCopy(image_obj, dest_obj, bmap_obj,
                                           image_size, logger=log)
        else:
            dest_str = "file '%s'" % os.path.basename(args.dest)
            writer = BmapCopy.BmapCopy(image_obj, dest_obj, bmap_obj,
                                       image_size, logger=log)
    except BmapCopy.Error as err:
        log.error(str(err))
        raise SystemExit(1)

    # Print the progress indicator while copying
    if not args.quiet:
        writer.set_progress_indicator(sys.stderr, "bmaptool: info: %d%% copied")

    start_time = time.time()
    if not bmap_obj:
        if args.nobmap:
            log.info("no bmap given, copy entire image to '%s'" % args.dest)
        else:
            log.error("please, use --nobmap option to flash without bmap")
            raise SystemExit(1)
    else:
        log.info("block map format version %s" % writer.bmap_version)
        log.info("%d blocks of size %d (%s), mapped %d blocks (%s or %.1f%%)"
                 % (writer.blocks_cnt, writer.block_size,
                    writer.image_size_human, writer.mapped_cnt,
                    writer.mapped_size_human, writer.mapped_percent))
        log.info("copying image '%s' to %s using bmap file '%s'"
                 % (os.path.basename(args.image), dest_str,
                    os.path.basename(bmap_path)))

    try:
        try:
            writer.copy(False, not args.no_verify)
        except BmapCopy.Error as err:
            log.error(str(err))
            raise SystemExit(1)

        # Synchronize the block device
        log.info("synchronizing '%s'" % args.dest)
        try:
            writer.sync()
        except BmapCopy.Error as err:
            log.error(str(err))
            raise SystemExit(1)
    except KeyboardInterrupt:
        log.error("the program is interrupted, exiting")
        raise SystemExit(1)

    copying_time = time.time() - start_time
    copying_speed = writer.mapped_size / copying_time
    log.info("copying time: %s, copying speed %s/sec"
             % (BmapHelpers.human_time(copying_time),
                BmapHelpers.human_size(copying_speed)))

    dest_obj.close()
    if bmap_obj:
        bmap_obj.close()
    image_obj.close()

def create_command(args, log):
    """
    Generate block map (AKA bmap) for an image. The idea is that while images
    files may generally be very large (e.g., 4GiB), they may nevertheless
    contain only little real data, e.g., 512MiB. This data are files,
    directories, file-system meta-data, partition table, etc. When copying the
    image to the target device, you do not have to copy all the 4GiB of data,
    you can copy only 512MiB of it, which is 4 times less, so copying should
    presumably be 4 times faster.

    The block map file is an XML file which contains a list of blocks which
    have to be copied to the target device. The other blocks are not used and
    there is no need to copy them. The XML file also contains some additional
    information like block size, image size, count of mapped blocks, etc. There
    are also many commentaries, so it is human-readable.

    The image has to be a sparse file. Generally, this means that when you
    generate this image file, you should start with a huge sparse file which
    contains a single hole spanning the entire file. Then you should partition
    it, write all the data (probably by means of loop-back mounting the image
    or parts of it), etc. The end result should be a sparse file where mapped
    areas represent useful parts of the image and holes represent useless parts
    of the image, which do not have to be copied when copying the image to the
    target device.
    """

    # Create and setup the output stream
    if args.output:
        try:
            output = open(args.output, "w+")
        except IOError as err:
            log.error("cannot open the output file '%s': %s"
                      % (args.output, err))
            raise SystemExit(1)
    else:
        try:
            # Create a temporary file for the bmap
            output = tempfile.TemporaryFile("w+")
        except IOError as err:
            log.error("cannot create a temporary file: %s" % err)
            raise SystemExit(1)

    try:
        creator = BmapCreate.BmapCreate(args.image, output)
        creator.generate(not args.no_checksum)
    except BmapCreate.Error as err:
        log.error(str(err))
        raise SystemExit(1)

    if not args.output:
        output.seek(0)
        sys.stdout.write(output.read())

    if creator.mapped_cnt == creator.blocks_cnt:
        log.warning("all %s are mapped, no holes in '%s'"
                    % (creator.image_size_human, args.image))
        log.warning("was the image handled incorrectly and holes "
                    "were expanded?")

def parse_arguments():
    """A helper function which parses the input arguments."""
    text = "Create block map (bmap) and copy files using bmap. The " \
           "documentation can be found here: " \
           "source.tizen.org/documentation/reference/bmaptool"
    parser = argparse.ArgumentParser(description = text, prog = 'bmaptool')

    # The --version option
    parser.add_argument("--version", action = "version",
                        version = "%(prog)s " + "%s" % VERSION)

    # The --quiet option
    text = "be quiet"
    parser.add_argument("-q", "--quiet", action = "store_true", help = text)

    subparsers = parser.add_subparsers(title = "subcommands")

    #
    # Create the parser for the "create" command
    #
    text = "generate bmap for an image file (which should be a sparse file)"
    parser_create = subparsers.add_parser("create", help = text)
    parser_create.set_defaults(func = create_command)

    # Mandatory command-line argument - image file
    text = "the image to generate bmap for"
    parser_create.add_argument("image", help = text)

    # The --output option
    text = "the output file name (otherwise stdout is used)"
    parser_create.add_argument("-o", "--output", help = text)

    # The --no-checksum option
    text = "do not generate the checksum for block ranges in the bmap"
    parser_create.add_argument("--no-checksum", action="store_true",
                               help = text)

    #
    # Create the parser for the "copy" command
    #
    text = "write an image to a block device using bmap"
    parser_copy = subparsers.add_parser("copy", help=text)
    parser_copy.set_defaults(func=copy_command)

    # The first positional argument - image file
    text = "the image file to copy. Supported formats: uncompressed, " + \
           ", ".join(TransRead.SUPPORTED_COMPRESSION_TYPES)
    parser_copy.add_argument("image", help=text)

    # The second positional argument - block device node
    text = "the destination file or device node to copy the image to"
    parser_copy.add_argument("dest", help=text)

    # The --bmap option
    text = "the block map file for the image"
    parser_copy.add_argument("--bmap", help=text)

    # The --nobmap option
    text = "allow copying without a bmap file"
    parser_copy.add_argument("--nobmap", action="store_true", help=text)

    # The --no-verify option
    text = "do not verify the data checksum while writing"
    parser_copy.add_argument("--no-verify", action="store_true", help=text)

    return parser.parse_args()

def setup_logger(loglevel):
    """
    A helper function which sets up and configures the logger. The log level is
    initialized to 'loglevel'. Returns the logger object.
    """

    # Esc-sequences for coloured output
    esc_red = '\033[91m'    # pylint: disable=W1401
    esc_yellow = '\033[93m' # pylint: disable=W1401
    esc_end = '\033[0m'     # pylint: disable=W1401

    # Change log level names to something less nicer than the default
    # all-capital 'INFO' etc.
    logging.addLevelName(logging.ERROR, esc_red + "ERROR" + esc_end)
    logging.addLevelName(logging.WARNING, esc_yellow + "WARNING" + esc_end)
    logging.addLevelName(logging.DEBUG, "debug")
    logging.addLevelName(logging.INFO, "info")

    log = logging.getLogger('bmap-logger')
    log.setLevel(loglevel)
    formatter = logging.Formatter("bmaptool: %(levelname)s: %(message)s")
    where = logging.StreamHandler(sys.stderr)
    where.setFormatter(formatter)
    log.addHandler(where)

    return log

def main():
    """Script entry point."""
    args = parse_arguments()

    if args.quiet:
        loglevel = logging.ERROR
    else:
        loglevel = logging.INFO

    log = setup_logger(loglevel)

    try:
        args.func(args, log)
    except MemoryError:
        log.error("Out of memory!")
        traceback.print_exc()

        log.info("The contents of /proc/meminfo:")
        with open('/proc/meminfo', 'rt') as file_obj:
            for line in file_obj:
                print line,

        log.info("The contents of /proc/self/status:")
        with open('/proc/self/status', 'rt') as file_obj:
            for line in file_obj:
                print line,

if __name__ == "__main__":
    sys.exit(main())
