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

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

"""Utilities and helper functions."""

import argparse
import glob
import os
import shlex
import signal
import subprocess
import sys
import time
import tempfile
import logging
import subprocess

from workloadmgr import exception
from workloadmgr import flags
from workloadmgr.openstack.common.gettextutils import _
from workloadmgr.openstack.common import excutils
from workloadmgr.openstack.common import importutils
from workloadmgr.openstack.common import lockutils
from workloadmgr.openstack.common import log as logging
from workloadmgr.openstack.common import timeutils


LOG = logging.getLogger(__name__)
ISO_TIME_FORMAT = "%Y-%m-%dT%H:%M:%S"
PERFECT_TIME_FORMAT = "%Y-%m-%dT%H:%M:%S.%f"
FLAGS = flags.FLAGS
_IS_NEUTRON_ATTEMPTED = False
_IS_NEUTRON = False

def _subprocess_setup():
    # Python installs a SIGPIPE handler by default. This is usually not what
    # non-Python subprocesses expect.
    signal.signal(signal.SIGPIPE, signal.SIG_DFL)


def execute(*cmd, **kwargs):
    """Helper method to execute command with optional retry.

    If you add a run_as_root=True command, don't forget to add the
    corresponding filter to etc/workloadmgr/rootwrap.d !

    :param cmd:                Passed to subprocess.Popen.
    :param process_input:      Send to opened process.
    :param check_exit_code:    Single bool, int, or list of allowed exit
                               codes.  Defaults to [0].  Raise
                               exception.ProcessExecutionError unless
                               program exits with one of these code.
    :param delay_on_retry:     True | False. Defaults to True. If set to
                               True, wait a short amount of time
                               before retrying.
    :param attempts:           How many times to retry cmd.
    :param run_as_root:        True | False. Defaults to False. If set to True,
                               the command is prefixed by the command specified
                               in the root_helper FLAG.

    :raises exception.Error: on receiving unknown arguments
    :raises exception.ProcessExecutionError:

    :returns: a tuple, (stdout, stderr) from the spawned process, or None if
             the command fails.
    """

    process_input = kwargs.pop('process_input', None)
    check_exit_code = kwargs.pop('check_exit_code', [0])
    ignore_exit_code = False
    if isinstance(check_exit_code, bool):
        ignore_exit_code = not check_exit_code
        check_exit_code = [0]
    elif isinstance(check_exit_code, int):
        check_exit_code = [check_exit_code]
    delay_on_retry = kwargs.pop('delay_on_retry', True)
    attempts = kwargs.pop('attempts', 1)
    run_as_root = kwargs.pop('run_as_root', False)
    shell = kwargs.pop('shell', False)

    if len(kwargs):
        raise exception.Error(_('Got unknown keyword args '
                                'to utils.execute: %r') % kwargs)

    if run_as_root:

        if FLAGS.rootwrap_config is not None or FLAGS.root_helper != 'sudo':
            LOG.deprecated(_('The root_helper option (which lets you specify '
                             'a root wrapper different from workloadmgr-rootwrap, '
                             'and defaults to using sudo) is now deprecated. '
                             'You should use the rootwrap_config option '
                             'instead.'))

        if (FLAGS.rootwrap_config is not None):
            cmd = ['sudo', '%s/bin/workloadmgr-rootwrap'%(EXEC_PREFIX),
                   FLAGS.rootwrap_config] + list(cmd)
        else:
            cmd = shlex.split(FLAGS.root_helper) + list(cmd)
    cmd = list(map(str, cmd))

    while attempts > 0:
        attempts -= 1
        try:
            LOG.debug(_('Running cmd (subprocess): %s'), ' '.join(cmd))
            _PIPE = subprocess.PIPE  # pylint: disable=E1101
            obj = subprocess.Popen(cmd,
                                   stdin=_PIPE,
                                   stdout=_PIPE,
                                   stderr=_PIPE,
                                   close_fds=True,
                                   preexec_fn=_subprocess_setup,
                                   shell=shell)
            result = None
            if process_input is not None:
                result = obj.communicate(bytes(process_input, 'utf-8'))
            else:
                try:
                    result = obj.communicate(timeout=30)
                except subprocess.TimeoutExpired:
                    if not obj.poll():
                        return ("Successfully started the process", None)
            obj.stdin.close()  # pylint: disable=E1101

            _returncode = obj.returncode  # pylint: disable=E1101
            if _returncode:
                LOG.debug(_('Result was %s') % _returncode)
                if not ignore_exit_code and _returncode not in check_exit_code:
                    (stdout, stderr) = result
                    raise exception.ProcessExecutionError(
                        exit_code=_returncode,
                        stdout=stdout,
                        stderr=stderr,
                        cmd=' '.join(cmd))
            return (str(result[0], encoding='utf-8'), str(result[1], encoding='utf-8'))
        except exception.ProcessExecutionError:
            if not attempts:
                raise
            else:
                LOG.debug(_('%r failed. Retrying.'), cmd)
                if delay_on_retry:
                    time.sleep(random.SystemRandom().randint(20, 200) / 100.0)
        finally:
            # NOTE(termie): this appears to be necessary to let the subprocess
            #               call clean something up in between calls, without
            #               it two execute calls in a row hangs the second one
            time.sleep(0)


def list_devices():
    devices = []
    for dev in os.listdir('/dev'):
        if dev.startswith('nbd'):
            # see if it's a partition
            if len(dev.split('p')) == 1:
                devices.append(dev)
    return devices

def find_available():
    devices = []
    busy = []
    for x in os.popen("cat /proc/partitions | grep nbd | awk '{print $4}'").readlines():
        busy.append(x.strip())

    for d in list_devices():
        if not d in busy:
            devices.append(d)
    return devices

def connect(image_file, read_only=True, secret=None):
    image_file = os.path.realpath(image_file)

    if not os.path.exists(image_file):
        return False

    devices = find_available()
    if len(devices) == 0:
        return False

    dev = devices[0]

    if read_only:
        read_only = '-r'
    else:
        read_only = ''

    cmd = "modprobe nbd".split()
    execute(*cmd, shell=False, run_as_root=True)

    if secret:
        cmd = "qemu-nbd %s -c /dev/%s "                       \
              "--object secret,id=sec0,data=%s "                 \
              "--image-opts driver=qcow2,encrypt.format=luks,"   \
              "encrypt.key-secret=sec0,file.filename=%s"         \
              %  (read_only, dev, secret, image_file)
    else:
        cmd = "qemu-nbd %s -c /dev/%s %s" % (read_only, dev, image_file)
    cmd = cmd.split()
    execute(*cmd, shell=False, run_as_root=True)
    cmd = "partprobe /dev/%s" % (dev)
    cmd = cmd.split()
    execute(*cmd, shell=False, run_as_root=True)

    return dev

def disconnect(dev=None):
    umount(dev)

    if dev == None:
        for dev in list_devices():
            disconnect(dev)
    else:
        cmd = 'qemu-nbd -d /dev/%s' % dev
        cmd = cmd.split()
        execute(*cmd, shell=False, run_as_root=True)

def mount(dev, partition=1, path=None):
    if partition:
        full_dev_path = '/dev/%sp%s' % (dev, partition)
    else:
        full_dev_path = '/dev/%s' % (dev)

    if path == None:
        import tempfile
        path = tempfile.mkdtemp()

    cmd = 'mount %s %s' % (full_dev_path, path)
    cmd = cmd.split()
    execute(*cmd, shell=False, run_as_root=True)

    return path

def find_mount(dev=None):
    if dev == None:
        mounts = []
        for dev in list_devices():
            m = find_mount(dev)
            if m != None and m not in mounts:
                mounts.append(m)

        return mounts

    else:
        mount = None

        sys_mount = os.popen('mount | grep %s' % dev).readline().strip().split(' ')
        if len(sys_mount) > 1:
            mount = {
                'dev': sys_mount[0],
                'mount': sys_mount[2],
                'type': sys_mount[3]
            }

        return mount

def umount(dev=None):
    m = find_mount(dev)

    if dev == None:
        for x in m:
            cmd = 'umount  -l %s' % x['mount']
            cmd = cmd.split()
            execute(*cmd, shell=False, run_as_root=True)
    elif m != None:
        cmd = 'umount -l %s' % m['mount']
        cmd = cmd.split()
        execute(*cmd, shell=False, run_as_root=True)

def create_overlay_file(backup_image, tmpdirname, secret=None):
    relpath = os.path.relpath(backup_image, os.path.dirname(backup_image))
    overlay_path = os.path.join(tmpdirname, relpath)
    if secret:
        cmd = ["qemu-img", "create", "-f", "qcow2",
               "--object", "secret,id=sec0,data=%s" % secret,
               "-b", 'json:{ "encrypt.key-secret": "sec0", "driver": "qcow2", "file": { "driver": "file", "filename": "%s" }}' % backup_image,
               "-o", "encrypt.format=luks,encrypt.key-secret=sec0", overlay_path]
    else:
        cmd = ["qemu-img", "create", "-f", "qcow2", "-b", backup_image, overlay_path]

    (out, err) = execute(*cmd, shell=False, run_as_root=False)
    if err:
        raise Exception(err)

    cmd = ["qemu-img", "info", overlay_path]
    (out, err) = execute(*cmd, shell=False, run_as_root=False)
    if err:
        raise Exception(err)

    return overlay_path


def main(pattern, backup_images, secret=None):
    stats_map = {
        "st_dev": "dev",
        "st_ino": "ino",
        "st_mode": "mode",
        "st_nlink": "nlink",
        "st_uid": "uid",
        "st_gid": "gid",
        "st_rdev": "rdev",
        "st_size": "size",
        "st_blksize": "blksize",
        "st_blocks": "blocks",
        "st_atime": "atime",
        "st_mtime": "mtime",
        "st_ctime": "ctime",
    }
    dt = []
    lt_drives = []
    umount()
    disconnect()
    snapshot_info = {}
    drives_info = {}
    with tempfile.TemporaryDirectory() as tmpdirname:
        for img in backup_images:
            snapshot_name = [i for i in img.split('/') if 'snapshot_' in i]
            snapshot_id = snapshot_name[0].lstrip('snapshot_') if snapshot_name else None
            if not snapshot_id:
                print("no snapshot id is associated with path: {}".format(img))
                break
            if snapshot_id not in snapshot_info:
                snapshot_info = {snapshot_id: []}
            partition_name = [i.split('_')[-1] for i in img.split('/') if 'vm_res_id_' in i]
            if not partition_name:
                break
            partition = '/dev/{}1'.format(partition_name[0])
            #print("Searching in %s:" % img)
            overlay_path = create_overlay_file(img, tmpdirname, secret=secret)
            try:
                # Create overlay file
                dev = connect(overlay_path, read_only="", secret=secret)
                try:
                    relpath = os.path.relpath(img, os.path.dirname(img))
                    mnt_path = os.path.join(tmpdirname, relpath+"-mnt")
                    os.mkdir(mnt_path)
                    mount(dev, partition=1, path=mnt_path)
                    #cwd = os.getcwd()
                    os.chdir(mnt_path)
                    for f in glob.glob(pattern):
                        #print(f)
                        statistics = os.stat(f)
                        drive_details = {}
                        for key, value in stats_map.items():
                            drive_details[value] = getattr(statistics, key, None)
                        drives_info[f] = drive_details
                        lt_drives.append(f)
                    #os.chdir(cwd)
                    disk_info = {partition: lt_drives}
                    disk_info.update(drives_info)
                    snapshot_info[snapshot_id].append(disk_info)
                    dt.append(snapshot_info)
                    umount(dev)
                    import time
                    time.sleep(10)
                finally:
                    umount(dev)
                    import time
                    time.sleep(10)
                    disconnect(dev)
                    os.rmdir(mnt_path)

            except Exception as ex:
                LOG.exception(ex)
    return dt


# Driver Code.
if __name__ == '__main__':
    parser = argparse.ArgumentParser(prog='filesearch',
                                     description='Searches for files in backup images.')
    parser.add_argument('--pattern', required=True, help='pattern help')
    parser.add_argument('--secret', default=None, help='secret help')
    parser.add_argument('backup-image', nargs='+', help='backup-image help')

    args = parser.parse_args()
    ret = main(args.pattern, args.__dict__['backup-image'], secret=args.secret)
    print(ret)