diamond

#!/usr/bin/env python
# coding=utf-8

import os
import sys
import configobj
if os.name != 'nt':
    import pwd
    import grp

try:
    from setproctitle import setproctitle
except ImportError:
    setproctitle = None

for path in [
    os.path.join('opt', 'diamond', 'lib'),
    os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'src'))
]:
    if os.path.exists(os.path.join(path, 'diamond', '__init__.py')):
        sys.path.append(path)
        break

from diamond.server import Server
from diamond.util import get_diamond_version
from diamond.utils.log import setup_logging

import multiprocessing
import optparse
import signal


def main():
    try:
        # Initialize Options

        defaults = {
            'skip_pidfile': False,
        }

        if os.name == 'nt':
            defaults['skip_pidfile'] = True

        parser = optparse.OptionParser()

        parser.add_option("-c", "--configfile",
                          dest="configfile",
                          default="/etc/diamond/diamond.conf",
                          help="config file")

        parser.add_option("-f", "--foreground",
                          dest="foreground",
                          default=False,
                          action="store_true",
                          help="run in foreground")

        parser.add_option("-l", "--log-stdout",
                          dest="log_stdout",
                          default=False,
                          action="store_true",
                          help="log to stdout")

        parser.add_option("-p", "--pidfile",
                          dest="pidfile",
                          default=None,
                          help="pid file")

        parser.add_option("-r", "--run",
                          dest="collector",
                          default=None,
                          help="run a given collector once and exit")

        parser.add_option("-v", "--version",
                          dest="version",
                          default=False,
                          action="store_true",
                          help="display the version and exit")

        parser.add_option("--skip-pidfile",
                          dest="skip_pidfile",
                          default=defaults['skip_pidfile'],
                          action="store_true",
                          help="Skip creating PID file")

        parser.add_option("-u", "--user",
                          dest="user",
                          default=None,
                          help="Change to specified unprivilegd user")
        parser.add_option("-g", "--group",
                          dest="group",
                          default=None,
                          help="Change to specified unprivilegd group")
        parser.add_option("--skip-change-user",
                          dest="skip_change_user",
                          default=False,
                          action="store_true",
                          help="Skip changing to an unprivilegd user")

        parser.add_option("--skip-fork",
                          dest="skip_fork",
                          default=False,
                          action="store_true",
                          help="Skip forking (damonizing) process")

        # Parse Command Line Args
        (options, args) = parser.parse_args()

        # Initial variables
        uid = -1
        gid = -1

        if options.version:
            print "Diamond version %s" % (get_diamond_version())
            sys.exit(0)

        # Initialize Config
        options.configfile = os.path.abspath(options.configfile)
        if os.path.exists(options.configfile):
            config = configobj.ConfigObj(options.configfile)
        else:
            print >> sys.stderr, "ERROR: Config file: %s does not exist." % (
                options.configfile)
            parser.print_help(sys.stderr)
            sys.exit(1)

        # Initialize Logging
        log = setup_logging(options.configfile, options.log_stdout)

    # Pass the exit up stream rather then handle it as an general exception
    except SystemExit, e:
        raise SystemExit

    except Exception, e:
        import traceback
        sys.stderr.write("Unhandled exception: %s" % str(e))
        sys.stderr.write("traceback: %s" % traceback.format_exc())
        sys.exit(1)

    # Switch to using the logging system
    try:
        # PID MANAGEMENT
        if not options.skip_pidfile:
            # Initialize Pid file
            if not options.pidfile:
                options.pidfile = str(config['server']['pid_file'])

            # Read existing pid file
            try:
                pf = file(options.pidfile, 'r')
                pid = int(pf.read().strip())
                pf.close()
            except (IOError, ValueError):
                pid = None

            # Check existing pid file
            if pid:
                # Check if pid is real
                if not os.path.exists("/".join(["/proc", str(pid), "cmdline"])):
                    # Pid is not real
                    os.unlink(options.pidfile)
                    pid = None
                    print >> sys.stderr, (
                        "WARN: Bogus pid file was found. I deleted it.")
                else:
                    print >> sys.stderr, (
                        "ERROR: Pidfile exists. Server already running?")
                    sys.exit(1)

            # Get final GIDs
            if os.name != 'nt':
                if options.group is not None:
                    gid = grp.getgrnam(options.group).gr_gid
                elif len(config['server']['group']):
                    gid = grp.getgrnam(config['server']['group']).gr_gid

            # Get final UID
            if os.name != 'nt':
                if options.user is not None:
                    uid = pwd.getpwnam(options.user).pw_uid
                elif len(config['server']['user']):
                    uid = pwd.getpwnam(config['server']['user']).pw_uid

            # Fix up pid permissions
            if not options.foreground and not options.collector:
                # Write pid file
                pid = str(os.getpid())
                try:
                    pf = file(options.pidfile, 'w+')
                except IOError, e:
                    print >> sys.stderr, "Failed to write PID file: %s" % (e)
                    sys.exit(1)
                pf.write("%s\n" % pid)
                pf.close()
                os.chown(options.pidfile, uid, gid)
                # Log
                log.debug("Wrote First PID file: %s" % (options.pidfile))

        # USER MANAGEMENT
        if not options.skip_change_user:
            # Switch user to specified user/group if required
            try:
                if gid != -1 and uid != -1:
                    # Manually set the groups since they aren't set by default

                    # Python 2.7+
                    if hasattr(os, 'initgroups'):
                        os.initgroups(pwd.getpwuid(uid).pw_name, gid)
                    # Python 2.6
                    else:
                        os.setgroups(
                            [e.gr_gid for e in grp.getgrall() if pwd.getpwuid(
                                uid).pw_name in e.gr_mem] + [gid]
                            )

                if gid != -1 and os.getgid() != gid:
                    # Set GID
                    os.setgid(gid)

                if uid != -1 and os.getuid() != uid:
                    # Set UID
                    os.setuid(uid)

            except Exception, e:
                print >> sys.stderr, "ERROR: Failed to set UID/GID. %s" % (e)
                sys.exit(1)

            # Log
            log.info('Changed UID: %d (%s) GID: %d (%s).' % (
                os.getuid(),
                config['server']['user'],
                os.getgid(),
                config['server']['group']))

        # DAEMONIZE MANAGEMENT
        if not options.skip_fork:
            # Detatch Process
            if not options.foreground and not options.collector:

                # Double fork to serverize process
                log.info('Detaching Process.')

                # Fork 1
                try:
                    pid = os.fork()
                    if pid > 0:
                        # Exit first paren
                        sys.exit(0)
                except OSError, e:
                    print >> sys.stderr, "Failed to fork process." % (e)
                    sys.exit(1)
                # Decouple from parent environmen
                os.setsid()
                os.umask(0o022)
                # Fork 2
                try:
                    pid = os.fork()
                    if pid > 0:
                        # Exit second paren
                        sys.exit(0)
                except OSError, e:
                    print >> sys.stderr, "Failed to fork process." % (e)
                    sys.exit(1)
                # Close file descriptors so that we can detach
                sys.stdout.close()
                sys.stderr.close()
                sys.stdin.close()
                os.close(0)
                os.close(1)
                os.close(2)
                sys.stdout = open(os.devnull, 'w')
                sys.stderr = open(os.devnull, 'w')

        # PID MANAGEMENT
        if not options.skip_pidfile:
            # Finish Initialize PID file
            if not options.foreground and not options.collector:
                # Write pid file
                pid = str(os.getpid())
                try:
                    pf = file(options.pidfile, 'w+')
                except IOError, e:
                    log.error("Failed to write child PID file: %s" % (e))
                    sys.exit(1)
                pf.write("%s\n" % pid)
                pf.close()
                # Log
                log.debug("Wrote child PID file: %s" % (options.pidfile))

        # Initialize Server
        server = Server(configfile=options.configfile)

        def shutdown_handler(signum, frame):
            log.info("Signal Received: %d" % (signum))
            # Delete Pidfile
            if not options.skip_pidfile and os.path.exists(options.pidfile):
                os.remove(options.pidfile)
                # Log
                log.debug("Removed PID file: %s" % (options.pidfile))
            for child in multiprocessing.active_children():
                if 'SyncManager' not in child.name:
                    # The SyncManager process will immediately shutdown once
                    # the parent (us) exits. If we explicitly terminate the
                    # SyncManager here, we can't guarantee that it will exit
                    # after all collector processes and the handler process have
                    # exited first. As a result, since the collector and handler
                    # processes push/retrieve items to/from the shared queue via
                    # the SyncManager, it's possible for those processes to
                    # terminate unexpectedly (crash). When this happens, when we
                    # call exit(), we just hang and never actually terminate
                    # gracefully. Therefore, do not explicitly terminate the
                    # SyncManager process here but continue to terminate all
                    # other processes (collectors and the handler) and perform a
                    # join() on each of them. This guarantees that the
                    # SyncManager is terminated last (implicitly as a result of
                    # us exiting).
                    child_debug = "Terminating and joining on: {0} ({1})"
                    log.debug(child_debug.format(child.name, child.pid))
                    child.terminate()
                    child.join()
            sys.exit(0)

        # Set the signal handlers
        signal.signal(signal.SIGINT, shutdown_handler)
        signal.signal(signal.SIGTERM, shutdown_handler)

        server.run()

    # Pass the exit up stream rather then handle it as an general exception
    except SystemExit, e:
        raise SystemExit

    except Exception, e:
        import traceback
        log.error("Unhandled exception: %s" % str(e))
        log.error("traceback: %s" % traceback.format_exc())
        sys.exit(1)


if __name__ == "__main__":
    if setproctitle:
        setproctitle(os.path.basename(__file__))
    main()