# Copyright 2011 Red Hat, Inc.
#
# 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.
"""Support for mounting images with qemu-nbd."""

import os
import random
import re
import subprocess
import time

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

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

from nova import utils
from oslo_concurrency import processutils
from nova.virt.disk.mount import api

LOG = logging.getLogger(__name__)

nbd_opts = [
    cfg.IntOpt(
        "nbd_timeout",
        default=10,
        help="Amount of time, in seconds, to wait for NBD device start up.",
    )
]

CONF = cfg.CONF
CONF.register_opts(nbd_opts)

NBD_DEVICE_RE = re.compile("nbd[0-9]+")


class NbdMount(api.Mount):
    """qemu-nbd support disk images."""

    mode = "nbd"

    def _detect_nbd_devices(self):
        """Detect nbd device files."""
        return list(filter(NBD_DEVICE_RE.match, os.listdir("/sys/block/")))

    def _disconnect_nbd_devices(self, dev):
        """Disconnect all nbd devices"""
        for each_dev in dev:
            LOG.debug("Release nbd device %s", each_dev)
            try:
                out, err = processutils.execute(
                    "sudo",
                    "nova-rootwrap",
                    CONF.rootwrap_config,
                    "qemu-nbd", "-d", "/dev/" + each_dev)
            except Exception as ex:
                LOG.warning(ex)
                return False
            if out or not err:
                return True
            else:
                LOG.warning("Failed to disconnect nbd device: {}, error: {}".format(each_dev, err))
                return False

    def _find_unused(self, devices):
        for device in devices:
            if not os.path.exists(os.path.join("/sys/block/", device, "pid")):
                if not os.path.exists("/var/lock/qemu-nbd-%s" % device):
                    return device
                else:
                    LOG.error(
                        "NBD error - /sys/block/{}/pid is absent "
                        "but /var/lock/qemu-nbd-{} present".format(device, device))
        LOG.warning("No free nbd devices")
        return None

    def _create_nbd_device_as_fdisk(self, nbd_dev, secret_uuid, overlay_file):
        if secret_uuid:
            cmd = ['sudo', 'nova-rootwrap', CONF.rootwrap_config,
                    'qemu-nbd', '-c', '/dev/'+nbd_dev,
                    '--object', 'secret,id=sec0,data={0}'.format(secret_uuid),
                    '--image-opts',
                    'driver=qcow2,encrypt.format=luks,encrypt.key-secret=sec0,file.filename={0}'.format(overlay_file)
                    ]
        else:
            cmd = ['sudo', 'nova-rootwrap', CONF.rootwrap_config,
                    'qemu-nbd', '-c', '/dev/%s' % (nbd_dev), overlay_file]

        try:
            _PIPE = subprocess.PIPE
            timeout = 30
            timeout_err = False
            for retry in range(3):
                out, err = None, None
                timeout_err = False
                with subprocess.Popen(cmd, stdin=_PIPE, stdout=_PIPE, stderr=_PIPE, close_fds=True) as obj:
                    try:
                        out, err = obj.communicate(timeout=timeout)
                        if not out:
                            return True
                    except Exception as ex:
                        # in case of subprocess.TimeoutExpired the exception is not
                        # being recognized, hence handling in common exception block
                        if "timed out" in str(ex).lower() or "TimeoutExpired" in str(type(ex)) or not obj.poll():
                            #sometimes out is '', err is '' which shows the nbd device is connected but in actual
                            # it's an TimeoutExpired Exception. In this case we will fire lsblk cmd to find out
                            # if the device is showing up or not.
                            lsblk_out, lsblk_err = processutils.execute('lsblk')
                            if nbd_dev in lsblk_out.split():
                                return True
                            timeout_err = True
                            timeout = timeout * 2
                            time.sleep(5)
                            self._disconnect_nbd_devices([nbd_dev])
                            LOG.warning("timeout occurred during qemu-nbd device connect execution, retrying...")
                            continue
                        else:
                            LOG.exception(ex)
                            return False
            if timeout_err:
                LOG.exception("timeout occurred during qemu-nbd device connect execution")
                return False
        except Exception as ex:
            LOG.exception(ex)
            raise

    def list_partitions(self, nbd_dev):
        devices = []
        for dev in os.listdir('/dev'):
            if dev.startswith(nbd_dev):
                if len(dev.split('p')) == 2:
                    devices.append(dev)
        return devices

    def _allocate_nbd(self):
        if not os.path.exists("/sys/block/nbd0"):
            LOG.error("nbd module not loaded")
            self.error = "nbd unavailable: module not loaded"
            return None

        devices = self._detect_nbd_devices()
        random.shuffle(devices)
        device = self._find_unused(devices)
        if not device:
            # really want to log this info, not raise
            self.error = "No free nbd devices"
            return None
        return os.path.join("/dev", device)

    @utils.synchronized("nbd-allocation-lock")
    def _inner_get_dev(self):
        device = self._allocate_nbd()
        if not device:
            return False

        # NOTE(mikal): qemu-nbd will return an error if the device file is
        # already in use.
        LOG.debug(
            "Get nbd device %(dev)s for %(imgfile)s",
            {"dev": device, "imgfile": self.image},
        )
        _out, err = utils.trycmd(
            "qemu-nbd",
            "-c",
            device,
            self.image,
            "--aio=native",
            "--nocache",
            run_as_root=True,
        )
        if err:
            self.error = "qemu-nbd error: %s" % err
            LOG.info("NBD mount error: %s", self.error)
            _out, err = utils.trycmd("qemu-nbd", "-d", device, run_as_root=True)
            if err:
                LOG.warning(
                    "Detaching from erroneous nbd device returned error: %s", err,
                )
            return False

        # NOTE(vish): this forks into another process, so give it a chance
        # to set up before continuing
        pidfile = "/sys/block/%s/pid" % os.path.basename(device)
        for _i in range(CONF.nbd_timeout):
            if os.path.exists(pidfile):
                self.device = device
                break
            else:
                self.error = "nbd device %s did not show up" % device
                LOG.info("NBD mount error: %s", self.error)

                # Cleanup
                _out, err = utils.trycmd("qemu-nbd", "-d", device, run_as_root=True)
                if err:
                    LOG.warning(
                        "Detaching from erroneous nbd device returned error: %s",
                        err,
                    )
                return False

        self.error = ""
        self.linked = True
        return True

    def get_dev(self):
        """Retry requests for NBD devices."""
        return self._get_dev_retry_helper()

    def unget_dev(self):
        if not self.linked:
            return
        LOG.debug("Release nbd device %s", self.device)
        utils.execute("qemu-nbd", "-d", self.device, run_as_root=True)
        self.linked = False
        self.device = None

    def flush_dev(self):
        """flush NBD block device buffer."""
        # Perform an explicit BLKFLSBUF to support older qemu-nbd(s).
        # Without this flush, when a nbd device gets re-used the
        # qemu-nbd intermittently hangs.
        if self.device:
            utils.execute("blockdev", "--flushbufs", self.device, run_as_root=True)