handlers.py

#!/usr/bin/env python
#coding=utf8

import inspect
import logging
from sjrpc.utils import RpcHandler, pure
from tql import TqlQuery
from conf import CCConf

def listed(func):
    func.__listed__ = True
    return func

class CCHandler(RpcHandler):
    '''
    Base class for handlers of CloudControl server.

    :cvar role_name: the name of the role binded to the handler
    :type role_name: str or None is not applicable
    '''
    
    role_name = None

    def __init__(self, server):
        self._server = server

    def __getitem__(self, name):
        logging.debug('Called method %s.%s' % (self.__class__.__name__, name))
        return super(CCHandler, self).__getitem__(name)

    @pure
    def list_commands(self):
        cmd_list = []

        for attr in dir(self):
            attr = getattr(self, attr, None)
            if getattr(attr, '__listed__', False):
                cmd = {}
                cmd['name'] = attr.__name__
                doc = inspect.getdoc(attr)
                if doc:
                    cmd['description'] = inspect.cleandoc(doc)
                cmd_list.append(cmd)

        return cmd_list

class OnlineCCHandler(CCHandler):

    def on_disconnect(self, connection):
        self._server.unregister(connection)

class HypervisorHandler(OnlineCCHandler):
    '''
    Handler binded to 'node' role.
    '''
    role_name = 'hypervisor'


class ClientHandler(OnlineCCHandler):
    '''
    Handler binded to 'cli' role.
    '''
    
    role_name = 'client'

    def _list(self, query):

        objects = []
        query = TqlQuery(query)

        for condition in query.iterconditions():
            tag_name = condition.name
            
            if tag_name in ('a', 'hv', 'h', 'id'):
                # Append all accounts:
                for login in self._server.conf.list_accounts():
                    tags = self._server.resolve_tags(login, query.tags)
                    objects.append(tags)

            if tag_name in ('vm', 'h', 'id'):
                hvs = [o for o in objects if o['role'] == 'hypervisor'
                               and o['status'] == 'online']
                if not hvs:
                    for hy in self._server.iter_connected_role('hypervisor'):
                        tags = self._server.resolve_tags(hy.login, query.tags)
                        hvs.append(tags)

                # Query the selected hypervisors:
                logging.debug('Querying %d hypervisors: '
                              '%s' % (len(hvs), hvs))
                async_calls = {}
                tags = tuple(query.tags)
                for hy_tags in hvs:
                    hy = self._server.get_connection(hy_tags['a'])
                    cid = hy.connection.async_call('list_vm', tags=tags)
                    async_calls[cid] = hy_tags

                logging.debug('Waiting for the response of hypervisors...')
                calls = frozenset(async_calls)
                responses = self._server.manager.wait(calls, timeout=5)
                logging.debug('Responses received from hypervisors.')

                for resp in responses:
                    if resp['error'] is None:
                        hy = async_calls[resp['id']]
                        for vm_tags in resp['return']:
                            vm_tags['role'] = 'vm'
                            vm_tags['id'] = '%s.%s' % (hy['a'], vm_tags['vm'])
                            vm_tags.update(hy['__static_user'])
                            objects.append(vm_tags)

            # Filter the tag
            objects = [o for o in objects if condition.check(o.get(tag_name))]

        # Before to return, filter with the asked tags for each objects:
        tags = query.tags
        tags.add('id')
        tags.add('role')
        final_objects = {}
        for obj in objects:
            for key in obj.keys():
                if key not in tags:
                    del obj[key]
            final_objects[obj['id']] = obj

        return final_objects.values()

    @pure
    @listed
    def list(self, query):
        '''
        List all objects registered on this instance.
        '''

        logging.debug('Executed list function with query %s' % query)
        return self._list(query)

    def _vm_action(self, query, method, *args, **kwargs):
        vms = self._list(query)
        hypervisors = list(self._server.iter_connected_role('hypervisor'))
        for hv in hypervisors:
            vm_to_start = []
            for vm in vms:
                if vm['role'] != 'vm':
                    pass
                elif vm['id'].split('.')[0] == hv.login:
                    vm_to_start.append(vm['vm'])
            if vm_to_start:
                hv.connection.call(method, vm_to_start, *args, **kwargs)

    @pure
    @listed
    def start(self, query):
        self._vm_action(query, 'start_vm')

    @pure
    @listed
    def stop(self, query, force=False):
        self._vm_action(query, 'stop_vm', force)

    @pure
    @listed
    def suspend(self, query):
        self._vm_action(query, 'suspend_vm')

    @pure
    @listed
    def resume(self, query):
        self._vm_action(query, 'resume_vm')

    @pure
    @listed


class AuthenticationError(Exception):
    pass


class WelcomeHandler(CCHandler):
    '''
    Default handler used on client connections of the server.

    :cvar ROLES: role name/handler mapping
    '''
    
    ROLES = {
        'client': ClientHandler,
        'hypervisor': HypervisorHandler,
    }

    @listed
    def authentify(self, connection, login, password):
        '''
        Authenticate the client.
        '''
        try:
            role = self._server.conf.authentify(login, password)
        except CCConf.UnknownAccount:
            raise AuthenticationError('Authentication failure (Unknown login)')

        if role is None:
            logging.info('New authentication from %s: failure' % login)
            raise AuthenticationError('Authentication failure')
        else:
            # If authentication is a success, ask tags to the server:
            self._server.register(login, role, connection)
            connection.set_handler(WelcomeHandler.ROLES.get(role)(self._server))
            logging.info('New authentication from %s: success' % login)
            
            return role