__init__.py

import logging
import os.path
from itertools import chain, imap
from subprocess import Popen, PIPE, STDOUT

from ccnode.tags import Tag, tag_inspector
from ccnode.plugins import Base as BasePlugin
from ccnode.host import tags


logger = logging.getLogger(__name__)


def disk_tag_value(disk_name):
    def size():
        s = open(os.path.join('/sys/block', disk_name, 'size')).read().strip()
        try:
            s = int(s)
            if s > 0:
                return s * 512
            else:
                return None
        except ValueError:
            return None

    return size


class Handler(BasePlugin):
    """Handler for host role."""
    def __init__(self, *args, **kwargs):
        BasePlugin.__init__(self, *args, **kwargs)

        # add plugin tags
        self.tag_db['__main__'].update(dict(
            (t.name, t) for t in tag_inspector(tags),
        ))

        # disk related tags
        self.tag_db['__main__'].update(dict((t.name, t) for t in imap(
            lambda d: Tag('disk%s_size' % d, disk_tag_value(d), 60),
            self.tag_db['__main__']['disk']._calculate_value().split(),
        )))

        # rpc handler
        self.rpc_handler.update(dict(
            execute_command=self.execute_command,
            node_shutdown=self.node_shutdown,
        ))

    def execute_command(self, command):
        """Execute an arbitrary shell command on the host.

        :param string command: shell command to run
        """
        # return stdout and stderr mixed in
        return Popen(command, shell=True, bufsize=-1, stdin=PIPE, stdout=PIPE,
                     stderr=STDOUT).communicate()[0] or None

    def node_shutdown(self, reboot=True, gracefull=True):
        """Halt/Reboot the node.

        :param bool reboot: halt/reboot the system
        :param bool gracefull: force the operation (gracefull == not force)
        """
        args = ['/sbin/shutdown']
        if reboot:
            args.append('-r')
            if gracefull:
                logger.info('Going to reboot the host...')
                args.append('-f')
            else:
                logger.info('Going to force the reboot of the host...')
                args.append('-n')
        else:
            # halt
            args.append('-h -P')
            if not gracefull:
                logger.info('Going to halt the host...')
                args.append('-n')
            else:
                logger.info('Going to force the halt of the host...')

        args.append('0')

        return self.execute_command(' '.join(args))