#!/usr/bin/env python

# Copyright 2011 OpenStack, LLC
#
# 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.

"""vm_vdi_cleaner.py - List or clean orphaned VDIs/instances on XenServer."""

import doctest
import os
import sys
import XenAPI


possible_topdir = os.getcwd()
if os.path.exists(os.path.join(possible_topdir, "nova", "__init__.py")):
        sys.path.insert(0, possible_topdir)


from nova import context
from nova import db
from nova import exception
from nova import flags
from nova.openstack.common import cfg
from nova.openstack.common import timeutils
from nova.virt.xenapi import driver as xenapi_driver


CONF = cfg.CONF
flags.DECLARE("resize_confirm_window", "nova.compute.manager")


ALLOWED_COMMANDS = ["list-vdis", "clean-vdis", "list-instances",
                    "clean-instances", "test"]


def call_xenapi(xenapi, method, *args):
    """Make a call to xapi"""
    return xenapi._session.call_xenapi(method, *args)


def find_orphaned_instances(xenapi):
    """Find and return a list of orphaned instances."""
    ctxt = context.get_admin_context(read_deleted="only")

    orphaned_instances = []

    for vm_ref, vm_rec in _get_applicable_vm_recs(xenapi):
        try:
            uuid = vm_rec['other_config']['nova_uuid']
            instance = db.api.instance_get_by_uuid(ctxt, uuid)
        except (KeyError, exception.InstanceNotFound):
            # NOTE(jk0): Err on the side of caution here. If we don't know
            # anything about the particular instance, ignore it.
            print_xen_object("INFO: Ignoring VM", vm_rec, indent_level=0)
            continue

        # NOTE(jk0): This would be triggered if a VM was deleted but the
        # actual deletion process failed somewhere along the line.
        is_active_and_deleting = (instance.vm_state == "active" and
                instance.task_state == "deleting")

        # NOTE(jk0): A zombie VM is an instance that is not active and hasn't
        # been updated in over the specified period.
        is_zombie_vm = (instance.vm_state != "active"
                and timeutils.is_older_than(instance.updated_at,
                        CONF.zombie_instance_updated_at_window))

        if is_active_and_deleting or is_zombie_vm:
            orphaned_instances.append((vm_ref, vm_rec, instance))

    return orphaned_instances


def cleanup_instance(xenapi, instance, vm_ref, vm_rec):
    """Delete orphaned instances."""
    xenapi._vmops._destroy(instance, vm_ref)


def _get_applicable_vm_recs(xenapi):
    """An 'applicable' VM is one that is not a template and not the control
    domain.
    """
    for vm_ref in call_xenapi(xenapi, 'VM.get_all'):
        try:
            vm_rec = call_xenapi(xenapi, 'VM.get_record', vm_ref)
        except XenAPI.Failure, e:
            if e.details[0] != 'HANDLE_INVALID':
                raise
            continue

        if vm_rec["is_a_template"] or vm_rec["is_control_domain"]:
            continue
        yield vm_ref, vm_rec


def print_xen_object(obj_type, obj, indent_level=0, spaces_per_indent=4):
    """Pretty-print a Xen object.

    Looks like:

        VM (abcd-abcd-abcd): 'name label here'
    """
    if not CONF.verbose:
        return
    uuid = obj["uuid"]
    try:
        name_label = obj["name_label"]
    except KeyError:
        name_label = ""
    msg = "%(obj_type)s (%(uuid)s) '%(name_label)s'" % locals()
    indent = " " * spaces_per_indent * indent_level
    print "".join([indent, msg])


def _find_vdis_connected_to_vm(xenapi, connected_vdi_uuids):
    """Find VDIs which are connected to VBDs which are connected to VMs."""
    def _is_null_ref(ref):
        return ref == "OpaqueRef:NULL"

    def _add_vdi_and_parents_to_connected(vdi_rec, indent_level):
        indent_level += 1

        vdi_and_parent_uuids = []
        cur_vdi_rec = vdi_rec
        while True:
            cur_vdi_uuid = cur_vdi_rec["uuid"]
            print_xen_object("VDI", vdi_rec, indent_level=indent_level)
            connected_vdi_uuids.add(cur_vdi_uuid)
            vdi_and_parent_uuids.append(cur_vdi_uuid)

            try:
                parent_vdi_uuid = vdi_rec["sm_config"]["vhd-parent"]
            except KeyError:
                parent_vdi_uuid = None

            # NOTE(sirp): VDI's can have themselves as a parent?!
            if parent_vdi_uuid and parent_vdi_uuid != cur_vdi_uuid:
                indent_level += 1
                cur_vdi_ref = call_xenapi(xenapi, 'VDI.get_by_uuid',
                    parent_vdi_uuid)
                try:
                    cur_vdi_rec = call_xenapi(xenapi, 'VDI.get_record',
                            cur_vdi_ref)
                except XenAPI.Failure, e:
                    if e.details[0] != 'HANDLE_INVALID':
                        raise
                    break
            else:
                break

    for vm_ref, vm_rec in _get_applicable_vm_recs(xenapi):
        indent_level = 0
        print_xen_object("VM", vm_rec, indent_level=indent_level)

        vbd_refs = vm_rec["VBDs"]
        for vbd_ref in vbd_refs:
            try:
                vbd_rec = call_xenapi(xenapi, 'VBD.get_record', vbd_ref)
            except XenAPI.Failure, e:
                if e.details[0] != 'HANDLE_INVALID':
                    raise
                continue

            indent_level = 1
            print_xen_object("VBD", vbd_rec, indent_level=indent_level)

            vbd_vdi_ref = vbd_rec["VDI"]

            if _is_null_ref(vbd_vdi_ref):
                continue

            try:
                vdi_rec = call_xenapi(xenapi, 'VDI.get_record', vbd_vdi_ref)
            except XenAPI.Failure, e:
                if e.details[0] != 'HANDLE_INVALID':
                    raise
                continue

            _add_vdi_and_parents_to_connected(vdi_rec, indent_level)


def _find_all_vdis_and_system_vdis(xenapi, all_vdi_uuids, connected_vdi_uuids):
    """Collects all VDIs and adds system VDIs to the connected set."""
    def _system_owned(vdi_rec):
        vdi_name = vdi_rec["name_label"]
        return (vdi_name.startswith("USB") or
                vdi_name.endswith(".iso") or
                vdi_rec["type"] == "system")

    for vdi_ref in call_xenapi(xenapi, 'VDI.get_all'):
        try:
            vdi_rec = call_xenapi(xenapi, 'VDI.get_record', vdi_ref)
        except XenAPI.Failure, e:
            if e.details[0] != 'HANDLE_INVALID':
                raise
            continue
        vdi_uuid = vdi_rec["uuid"]
        all_vdi_uuids.add(vdi_uuid)

        # System owned and non-managed VDIs should be considered 'connected'
        # for our purposes.
        if _system_owned(vdi_rec):
            print_xen_object("SYSTEM VDI", vdi_rec, indent_level=0)
            connected_vdi_uuids.add(vdi_uuid)
        elif not vdi_rec["managed"]:
            print_xen_object("UNMANAGED VDI", vdi_rec, indent_level=0)
            connected_vdi_uuids.add(vdi_uuid)


def find_orphaned_vdi_uuids(xenapi):
    """Walk VM -> VBD -> VDI change and accumulate connected VDIs."""
    connected_vdi_uuids = set()

    _find_vdis_connected_to_vm(xenapi, connected_vdi_uuids)

    all_vdi_uuids = set()
    _find_all_vdis_and_system_vdis(xenapi, all_vdi_uuids, connected_vdi_uuids)

    orphaned_vdi_uuids = all_vdi_uuids - connected_vdi_uuids
    return orphaned_vdi_uuids


def list_orphaned_vdis(vdi_uuids):
    """List orphaned VDIs."""
    for vdi_uuid in vdi_uuids:
        if CONF.verbose:
            print "ORPHANED VDI (%s)" % vdi_uuid
        else:
            print vdi_uuid


def clean_orphaned_vdis(xenapi, vdi_uuids):
    """Clean orphaned VDIs."""
    for vdi_uuid in vdi_uuids:
        if CONF.verbose:
            print "CLEANING VDI (%s)" % vdi_uuid

        vdi_ref = call_xenapi(xenapi, 'VDI.get_by_uuid', vdi_uuid)
        try:
            call_xenapi(xenapi, 'VDI.destroy', vdi_ref)
        except XenAPI.Failure, exc:
            print >> sys.stderr, "Skipping %s: %s" % (vdi_uuid, exc)


def list_orphaned_instances(orphaned_instances):
    """List orphaned instances."""
    for vm_ref, vm_rec, orphaned_instance in orphaned_instances:
        if CONF.verbose:
            print "ORPHANED INSTANCE (%s)" % orphaned_instance.name
        else:
            print orphaned_instance.name


def clean_orphaned_instances(xenapi, orphaned_instances):
    """Clean orphaned instances."""
    for vm_ref, vm_rec, instance in orphaned_instances:
        if CONF.verbose:
            print "CLEANING INSTANCE (%s)" % instance.name

        cleanup_instance(xenapi, instance, vm_ref, vm_rec)


def main():
    """Main loop."""
    args = CONF(args=sys.argv,
                usage='%prog [options] [' + '|'.join(ALLOWED_COMMANDS) + ']')
    if len(args) < 2:
        CONF.print_usage()
        sys.exit(1)

    command = args[1]

    if CONF.zombie_instance_updated_at_window < CONF.resize_confirm_window:
        raise Exception("`zombie_instance_updated_at_window` has to be longer"
                " than `resize_confirm_window`.")

    xenapi = xenapi_driver.XenAPIDriver()

    if command == "list-vdis":
        if CONF.verbose:
            print "Connected VDIs:\n"
        orphaned_vdi_uuids = find_orphaned_vdi_uuids(xenapi)
        if CONF.verbose:
            print "\nOrphaned VDIs:\n"
        list_orphaned_vdis(orphaned_vdi_uuids)
    elif command == "clean-vdis":
        orphaned_vdi_uuids = find_orphaned_vdi_uuids(xenapi)
        clean_orphaned_vdis(xenapi, orphaned_vdi_uuids)
    elif command == "list-instances":
        orphaned_instances = find_orphaned_instances(xenapi)
        list_orphaned_instances(orphaned_instances)
    elif command == "clean-instances":
        orphaned_instances = find_orphaned_instances(xenapi)
        clean_orphaned_instances(xenapi, orphaned_instances)
    elif command == "test":
        doctest.testmod()
    else:
        print "Unknown command '%s'" % command
        sys.exit(1)


if __name__ == "__main__":
    main()