handlers.py

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

import inspect
import logging
from sjrpc.utils import RpcHandler, pure
from tql import TqlQuery, TqlCondition, TqlLimit
from conf import CCConf
from exceptions import AlreadyRegistered, AuthenticationError

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():
            if isinstance(condition, TqlCondition):
                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['con'] != 'offline']
                    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))]

            elif isinstance(condition, TqlLimit):
                objects = objects[:condition.number]
                
        # 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 + '$vm')
        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 destroy(self, query):
        self.stop(query, force=True)

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

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

    @pure
    @listed
    def account_setpassword(self, login, password, method='ssha'):
        '''
        Define a new password for specified user.
        '''
        self._server.conf.set_password(login, password, method)

    @pure
    @listed
    def account_create(self, login, password, role):
        '''
        Create a new account with specified login.
        '''
        self._server.conf.create_account(login, password, role)

    @pure
    @listed
    def account_addtag(self, login, tag):
        '''
        Add a tag to the account with specified login.
        '''
        self._server.conf.add_tag(login, tag)

    @pure
    @listed
    def account_removetag(self, login, tag):
        '''
        Remove a tag of the account with specified login.
        '''
        self._server.conf.remove_tag(login, tag)

    @pure
    @listed
    def account_delete(self, login):
        '''
        Delete the account with specified login.
        '''
        self._server.conf.remove_account(login)

    @pure
    def proxy_client(self, login, command, *args, **kwargs):
        client = self._server.get_connection(login)
        return client.connection.call(command, *args, **kwargs)

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, try to register the client:
            try:
                self._server.register(login, role, connection)
            except AlreadyRegistered:
                raise AuthenticationError('Authentication failure '
                                          '(already connected)')
            connection.set_handler(WelcomeHandler.ROLES.get(role)(self._server))
            logging.info('New authentication from %s: success' % login)
            
            return role