# Copyright 2010 United States Government as represented by the
# Administrator of the National Aeronautics and Space Administration.
#
# Copyright 2011, Piston Cloud Computing, Inc.
#
# All Rights Reserved.
#
#    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.
"""
Utility methods to resize, repartition, and modify disk images.

Includes injection of SSH PGP keys into authorized_keys file.

"""

import os
import random
import tempfile

if os.name != 'nt':
    import crypt

from oslo.config import cfg
from oslo.serialization import jsonutils

from nova import exception
from nova.i18n import _
from nova.i18n import _LE
from nova.i18n import _LW
from nova.openstack.common import log as logging
from nova.openstack.common import processutils
from nova import utils
from nova.virt.disk.mount import api as mount
from nova.virt.disk.vfs import api as vfs
from nova.virt import images


LOG = logging.getLogger(__name__)

disk_opts = [
    # NOTE(yamahata): ListOpt won't work because the command may include a
    #                 comma. For example:
    #
    #                 mkfs.ext3 -O dir_index,extent -E stride=8,stripe-width=16
    #                           --label %(fs_label)s %(target)s
    #
    #                 list arguments are comma separated and there is no way to
    #                 escape 
such commas.
    #
    cfg.MultiStrOpt('virt_mkfs',
                    default=[],
                    help='Name of the mkfs commands for ephemeral device. '
                         'The format is <os_type>=<mkfs command>'),

    cfg.B
oolOpt('resize_fs_using_block_device',
                default=False,
                help='Attempt to resize the filesystem by accessing the '
                     'image over a block device. This is done by the host '
                     'and may not be necessary if the image contains a recent '
                     'version of cloud-init. Possible mechanisms require '
                     'the nbd driver (for qcow and raw), or loop (for raw).'),
    ]

CONF = cfg.CONF
CONF.register_opts(disk_opts)
CONF.import_opt('default_ephemeral_format', 'nova.virt.driver')

_MKFS_COMMAND = {}
_DEFAULT_MKFS_COMMAND = None

FS_FORMAT_EXT2 = "ext2"
FS_FORMAT_EXT3 = "ext3"
FS_FORMAT_EXT4 = "ext4"
FS_FORMAT_XFS = "xfs"
FS_FORMAT_NTFS = "ntfs"
FS_FORMAT_VFAT = "vfat"

_DEFAULT_FS_BY_OSTYPE = {'linux': FS_FORMAT_EXT3,
                         'windows': FS_FORMAT_NTFS}

for s in CONF.virt_mkfs:
    # NOTE(yamahata): mkfs command may includes '=' for its options.
    #                 So item.partition('=') doesn't work here
    os_type, mkfs_command = s.split('=', 1)
    if os_type:
        _MKFS_COMMAND[os_type] = mkfs_command
    if os_type == 'default':
        _DEFAULT_MKFS_COMMAND = mkfs_command


def get_fs_type_for_os_type(os_type):
    return os_type if _MKFS_COMMAND.get(os_type) else 'default'


def mkfs(os_type, fs_label, target, run_as_root=True, specified_fs=None):
    """Format a file or block device using
       a user provided command for each os type.
       If user has not provided any configuration,
       format type will be used according to a
       default_ephemeral_format configuration
       or a system defaults.
    """

    mkfs_command = (_MKFS_COMMAND.get(os_type, _DEFAULT_MKFS_COMMAND) or
                    '') % {'fs_label': fs_label, 'target': target}
    if mkfs_command:
        utils.execute(*mkfs_command.split(), run_as_root=run_as_root)
    else:
        if not specified_fs:
            specified_fs = CONF.default_ephemeral_format
            if not specified_fs:
                specified_fs = _DEFAULT_FS_BY_OSTYPE.get(os_type, 'ext3')

        utils.mkfs(specified_fs, target, fs_label, run_as_root=run_as_root)


def resize2fs(image, check_exit_code=False, run_as_root=False):
    try:
        utils.execute('e2fsck',
                      '-fp',
                      image,
                      check_exit_code=[0, 1, 2],
                      run_as_root=run_as_root)
    except processutils.ProcessExecutionError as exc:
        LOG.debug("Checking the file system with e2fsck has failed, "
                  "the resize will be aborted. (%s)", exc)
    else:
        utils.execute('resize2fs',
                      image,
                      check_exit_code=check_exit_code,
                      run_as_root=run_as_root)


def get_disk_size(path):
    """Get the (virtual) size of a disk image

    :param path: Path to the disk image
    :returns: Size (in bytes) of the given disk image as it would be seen
              by a virtual machine.
    """
    return images.qemu_img_info(path).virtual_size


def extend(image, size, use_cow=False):
    """Increase image to size."""
    if not can_resize_image(image, size):
        return

    utils.execute('qemu-img', 'resize', image, size)

    # if we can't access the filesystem, we can't do anything more
    if not is_image_partitionless(image, use_cow):
        return

    def safe_resize2fs(dev, run_as_root=False, finally_call=lambda: None):
        try:
            resize2fs(dev, run_as_root=run_as_root, check_exit_code=[0])
        except processutils.ProcessExecutionError as exc:
            LOG.debug("Resizing the file system with resize2fs "
                      "has failed with error: %s", exc)
        finally:
            finally_call()

    # NOTE(vish): attempts to resize filesystem
    if use_cow:
        if CONF.resize_fs_using_block_device:
            # in case of non-raw disks we can't just resize the image, but
            # rather the mounted device instead
            mounter = mount.Mount.instance_for_format(
                image, None, None, 'qcow2')
            if mounter.get_dev():
                safe_resize2fs(mounter.device,
                               run_as_root=True,
                               finally_call=mounter.unget_dev)
    else:
        safe_resize2fs(image)


def can_resize_image(image, size):
    """Check whether we can resize the container image file."""
    LOG.debug('Checking if we can resize image %(image)s. '
              'size=%(size)s', {'image': image, 'size': size})

    # Check that we're increasing the size
    virt_size = get_disk_size(image)
    if virt_size >= size:
        LOG.debug('Cannot resize image %s to a smaller size.',
                  image)
        return False
    return True


def is_image_partitionless(image, use_cow=False):
    """Check whether we can resize contained file system."""
    LOG.debug('Checking if we can resize filesystem inside %(image)s. '
              'CoW=%(use_cow)s', {'image': image, 'use_cow': use_cow})

    # Check the image is unpartitioned
    if use_cow:
        try:
            fs = vfs.VFS.instance_for_image(image, 'qcow2', None)
            fs.setup()
            fs.teardown()
        except exception.NovaException as e:
            LOG.debug('Unable to mount image %(image)s with '
                      'error %(error)s. Cannot resize.',
                      {'image': image,
                       'error': e})
            return False
    else:
        # For raw, we can directly inspect the file system
        try:
            utils.execute('e2label', image)
        except processutils.ProcessExecutionError as e:
            LOG.debug('Unable to determine label for image %(image)s with '
                      'error %(error)s. Cannot resize.',
                      {'image': image,
                       'error': e})
            return False

    return True


class _DiskImage(object):
    """Provide operations on a disk image file."""

    tmp_prefix = 'openstack-disk-mount-tmp'

    def __init__(self, image, partition=N