handlers.py

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

import inspect
import logging
from sjrpc.utils import RpcHandler, pure
from sjrpc.core import RpcError
from conf import CCConf
from exceptions import (AlreadyRegistered, AuthenticationError, RightError,
                        ReservedTagError, BadObjectError, BadRoleError,
                        NotConnectedAccountError)

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


class Reporter(object):
    '''
    Simple class used to report error, warning and success of command execution.
    '''

    def __init__(self):
        self._reports = {}

    def get_dict(self):
        return self._reports.copy()

    def success(self, oid, message):
        self._reports[oid] = ('success', message)
        
    def warn(self, oid, message):
        self._reports[oid] = ('warn', message)
        
    def error(self, oid, message):
        self._reports[oid] = ('error', message)


class CCHandler(RpcHandler):
    '''
    Base class for handlers of CloudControl server.
    '''
    
    def __init__(self, server):
        self._server = server

    def __getitem__(self, name):
        if name.startswith('_'):
            # Filter the private members access:
            raise KeyError('Remote name %s is private.' % repr(name))
        else:
            logging.debug('Called %s.%s', self.__class__.__name__, name)
            return super(CCHandler, self).__getitem__(name)


    def list_commands(self, conn):
        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, conn):
        self._server.unregister(conn)

    def _check(self, conn, method, tql):
        client = self._server.search_client_by_connection(conn)
        allow = self._server.check(client, method, tql)
        if not allow:
            raise RightError('You are not allowed to do this action.')


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

    @listed
    def register(self, conn, obj_id, role):
        '''
        Register an object managed by the calling node.

        .. note:
           the obj_id argument passed to this handler is the object id of the
           registered object (not the fully qualified id, the server will
           preprend the id by "node_id." itself).

        :param obj_id: the id of the object to register
        :param role: the role of the object to register
        '''

        client = self._server.search_client_by_connection(conn)
        self._server.sub_register(client.login, obj_id, role)

    @listed
    def unregister(self, conn, obj_id):
        '''
        Unregister an object managed by the calling node.

        .. note:
           the obj_id argument passed to this handler is the object id of the
           unregistered object (not the fully qualified id, the server will
           preprend the id by "node_id." itself).

        :param obj_id: the id of the object to unregister
        '''

        client = self._server.search_client_by_connection(conn)
        self._server.sub_unregister(client.login, obj_id)


class CliHandler(OnlineCCHandler):
    '''
    Handler binded to 'cli' role.
    '''

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

        self._check(conn, 'list', query)
        logging.debug('Executed list function with query %s', query)
        return self._server.list(query)

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

    @listed
    def start(self, conn, query):
        self._check(conn, 'start', query)
        self._vm_action(query, 'vm_start')

    @listed
    def stop(self, conn, query):
        self._check(conn, 'stop', query)
        self._vm_action(query, 'vm_stop', force=False)

    @listed
    def destroy(self, conn, query):
        self._check(conn, 'destroy', query)
        self._vm_action(query, 'vm_stop', force=True)

    @listed
    def pause(self, conn, query):
        self._check(conn, 'pause', query)
        self._vm_action(query, 'vm_suspend')

    @listed
    def resume(self, conn, query):
        self._check(conn, 'resume', query)
        self._vm_action(query, 'vm_resume')

    @listed
    def passwd(self, conn, query, password, method='ssha'):
        '''
        Define a new password for specified user.
        '''

        self._check(conn, 'passwd', query)
        objects = self._server.list(query, show=set(('a',)))
        errs = Reporter()
        with self._server.conf:
            for obj in objects:
                if 'a' not in obj:
                    errs.error(obj['id'], 'not an account')
                    continue

                self._server.conf.set_password(obj['a'], password, method)

        return errs.get_dict()

    @listed
    def addaccount(self, conn, login, role, password=None):
        '''
        Create a new account with specified login.
        '''
        if role in WelcomeHandler.ROLES:
            self._server.conf.create_account(login, role, password)
        else:
            raise BadRoleError('%r is not a legal role.' % role)

    @listed
    def addtag(self, conn, query, tag_name, tag_value):
        '''
        Add a tag to the account which match the specified query.
        '''

        self._check(conn, 'addtag', query)

        if tag_name in self._server.RESERVED_TAGS:
            raise ReservedTagError('Tag %r is read-only' % tag_name)

        objects = self._server.list(query, show=set(('a',)))
        errs = Reporter()
        with self._server.conf:
            for obj in objects:
                if 'a' not in obj:
                    errs.error(obj['id'], 'not an account')
                    continue
                tags = self._server.conf.show(obj['a'])['tags']
                if tag_name in tags:
                    errs.warn(obj['id'], 'tag already exists, changed from %r'
                                         ' to %r' % (tags[tag_name], tag_value))
                else:
                    errs.success(obj['id'], 'tag created')
                self._server.conf.add_tag(obj['a'], tag_name, tag_value)

        return errs.get_dict()

    @listed
    def deltag(self, conn, query, tag_name):
        '''
        Remove a tag of the account with specified login.
        '''

        self._check(conn, 'deltag', query)

        if tag_name in self._server.RESERVED_TAGS:
            raise ReservedTagError('Tag %r is read-only' % tag_name)

        objects = self._server.list(query, show=set(('a',)))
        errs = Reporter()
        with self._server.conf:
            for obj in objects:
                if 'a' not in obj:
                    errs.error(obj['id'], 'not an account')
                    continue
                tags = self._server.conf.show(obj['a'])['tags']
                if tag_name in tags:
                    errs.success(obj['id'], 'tag deleted')
                else:
                    errs.warn(obj['id'], 'unknown tag')
                self._server.conf.remove_tag(obj['a'], tag_name)

        return errs.get_dict()

    @listed
    def tags(self, conn, query):
        '''
        Return all static tags attached to the specified login.
        '''

        self._check(conn, 'tags', query)
        objects = self._server.list(query, show=set(('a',)))
        for obj in objects:
            if 'a' not in obj:
                raise BadObjectError('All objects must have the "a" tag.')
        tags = []
        for obj in objects:
            otags = self._server.conf.show(obj['a'])['tags']
            otags.update({'id': obj['id']})
            tags.append(otags)
        return tags

    @listed
    def delaccount(self, conn, query):
        '''
        Delete the account with specified login.
        '''

        self._check(conn, 'delaccount', query)
        objects = self._server.list(query, show=set(('a',)))
        errs = Reporter()
        with self._server.conf:
            for obj in objects:
                if 'a' not in obj:
                    errs.error(obj['id'], 'not an account')
                    continue
                try:
                    self._server.conf.remove_account(obj['a'])
                except CCConf.UnknownAccount:
                    errs.error(obj['id'], 'unknown account')
                else:
                    errs.success(obj['id'], 'account deleted')
                
        return errs.get_dict()

    @listed
    def close(self, conn, query):
        '''
        Close an account without deleting it.
        '''

        self._check(conn, 'close', query)
        objects = self._server.list(query, show=set(('a',)))
        errs = Reporter()
        with self._server.conf:
            for obj in objects:
                if 'a' not in obj:
                    errs.error(obj['id'], 'not an account')
                    continue
                tags = self._server.conf.show(obj['a'])['tags']
                if 'close' in tags:
                    errs.warn(obj['id'], 'account already closed')
                else:
                    errs.success(obj['id'], 'closed')
                self._server.conf.add_tag(obj['a'], 'close', 'yes')

        return errs.get_dict()

    @listed
    def declose(self, conn, query):
        '''
        Re-open an closed account.
        '''

        self._check(conn, 'declose', query)
        objects = self._server.list(query, show=set(('a',)))
        errs = Reporter()
        with self._server.conf:
            for obj in objects:
                if 'a' not in obj:
                    errs.error(obj['id'], 'not an account')
                    continue
                tags = self._server.conf.show(obj['a'])['tags']
                if 'close' in tags:
                    errs.success(obj['id'], 'account declosed')
                else:
                    errs.warn(obj['id'], 'account not closed')
                self._server.conf.remove_tag(obj['a'], 'close')

        return errs.get_dict()

    @listed
    def kill(self, conn, query):
        '''
        Disconnect all connected accounts selected by query.
        '''

        self._check(conn, 'kill', query)
        objects = self._server.list(query, show=set(('a',)))
        errs = Reporter()
        with self._server.conf:
            for obj in objects:
                print obj
                if 'a' not in obj:
                    errs.error(obj['id'], 'not an account')
                    continue
                try:
                    self._server.kill(obj['a'])
                except NotConnectedAccountError:
                    errs.error(obj['id'], 'account is not connected')
                else:
                    errs.success(obj['id'], 'account killed')

        return errs.get_dict()

    @listed
    def rights(self, conn, query):
        '''
        Get the rights of an object set.
        '''

        self._check(conn, 'rights', query)
        objects = self._server.list(query, show=set(('a',)))
        rules = {}
        for obj in objects:
            if 'a' in obj:
                rules[obj['a']] = self._server.conf.show(obj['a'])['rights']
            else:
                raise BadObjectError('All objects must have the "a" tag.')
        
        return rules

    @listed
    def addright(self, conn, query, tql, method=None, allow=True, index=None):
        '''
        Add a right rule to the selected objects.
        '''

        self._check(conn, 'addright', query)
        objects = self._server.list(query, show=set(('a',)))
        errs = Reporter()
        with self._server.conf:
            for obj in objects:
                if 'a' not in obj:
                    errs.error(obj['id'], 'not an account')
                    continue
                try:
                    self._server.conf.add_right(obj['a'], tql, method,
                                                allow, index)
                except conf.UnknownAccount:
                    errs.error(obj['id'], 'unknown account')
                else:
                    errs.success(obj['id'], 'right rule added')
                
        return errs.get_dict()

    @listed
    def delright(self, conn, query, index):
        '''
        Remove a right rule from the selected objects.
        '''

        self._check(conn, 'delright', query)
        objects = self._server.list(query, show=set(('a',)))
        errs = Reporter()
        with self._server.conf:
            for obj in objects:
                if 'a' not in obj:
                    errs.error(obj['id'], 'not an account')
                    continue
                try:
                    self._server.conf.remove_right(obj['a'], index)
                except conf.UnknownAccount:
                    errs.error(obj['id'], 'unknown account')
                except conf.OutOfRangeIndexError:
                    errs.error(obj['id'], 'index out of range')
                else:
                    errs.success(obj['id'], 'right rule deleted')
                
        return errs.get_dict()

    @listed
    def execute(self, conn, query, command):
        '''
        Execute command on matched objects (must be roles hv or host).

        :param query: the tql query to select objects.
        :param command: the command to execute on each object
        :return: a dict where key is the id of a selected object, and the value
            a tuple (errcode, message) where errcode is (success|error|warn) and
            message an error message or the output of the command in case of
            success.
        '''

        self._check(conn, 'execute', query)
        objects = self._server.list(query, show=set(('r',)))
        errs = Reporter()
        for obj in objects:
            if obj['r'] not in ('hv', 'host'):
                errs.error(obj['id'], 'bad role')
                continue
            try:
                objcon = self._server.get_connection(obj['id'])
            except KeyError:
                errs.error(obj['id'], 'node not connected')
            else:
                returned = objcon.connection.call('execute_command', command)
                errs.success(obj['id'], returned)
                
        return errs.get_dict()

    @listed
    def shutdown(self, conn, query, reboot=True, gracefull=True):
        '''
        Execute a shutdown on selected objects (must be roles hv or host).

        :param query: the tql query to select objects.
        :param reboot: reboot the host instead of just shut it off
        :param gracefull: properly shutdown the host
        :return: a dict where key is the id of a selected object, and the value
            a tuple (errcode, message) where errcode is (success|error|warn) and
            message an error message.
        '''

        self._check(conn, 'execute', query)
        objects = self._server.list(query, show=set(('r',)))
        errs = Reporter()
        for obj in objects:
            if obj['r'] not in ('hv', 'host'):
                errs.error(obj['id'], 'bad role')
                continue
            try:
                objcon = self._server.get_connection(obj['id'])
            except KeyError:
                errs.error(obj['id'], 'node not connected')
            else:
                try:
                    objcon.connection.call('node_shutdown',
                                                      reboot, gracefull)
                except RpcError as err:
                    errs.error(obj['id'], '%s (exc: %s)' % (err.message,
                                                            err.exception))
                else:
                    errs.success(obj['id'], 'ok')
                
        return errs.get_dict()

    @listed
    def dbstats(self, conn):
        '''
        Return statistics about current database status.
        '''
        return self._server.objects.stats()

    def proxy_client(self, conn, 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 = {
        'cli': CliHandler,
        'hv': HypervisorHandler,
        'host': None,
    }

    @listed
    def authentify(self, conn, login, password):
        '''
        Authenticate the client.
        '''

        conf = self._server.conf
        with conf:
            try:
                role = self._server.conf.authentify(login, password)
            except CCConf.UnknownAccount:
                raise AuthenticationError('Unknown login')
            else:
                if 'close' in self._server.conf.show(login)['tags']:
                    raise AuthenticationError('Account is closed')

        if role not in WelcomeHandler.ROLES:
            raise BadRoleError('%r is not a legal role' % role)

        if role is None:
            logging.info('New authentication from %s: failure',
                         login.encode('ascii', 'ignore'))
            raise AuthenticationError('Bad login/password')
        else:
            # If authentication is a success, try to register the client:
            try:
                self._server.register(login, role, conn)
            except AlreadyRegistered:
                raise AuthenticationError('Already connected')
            conn.set_handler(WelcomeHandler.ROLES.get(role)(self._server))
            
            return role