# vim: tabstop=4 shiftwidth=4 softtabstop=4

# Copyright (c) 2013 TrilioData, Inc.
# All Rights Reserved.

"""
Handling of VM disk images.
"""
import re
import time
import os
import subprocess
from contego import utils
from distutils.spawn import find_executable

from oslo_concurrency import processutils

try:
    from oslo_log import log as logging
except ImportError:
    from nova.openstack.common import log as logging

try:
    from oslo_config import cfg
except ImportError:
    from oslo.config import cfg

from nova import utils
from contego.utils import sanitize_message

import contego.privsep.path

LOG = logging.getLogger(__name__)

contego_qemuimage_opts = [
    cfg.BoolOpt(
        "force_raw_images", default=True, help="Force backing images to raw format",
    )
]

CONF = cfg.CONF
if "force_raw_images" not in list(CONF.keys()):
    CONF.register_opts(contego_qemuimage_opts)


def to_bytes(text, default=0):
    """Try to turn a string into a number of bytes. Looks at the last
    characters of the text to determine what conversion is needed to
    turn the input text into a byte number.

    Supports: B/b, K/k, M/m, G/g, T/t, KiB, MiB, GiB, TiB,
              (or the same with b/B on the end)

    """

    BYTE_MULTIPLIERS = {
        "": 1,
        "t": 1024 ** 4,
        "g": 1024 ** 3,
        "m": 1024 ** 2,
        "k": 1024,
    }

    # Take off everything not number 'like' (which should leave
    # only the byte 'identifier' left)
    mult_key_org = text.lstrip("-1234567890.")
    mult_key = mult_key_org.strip().lower()
    mult_key_len = len(mult_key)

    if mult_key.endswith("ib"):
        mult_key = mult_key[0:-2]

    if mult_key.endswith("b"):
        mult_key = mult_key[0:-1]
    try:
        multiplier = BYTE_MULTIPLIERS[mult_key]
        if mult_key_len:
            # Empty cases shouldn't cause text[0:-0]
            text = text[0:-mult_key_len].strip()
        return int(float(text) * multiplier)
    except KeyError:
        msg = ("Unknown byte multiplier: %s") % mult_key_org
        raise TypeError(msg)
    except ValueError:
        return default


class QemuImgInfo(object):
    BACKING_FILE_RE = re.compile(
        (r"^(.*?)\s*\(actual\s+path\s*:" r"\s+(.*?)\)\s*$"), re.I
    )
    TOP_LEVEL_RE = re.compile(r"^([\w\d\s\_\-]+):(.*)$")
    SIZE_RE = re.compile(r"\(\s*(\d+)\s+bytes\s*\)", re.I)

    def __init__(self, cmd_output=None):
        details = self._parse(cmd_output)
        self.image = details.get("image")
        self.backing_file = details.get("backing_file")
        self.file_format = details.get("file_format", "raw")
        self.virtual_size = details.get("virtual_size")
        self.cluster_size = details.get("cluster_size")
        self.disk_size = details.get("disk_size")
        self.snapshots = details.get("snapshot_list", [])
        self.encryption = details.get("encrypted")

    def __str__(self):
        lines = [
            "image: %s" % self.image,
            "file_format: %s" % self.file_format,
            "virtual_size: %s" % self.virtual_size,
            "disk_size: %s" % self.disk_size,
            "cluster_size: %s" % self.cluster_size,
            "backing_file: %s" % self.backing_file,
        ]
        if self.snapshots:
            lines.append("snapshots: %s" % self.snapshots)
        return "\n".join(lines)

    def _canonicalize(self, field):
        # Standardize on underscores/lc/no dash and no spaces
        # since qemu seems to have mixed outputs here... and
        # this format allows for better integration with python
        # - ie for usage in kwargs and such...
        field = field.lower().strip()
        for c in (" ", "-"):
            field = field.replace(c, "_")
        return field

    def _extract_bytes(self, details):
        # Replace it with the byte amount
        real_size = self.SIZE_RE.search(details)
        if real_size:
            details = real_size.group(1)
        try:
            details = to_bytes(details)
        except (TypeError, ValueError):
            pass
        return details

    def _extract_details(self, root_cmd, root_details, lines_after):
        consumed_lines = 0
        real_details = root_details
        if root_cmd == "backing_file":
            # Replace it with the real backing file
            backing_match = self.BACKING_FILE_RE.match(root_details)
            if backing_match:
                real_details = backing_match.group(2).strip()
        elif root_cmd in ["virtual_size", "cluster_size", "disk_size"]:
            # Replace it with the byte amount (if we can convert it)
            real_details = self._extract_bytes(root_details)
        elif root_cmd == "file_format":
            real_details = real_details.strip().lower()
        elif root_cmd == "snapshot_list":
            # Next line should be a header, starting with 'ID'
            # In qemu-img version 7.2, The Snapshot_list attribute Headers
            # has some spaces , so need to trim it.
            if not lines_after or not lines_after[0].strip().startswith("ID"):
                msg = "Snapshot list encountered but no header found!"
                raise ValueError(msg)
            consumed_lines += 1
            possible_contents = lines_after[1:]
            real_details = []
            # This is the sprintf pattern we will try to match
            # "%-10s%-20s%7s%20s%15s"
            # ID TAG VM SIZE DATE VM CLOCK (current header)
            for line in possible_contents:
                line_pieces = line.split(None)
                if len(line_pieces) != 6:
                    break
                else:
                    # Check against this pattern in the final position
                    # "%02d:%02d:%02d.%03d"
                    date_pieces = line_pieces[5].split(":")
                    if len(date_pieces) != 3:
                        break
                    real_details.append(
                        {
                            "id": line_pieces[0],
                            "tag": line_pieces[1],
                            "vm_size": line_pieces[2],
                            "date": line_pieces[3],
                            "vm_clock": line_pieces[4] + " " + line_pieces[5],
                        }
                    )
                    consumed_lines += 1
        return (real_details, consumed_lines)

    def _parse(self, cmd_output):
        # Analysis done of qemu-img.c to figure out what is going on here
        # Find all points start with some chars and then a ':' then a newline
        # and then handle the results of those 'top level' items in a separate
        # function.
        #
        # TODO(harlowja): newer versions might have a json output format
        #                 we should switch to that whenever possible.
        #                 see: http://bit.ly/XLJXDX
        if not cmd_output:
            cmd_output = ""
        contents = {}
        lines = cmd_output.splitlines()
        i = 0
        line_am = len(lines)
        while i < line_am:
            line = lines[i]
            if not line.strip():
                i += 1
                continue
            consumed_lines = 0
            top_level = self.TOP_LEVEL_RE.match(line)
            if top_level:
                root = self._canonicalize(top_level.group(1))
                if not root:
                    i += 1
                    continue
                root_details = top_level.group(2).strip()
                details, consumed_lines = self._extract_details(
                    root, root_details, lines[i + 1 :]
                )
                contents[root] = details
            i += consumed_lines + 1
        return contents


def copy_to_temp(src, dst):
    out, err = processutils.execute(
        "env", "LC_ALL=C", "LANG=C", "cp", src, dst, run_as_root=False
    )
    if err == "":
        return True
    return False


def qemu_img_info(path):
    """Return an object containing the parsed output from qemu-img info."""
    try:
        out, err = processutils.execute(
            "env",
            "LC_ALL=C",
            "LANG=C",
            "qemu-img",
            "info",
            "--image-opts",
            "file.locking=off,file.filename=%s" % path,
            run_as_root=False,
        )
        if err:
            raise Exception(err)
        return QemuImgInfo(out)
    except Exception as ex:
        LOG.warning(ex)
        time.sleep(30)
        try:
            out, err = processutils.execute(
                "env",
                "LC_ALL=C",
                "LANG=C",
                "qemu-img",
                "info",
                "--image-opts",
                "file.locking=off,file.filename=%s" % path,
                run_as_root=False,
            )
            if err:
                raise Exception(err)
            return QemuImgInfo(out)
        except Exception as ex:
            # qemu-img version < 2.10 do not have file.locking flag
            LOG.warning(ex)
            out, err = processutils.execute(
                "env",
                "LC_ALL=C",
                "LANG=C",
                "qemu-img",
                "info",
                path,
                run_as_root=False,
            )
            if err:
                raise Exception(err)

            return QemuImgInfo(out)


def qemu_check_n_resize(path, vsize):
    """Return an object containing the parsed output from qemu-img info."""
    try:
        out, err = processutils.execute(
            "env",
            "LC_ALL=C",
            "LANG=C",
            "qemu-img",
            "check",
            "-r",
            "all",
            path,
            run_as_root=False,
        )
    except Exception as ex:
        LOG.exception(ex)
        # Sometimes it fails to get "consistent read" lock.
        time.sleep(30)
        out, err = processutils.execute(
            "env",
            "LC_ALL=C",
            "LANG=C",
            "qemu-img",
            "check",
            "-r",
            "all",
            path,
            run_as_root=False,
        )

    resize_image(path, vsize)


def convert_image(source, dest, out_format, run_as_root=False):
    """Convert image to other format."""
    cmd = ("env", "LC_ALL=C", "LANG=C", "qemu-img", "convert") + (
        "-O",
        out_format,
        source,
        dest,
    )
    try:
        processutils.execute(*cmd, run_as_root=run_as_root)
    except processutils.ProcessExecutionError as ex:
        if hasattr(ex, 'cmd'):
            ex.cmd = sanitize_message(ex.cmd)
        LOG.exception(ex)
        raise ex
    except Exception as ex:
        LOG.exception(ex)
        raise ex


def rebase_qcow2(backing_file_base, backing_file_top, run_as_root=False):
    """rebase the backing_file_top to backing_file_base using unsafe mode
    :param backing_file_base: backing file to rebase to
    :param backing_file_top: top file to rebase
    """
    base_details = qemu_img_info(backing_file_base)
    try:
        processutils.execute(
            "qemu-img",
            "rebase",
            "-u",
            "-F",
            base_details.file_format,
            "-b",
            backing_file_base,
            backing_file_top,
            run_as_root=run_as_root,
        )
    except Exception as ex:
        LOG.exception("qemu-img Errored. Retrying: %s", ex)
        time.sleep(30)
        processutils.execute(
            "qemu-img",
            "rebase",
            "-u",
            "-F",
            base_details.file_format,
            "-b",
            backing_file_base,
            backing_file_top,
            run_as_root=run_as_root,
        )


def commit_qcow2(backing_file_top, secret_uuid=None, run_as_root=False):
    """rebase the backing_file_top to backing_file_base
     :param backing_file_top: top file to commit from to its base
     qemu-img commit --object secret,id=sec0,da
     ta=123456789 --image-opts driver=qcow2,encrypt.key-secret=sec0,encrypt.format=luks,fil
     e.filename=/var/trilio/triliovault-mounts/MTkyLjE2OC4xLjM0Oi9tbnQvdHZhdWx0LzQyNDM2/wor
     kload_9bb70b37-1693-45fd-b649-ef0ae3442bf8/snapshot_00754809-2485-476d-b0de-5e62581593
     f4/vm_id_4e56a99d-2eae-470f-8331-3e617a3de5f3/vm_res_id_52f0a95e-797e-45b4-a8d7-9bfa56
     2b270d_vdb/6e52ffe5-5920-4218-87df-8eb29b211f06 -p
    """
    try:
        exit_code = 0
        ex_obj = None
        if not secret_uuid:
            processutils.execute("qemu-img", "commit", "-f", "qcow2", backing_file_top, run_as_root=run_as_root)
        else:
            processutils.execute(
                        'qemu-img',
                        'commit',
                        '--object',
                        'secret,id=sec0,data={0}'.format(secret_uuid),
                        '--image-opts',
                        'driver=qcow2,encrypt.key-secret=sec0,encrypt.format=luks,file.filename={0}'.format(backing_file_top),
                        '-p',
                        run_as_root=run_as_root
                    )
    except processutils.ProcessExecutionError as ex:
        if hasattr(ex, 'cmd'):
            ex.cmd = sanitize_message(ex.cmd)
        LOG.exception(ex)
        if hasattr(ex, 'exit_code'):
            exit_code = int(ex.exit_code)
            ex_obj = ex
    except Exception as ex:
        if hasattr(ex, 'exit_code'):
            exit_code = int(ex.exit_code)
            ex_obj = ex

    if exit_code and exit_code in [1, 13]:
        try:
            user_id = str(os.getuid())
            processutils.execute(
                "chown",
                user_id + ":" + user_id,
                backing_file_top,
                run_as_root=False,
            )
            qemuinfo = qemu_img_info(backing_file_top)
            processutils.execute(
                "chown",
                user_id + ":" + user_id,
                qemuinfo.backing_file,
                run_as_root=False,
            )
        except processutils.ProcessExecutionError as ex:
            if hasattr(ex, 'cmd'):
                ex.cmd = sanitize_message(ex.cmd)
            LOG.exception(ex)
        except Exception as ex:
            LOG.exception(ex)
            try:
                user_id = os.getuid()
                contego.privsep.path.chown(backing_file_top, uid=user_id, gid=user_id)
                qemuinfo = qemu_img_info(backing_file_top)
                contego.privsep.path.chown(qemuinfo.backing_file, uid=user_id, gid=user_id)
            except Exception as ex:
                LOG.exception(ex)
        if not secret_uuid:
            processutils.execute(
                "qemu-img", "commit", "-f", "qcow2", backing_file_top, run_as_root=run_as_root
            )
        else:
            try:
                processutils.execute(
                        'qemu-img',
                        'commit',
                        '--object',
                        'secret,id=sec0,data={0}'.format(secret_uuid),
                        '--image-opts',
                        'driver=qcow2,encrypt.key-secret=sec0,encrypt.format=luks,file.filename={0}'.format(backing_file_top),
                        '-p',
                        run_as_root=run_as_root
                    )
            except processutils.ProcessExecutionError as ex:
                if hasattr(ex, 'cmd'):
                    ex.cmd = sanitize_message(ex.cmd)
                LOG.exception(ex)
                raise ex
            except Exception as ex:
                LOG.exception(ex)
                raise ex
    else:
        if ex_obj:
            raise ex_obj



def resize_image(path, new_size, run_as_root=False):
    """rebase the backing_file_top to backing_file_base
     :param backing_file_top: top file to commit from to its base
    """
    processutils.execute("qemu-img", "resize", "-f", "qcow2", path, new_size, run_as_root=run_as_root)


def create_cow_image(backing_file, path, size=None, payload=None, preallocation=None):
    """Create COW image

    Creates a COW image with the given backing file

    :param backing_file: Existing image on which to base the COW image
    :param path: Desired location of the COW image
    """
    base_cmd = ["qemu-img", "create", "-f", "qcow2"]
    cow_opts = []
    if backing_file:
        # Verify if qemu-img version is equal or above 2.10
        # to append -F as an option
        cow_opts += ["backing_file=%s" % backing_file]
        base_details = qemu_img_info(backing_file)

        try:
            version = processutils.execute("qemu-img",
                    "--version")[0].split('version')[1].split()[0]
            version = tuple(version.split('.'))
            if version >= ('2','10','0'):
                base_cmd += ["-F", base_details.file_format]
        except Exception as ex:
            LOG.exception(ex)
    else:
        base_details = None
    # Explicitly inherit the value of 'cluster_size' property of a qcow2
    # overlay image from its backing file. This can be useful in cases
    # when people create a base image with a non-default 'cluster_size'
    # value or cases when images were created with very old QEMU
    # versions which had a different default 'cluster_size'.
    if base_details and base_details.cluster_size is not None:
        cow_opts += ["cluster_size=%s" % base_details.cluster_size]
    if size is not None:
        cow_opts += ["size=%s" % size]
    if cow_opts:
        # Format as a comma separated list
        csv_opts = ",".join(cow_opts)
        cow_opts = ["-o", csv_opts]
    if preallocation:
        cow_opts += ["-o", "preallocation=metadata"]
    cmd = base_cmd + cow_opts + [path]
    if payload:
        objectdef = ['--object', 'secret,id=sec0,data={0}'.format(payload)]
        cow_opts = ['-o', "encrypt.format=luks,encrypt.key-secret=sec0"]
        if backing_file:
            backing_opts = ['-b', 'json:{ "encrypt.key-secret": "sec0", "driver": "qcow2", "file": { "driver": "file", "filename":"'+backing_file+'"}}']
            cmd = base_cmd + objectdef + backing_opts + cow_opts + [path]
        else:
            cmd = base_cmd + objectdef + cow_opts + [path] + [str(size)]
    # TODO: add here
    try:
        processutils.execute(*cmd, attempts=3)
    except processutils.ProcessExecutionError as ex:
        if hasattr(ex, 'cmd'):
            ex.cmd = sanitize_message(ex.cmd)
        LOG.exception(ex)
        raise ex
    except Exception as ex:
        LOG.exception(ex)
        raise ex



def shrink_target_encrypted_backup_qcow2_disk(cmdspec):
    try:
        stdout, stderr = processutils.execute(*cmdspec, attempts=3)
        if stderr == "":
            LOG.info("Successfully Shrinked Target Encrypted Qcow2 Image.")
    except processutils.ProcessExecutionError as ex:
        if hasattr(ex, 'cmd'):
            ex.cmd = sanitize_message(ex.cmd)
        LOG.exception("Error In Shrinking Target Disk: {0}".format(ex))
    except Exception as ex:
        LOG.exception("Error In Shrinking Target Disk: {0}".format(ex))

def get_disk_backing_file(path, basename=True, format=None):
    """Get the backing file of a disk image

    :param path: Path to the disk image
    :returns: a path to the image's backing store
    """
    backing_file = qemu_img_info(path).backing_file
    if backing_file and basename:
        backing_file = os.path.basename(backing_file)

    return backing_file

def qemu_integrity_check(disk, params=None):
    """Validate integrity of disk and it's backing chain."""
    try:
        qemu_img_exec = find_executable('qemu-img')
        img_info = qemu_img_info(disk)
        if params:
            secret_uuid = params.get('secret_uuid')
        status = None
        def _create_encrypted_disk_check_cmd(data, disk, backing_file=None):
            disk_check_cmd = [
                    qemu_img_exec,
                    'check',
                    "--object",
                    "secret,id=sec0,data={0}".format(data),
                    "--image-opts",
                    "driver=qcow2,encrypt.key-secret=sec0,file.filename={0}".format(disk),
                    ]

            return disk_check_cmd


        if hasattr(img_info, 'encryption') and img_info.encryption == 'yes':
            try:
                status = subprocess.call(_create_encrypted_disk_check_cmd(data=secret_uuid, disk=disk))
            except Exception as Ex:
                LOG.info('Got Exception in Disk Check:- {0} Retrying...'.format(Ex))
                status = subprocess.call(_create_encrypted_disk_check_cmd(data=secret_uuid, disk=disk))
        else:
            status = subprocess.call([qemu_img_exec, 'check', disk])

        if status is None:
            LOG.warning("Received none during disk check, \
                         failed to check disk integrity")
            return None
        
        # In case of success with return code 0, procceding further 
        # with backing chain check
        if status == 0:
            backing_chain_status = subprocess.call([qemu_img_exec, 'info',
                                                    '--backing-chain', disk])
            if backing_chain_status:
                LOG.exception("Backing chain is broken for the disk: %s" \
                              % (disk))
                return False
        else:
            LOG.exception("Disk integrity check failed for disk: %s" %(disk))
            return False
        return True
    except Exception as ex:
        LOG.exception(ex)
        return False