#!/usr/bin/env python3
#
# Copyright (C) 2017 Citrix, Inc.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# 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.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#
#
# A XenAPI plugin to calculate dom0 diskspace
#


import logging
import os
import re
import shutil

import xcp.cmd as cmd
import xcp.logger as logger
import XenAPIPlugin

UUID_REGEX = re.compile(
    "^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$", re.IGNORECASE
)


def host_disk_space(session, args):
    """Return the available disk space of dom0 filesystem, in bytes"""
    st = os.statvfs("/")
    return str(st.f_bfree * st.f_frsize)


def check_patch_upload(session, args):
    """Compare the disk space available in dom0 with the size provided and return true if size (in bytes) provided is less than available, false otherwise"""
    logger.logToSyslog(level=logging.INFO)
    if "size" not in args:
        logger.critical("Missing argument 'size'")
        raise Exception("MISSING_SIZE")
    size = int(args["size"])
    available_size = int(host_disk_space(None, None))
    return str(size < available_size)


def get_required_space(session, args):
    """Calculates the required space for uploading hotfix, given the hotfix size"""
    logger.logToSyslog(level=logging.INFO)
    if "size" not in args:
        logger.critical("Missing argument 'size'")
        raise Exception("MISSING_SIZE")
    return args["size"]


def get_file_size(f):
    """Return size of regular file, and 0 for everything else."""
    try:
        if os.path.isfile(f):
            return os.stat(f).st_size
        else:
            return 0
    except OSError:
        return 0


def is_this_uuid(s):
    """Checks if the string conforms to UUID format."""
    return UUID_REGEX.match(s) is not None


def get_hotfix_files(hotfix_dir):
    if os.path.exists("/var/patch"):
        for f in os.listdir(hotfix_dir):
            if is_this_uuid(f):
                yield os.path.join(hotfix_dir, f)


def get_directory_size(dir):
    """Return size of the directory"""
    if not os.path.isdir(dir):
        return 0

    (rc, out) = cmd.runCmd(["/usr/bin/du", "-sb", dir], with_stdout=True)
    if rc == 0:
        return int(out.split()[0])

    return 0


def is_this_pool_patch_uuid(session, s):
    if is_this_uuid(s):
        try:
            p_ref = session.xenapi.pool_patch.get_by_uuid(s)
            return bool(p_ref)
        except:
            return False

    return False


def get_patch_filename(session, uuid):
    if not is_this_pool_patch_uuid(session, uuid):
        return ""

    (rc, out) = cmd.runCmd(["/opt/xensource/bin/hfx_filename", uuid], with_stdout=True)
    if rc == 0:
        # This file might exists only on the master
        file_name = out.strip()
        if os.path.exists(file_name):
            return file_name

    return ""


def get_reclaimable_disk_space(session, args):
    """Return the disk-space consumed by hotfix residual files on the host."""

    if "exclude-hfx-uuid" in args:
        hfx_filename = get_patch_filename(session, args["exclude-hfx-uuid"])
    else:
        hfx_filename = ""

    sizes = (
        get_file_size(f) for f in get_hotfix_files("/var/patch") if f != hfx_filename
    )
    patch_backup_size = get_directory_size("/opt/xensource/patch-backup")
    return str(sum(sizes) + patch_backup_size)


def cleanup_disk_space(session, args):
    """Deletes the hotfix residual files."""

    logger.logToSyslog(level=logging.INFO)

    if "exclude-hfx-uuid" in args:
        hfx_filename = get_patch_filename(session, args["exclude-hfx-uuid"])
    else:
        hfx_filename = ""

    # Delete the hotfix files from /var/patch/
    if os.path.exists("/var/patch"):
        for f in os.listdir("/var/patch"):
            if is_this_uuid(f):
                f_name = os.path.join("/var/patch", f)
                if f_name != hfx_filename:
                    try:
                        logger.info("Removing " + f_name)
                        os.remove(f_name)
                    except OSError:
                        logger.critical("Failed to remove " + f_name)

    # Delete patch-backup directory if it exists
    try:
        if os.path.exists("/opt/xensource/patch-backup"):
            logger.info("Removing /opt/xensource/patch-backup directory")
            shutil.rmtree("/opt/xensource/patch-backup")
    except OSError:
        logger.critical("Failed to remove /opt/xensource/patch-backup directory")

    return ""


if __name__ == "__main__":
    XenAPIPlugin.dispatch(
        {
            "get_avail_host_disk_space": host_disk_space,
            "get_reclaimable_disk_space": get_reclaimable_disk_space,
            "check_patch_upload": check_patch_upload,
            "cleanup_disk_space": cleanup_disk_space,
            "get_required_space": get_required_space,
        }
    )
