#!/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 optparse
import sys
import XenAPI

from nova import context
from nova import db
from nova import exception
from nova import flags
from nova.openstack.common import timeutils

from nova.virt.xenapi import connection as xenapi_conn


flags.DECLARE("resize_confirm_window", "nova.compute.manager")
flags.DECLARE("xenapi_connection_url", "nova.virt.xenapi.connection")
flags.DECLARE("xenapi_connection_username", "nova.virt.xenapi.connection")
flags.DECLARE("xenapi_connection_password", "nova.virt.xenapi.connection")

FLAGS = flags.FLAGS
# NOTE(sirp): Nova futzs with the sys.argv in order to provide default
# flagfile. To isolate this awful practice, we're supplying a dummy
# argument list.
dummy = ["fakearg"]
flags.parse_args(dummy)


class UnrecognizedNameLabel(Exception):
    pass


def parse_options():
    """Generate command line options."""

    ALLOWED_COMMANDS = ["list-vdis", "clean-vdis", "list-instances",
                        "clean-instances", "test"]
    arg_str = "|".join(ALLOWED_COMMANDS)
    parser = optparse.OptionParser("%prog [options] [" + arg_str + "]")
    parser.add_option("--verbose", action="store_true")

    options, args = parser.parse_args()

    if not args:
        parser.print_usage()
        sys.exit(1)

    return options, args


def find_orphaned_instances(session, verbose=False):
    """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(session):
        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,
                    verbose=verbose)
            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,
                        FLAGS.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(session, instance, vm_ref, vm_rec):
    """Delete orphaned instances."""
    network_info = None
    connection = xenapi_conn.get_connection(_)
    if vm_rec['power_state'] == 'Running':
        session.xenapi.VM.hard_shutdown(vm_ref)
    connection._vmops._destroy(instance, vm_ref, shutdown=False)


def _get_applicable_vm_recs(session):
    """An 'applicable' VM is one that is not a template and not the control
    domain.
    """
    for vm_ref in session.xenapi.VM.get_all():
        try:
            vm_rec = session.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,
                     verbose=False):
    """Pretty-print a Xen object.

    Looks like:

        VM (abcd-abcd-abcd): 'name label here'
    """
    if not 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(session, connected_vdi_uuids, verbose=False):
    """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,
                             verbose=verbose)
            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 = session.xenapi.VDI.get_by_uuid(
                    parent_vdi_uuid)
                try:
                    cur_vdi_rec = session.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(session):
        indent_level = 0
        print_xen_object("VM", vm_rec, indent_level=indent_level,
                         verbose=verbose)

        vbd_refs = vm_rec["VBDs"]
        for vbd_ref in vbd_refs:
            try:
                vbd_rec = session.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,
                             verbose=verbose)

            vbd_vdi_ref = vbd_rec["VDI"]

            if _is_null_ref(vbd_vdi_ref):
                continue

            try:
                vdi_rec = session.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(session, all_vdi_uuids, connected_vdi_uuids,
                                   verbose=False):
    """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 session.xenapi.VDI.get_all():
        try:
            vdi_rec = session.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,
                             verbose=verbose)
            connected_vdi_uuids.add(vdi_uuid)
        elif not vdi_rec["managed"]:
            print_xen_object("UNMANAGED VDI", vdi_rec, indent_level=0,
                             verbose=verbose)
            connected_vdi_uuids.add(vdi_uuid)


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

    _find_vdis_connected_to_vm(session, connected_vdi_uuids, verbose=verbose)

    all_vdi_uuids = set()
    _find_all_vdis_and_system_vdis(session, all_vdi_uuids, connected_vdi_uuids,
                                   verbose=verbose)

    orphaned_vdi_uuids = all_vdi_uuids - connected_vdi_uuids
    return orphaned_vdi_uuids


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


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

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


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


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

        cleanup_instance(session, instance, vm_ref, vm_rec)


def main():
    """Main loop."""
    options, args = parse_options()
    verbose = options.verbose
    command = args[0]

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

    session = XenAPI.Session(FLAGS.xenapi_connection_url)
    session.xenapi.login_with_password(FLAGS.xenapi_connection_username,
            FLAGS.xenapi_connection_password)

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


if __name__ == "__main__":
    main()
