cc-server

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

from cloudcontrol.common.helpers.logger import patch_logging; patch_logging()

import os
import sys
import atexit
import logging
import logging.handlers
import ConfigParser
import signal
from optparse import OptionParser
from pwd import getpwnam
from grp import getgrnam

from daemon import DaemonContext

from cloudcontrol.server.server import CCServer
from cloudcontrol.server import __version__


DEFAULT_CONFIG_FILE = '/etc/cc-server.conf'
DEFAULT_UMASK = 0o0177
DEFAULT_CONFIGURATION = {
    'daemonize': False,
    'user': '',
    'group': '',
    'pidfile': '',
    'umask': '077',
    'port': 1984,
    'debug': False,
    'account_db': None, # None = mandatory option
    'interface': '127.0.0.1',
    'ssl_cert': None,
    'ssl_key': None,
    'maxcon': 600,
    'maxidle': 120,
}


class EncodingFormatter(logging.Formatter, object):

    def __init__(self, fmt=None, datefmt=None, encoding='utf-8'):
        super(EncodingFormatter, self).__init__(fmt, datefmt)
        self._encoding = encoding

    def formatException(self, ei):
        expt = super(EncodingFormatter, self).formatException( ei)
        if isinstance(expt, str):
            expt = expt.decode(self._encoding, 'replace')
        return expt

    def format(self, record):
        msg = super(EncodingFormatter, self).format(record)
        if isinstance(msg, unicode):
            msg = msg.encode(self._encoding, 'replace')
        return msg


def run_server(options):

    # Setup logging facility:
    level = logging.INFO
    if options['debug'] in (True, 'yes', 'true'):
        level = logging.DEBUG

    logger = logging.getLogger()
    logger.setLevel(level)

    if options['stdout']:
        handler = logging.StreamHandler()
        fmt = EncodingFormatter('[%(asctime)s] '
                                '\x1B[30;47m%(name)s\x1B[0m '
                                '\x1B[30;42m%(levelname)s\x1B[0m: '
                                '%(message)s')
    else:
        facility = logging.handlers.SysLogHandler.LOG_DAEMON
        handler = logging.handlers.SysLogHandler(address='/dev/log',
                                                 facility=facility)
        fmt = EncodingFormatter('%(name)s: %(levelname)s %(message)s')

    handler.setFormatter(fmt)
    logger.addHandler(handler)

    server = CCServer(logger.getChild('cc-server'),
                      conf_dir=options['account_db'],
                      maxcon=int(options['maxcon']),
                      maxidle=int(options['maxidle']),
                      port=int(options['port']),
                      address=options['interface'],
                      keyfile=options['ssl_key'],
                      certfile=options['ssl_cert'])

    def shutdown_handler(signum, frame):
        '''
        Handler called when SIGINT is emitted.
        '''

        server.rpc.shutdown()
        logging.info('Server properly exited by SIGINT')

    watcher_sigint = server.rpc.loop.signal(signal.SIGINT, shutdown_handler)
    watcher_sigstop = server.rpc.loop.signal(signal.SIGTERM, shutdown_handler)
    watcher_sigint.start()
    watcher_sigstop.start()

    try:
        server.run()
    except Exception as err:
        logging.critical('Server failed: %s', err)
        import traceback
        traceback.print_exc(file=sys.stdout)
        sys.exit(3)

if __name__ == '__main__':

    op = OptionParser(version='%%prog v%s' % __version__)
    op.add_option('-c', '--config', default=DEFAULT_CONFIG_FILE,
                  help='configuration file (default: %default)')

    op.add_option('-d', '--daemonize', action='store_true',
                  help='run as a daemon')
    op.add_option('-f', '--foreground', action='store_false', dest='daemonize',
                  help='don\'t run as a daemon')

    op.add_option('-p', '--pidfile', help='write pidfile to the path')
    op.add_option('-k', '--umask', help='set the umask of the process')
    op.add_option('-u', '--user', help='run as user')
    op.add_option('-g', '--group', help='run as group')

    op.add_option('-s', '--stdout', action='store_true', default=False,
                  help='log in stdout instead of syslog')

    cliopts, args = op.parse_args()

    # Reading the config file:
    config = ConfigParser.SafeConfigParser()
    config.read(cliopts.config)
    try:
        options = dict(config.items('server'))
    except ConfigParser.NoSectionError:
        sys.stderr.write("Configuration error: 'server' section must exist "
                         "in '%s'\n" % cliopts.config)
        sys.exit(1)

    # Applying default config file options:
    for opt, default in DEFAULT_CONFIGURATION.iteritems():
        if opt not in options or not options[opt]:
            if default is None:
                sys.stderr.write("Configuration error: you must specify '%s' "
                                 "option in '%s' !\n" % (opt, cliopts.config))
                sys.exit(1)
            else:
                options[opt] = default

    # Merge cli options and .conf file options:
    for opt in ('daemonize', 'pidfile', 'umask', 'user', 'group', 'stdout'):
        if getattr(cliopts, opt) is not None:
            options[opt] = getattr(cliopts, opt)

    # Create option set for the daemonization process:
    daemon_opts = {}

    daemonize = options['daemonize']
    if isinstance(daemonize, str):
        daemonize = daemonize in ('yes', 'true')
    daemon_opts['detach_process'] = daemonize

    daemon_opts['umask'] = int(options['umask'], 8)

    if options['user']:
        if options['user'].isdigit():
            daemon_opts['uid'] = int(options['user'])
        else:
            daemon_opts['uid'] = getpwnam(options['user']).pw_uid

    if options['group']:
        if options['group'].isdigit():
            daemon_opts['gid'] = int(options['group'])
        else:
            daemon_opts['gid'] = getgrnam(options['group']).gr_gid

    if not daemonize:
        daemon_opts['stderr'] = sys.stderr
        daemon_opts['stdout'] = sys.stderr

    # I've to write myself the pidfile because the daemon library write it
    # after the privilege downgrade:
    if options['pidfile']:
        pidfile = open(cliopts.pidfile, 'w')
        daemon_opts['files_preserve'] = [pidfile]
    else:
        pidfile = None

    with DaemonContext(**daemon_opts):
        if pidfile is not None:
            pidfile.write('%s' % os.getpid())
            pidfile.flush()

            @atexit.register
            def clean_pidfile():
                pidfile.seek(0)
                pidfile.truncate()
                pidfile.flush()

        run_server(options)