#!/usr/bin/python
#
# Register ganeti KVM instances with systemd-machined
#
# Copyright (c) 2015 Skroutz S.A.
#
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
# * Redistributions of source code must retain the above copyright
#   notice, this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in the
#   documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Register Ganeti KVM instances with systemd-machined

This script is meant to be run as a hook or directly. When run as a hook, it
will register a newly-created, started, rebooted, migrated or failed-over KVM
instance with systemd-machined.

All instances are placed in ganeti.slice, and a transient scope unit is created
for each instance via machined.

Instance UUIDs are obtained from the KVM command contained in the instance
runtime.

"""

import os
import sys
import socket

import dbus
from uuid import UUID

# Ganeti private module location
sys.path.append('/usr/share/ganeti/default')

from ganeti.ssconf import SimpleStore
from ganeti.hypervisor import GetHypervisor
from ganeti.hypervisor.hv_kvm import _AnalyzeSerializedRuntime as load_runtime
from ganeti.constants import HT_KVM
from ganeti.errors import ConfigurationError


MACHINED_MANAGER = 'org.freedesktop.machine1.Manager'


class InstanceNotRunning(Exception):
    """Instance not running error"""
    pass


def get_kvm_instance_pid(name):
    """Get a KVM instance PID"""
    kvm = GetHypervisor(HT_KVM)

    # pylint: disable=protected-access
    _, pid, alive = kvm._InstancePidAlive(name)

    if not alive:
        raise InstanceNotRunning

    return pid


def get_kvm_instance_uuid(name):
    """Get a KVM instance UUID from the runtime

    Since the UUID is not available via ssconf, this function figures it out
    from the runtime's command arguments.
    """
    kvm = GetHypervisor(HT_KVM)

    # pylint: disable=protected-access
    cmd = load_runtime(kvm._ReadKVMRuntime(name))[0]

    # Find '-uuid' and get its parameter
    idx = cmd.index('-uuid')
    if idx >= 0:
        return UUID(cmd[idx + 1])

    return None


def register_instance(name, uuid, pid):
    """Register a single instance.

    This makes sense only for KVM and LXC instances currently
    """
    bus = dbus.SystemBus()
    machined = bus.get_object('org.freedesktop.machine1',
                              '/org/freedesktop/machine1')

    if isinstance(uuid, UUID):
        _uuid = dbus.ByteArray(uuid.bytes)
    else:
        _uuid = []

    machined.CreateMachine(name,
                           _uuid,
                           "ganeti",
                           "vm",
                           pid,
                           "",
                           [("Slice", "ganeti.slice"),
                            ("Description", "%s started via Ganeti" % name)],
                           dbus_interface=MACHINED_MANAGER)


def register_all_kvm_instances():
    """Register all KVM instances not already registered"""
    conf = SimpleStore()
    try:
        enabled_hvs = conf.GetHypervisorList()
    except ConfigurationError:
        # Cluster is probably not initialized
        sys.stderr.write("E: failed to get hypervisor list"
                         " (not part of a cluster?)\n")
        sys.exit(1)

    if not HT_KVM in enabled_hvs:
        sys.exit(0)

    kvm = GetHypervisor(HT_KVM)

    bus = dbus.SystemBus()
    machined = bus.get_object('org.freedesktop.machine1',
                              '/org/freedesktop/machine1')

    registered = [str(x[0]) for x in
                  machined.ListMachines(dbus_interface=MACHINED_MANAGER)]

    for instance in kvm.ListInstances():
        if instance in registered:
            print "Skipping %s, already registered" % instance
            continue

        try:
            pid = get_kvm_instance_pid(instance)
        except InstanceNotRunning:
            print "Skipping %s, not running" % instance
            continue

        uuid = get_kvm_instance_uuid(instance)
        print "Registering %s (%s, PID: %d)" % (instance, uuid, pid)
        register_instance(instance, uuid, pid)


def main():
    """Register a single instance or all instances

    When run as a hook, register the instance the hooks is about. Otherwise,
    register all running instances.

    """

    # Exit if PID 1 is not systemd
    if not os.path.isdir("/run/systemd/system"):
        sys.exit(0)

    # Are we being called as a hook?
    opcode = os.environ.get('GANETI_OP_CODE')
    if opcode is None:
        # No, let us just register all instances
        register_all_kvm_instances()
        sys.exit(0)

    if os.environ['GANETI_HOOKS_PHASE'] != "post":
        sys.exit(0)

    if os.environ['GANETI_HOOKS_VERSION'] != "2":
        print "Invalid hooks version: %s" % os.environ['GANETI_HOOKS_VERSION']
        sys.exit(0)

    if os.environ['GANETI_INSTANCE_HYPERVISOR'] != HT_KVM:
        sys.exit(0)

    if opcode in ('OP_INSTANCE_CREATE', 'OP_INSTANCE_STARTUP',
                  'OP_INSTANCE_REBOOT'):
        primary = os.environ['GANETI_INSTANCE_PRIMARY']
    elif opcode in ('OP_INSTANCE_MIGRATE', 'OP_INSTANCE_FAILOVER'):
        primary = os.environ['GANETI_NEW_PRIMARY']
    else:
        print "Don't know how to handle %s" % opcode
        sys.exit(0)

    if primary != socket.getfqdn():
        sys.exit(0)

    name = os.environ['GANETI_INSTANCE_NAME']
    try:
        pid = get_kvm_instance_pid(name)
    except InstanceNotRunning:
        sys.exit(0)
    uuid = get_kvm_instance_uuid(name)
    register_instance(name, uuid, pid)
    print "Registered %s with systemd-machined" % name


if __name__ == "__main__":
    main()