#!/usr/bin/python

# daemonize
#
# Daemonizes and runs a user-specified command, optionally restarting it if
# it terminates.
#
# Usage examples (fairly trivial):
#     daemonize -o ~/out -n 5 /bin/ls
#         (runs "ls" six times, writing the output to ~/out, then exits)
#
#     daemonize -n 5 -d /bin/ls
#         (same as above, but outputs to the console)
#
#     daemonize -n 1 -t 2 -d /bin/ls
#         (same as above, but only allows one restart every 2 seconds)
#
# Possible enhancements:
#     Provide options to chroot() and setuid() before the child is spawned.
#     Don't sleep() to satisfy limits when a SIGHUP is received, i.e. restart
#         immediately.
#
# Adapted from the ASPN Python Cookbook at:
#     http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/278731
#
# v0.03-1 - Only unlink() PID files if they were actually created in the first
#           place
#           by simon@grimsworld.org.uk on 23/06/2005
#
# v0.03   - Options to create PID files added
#           by simon@grimsworld.org.uk on 23/06/2005
#
# v0.02   - Logfiles are now reopened when a SIGHUP is received (in fact,
#           whenever the command is restarted). Also fixed HangupException bug.
#           by simon@grimsworld.org.uk on 16/06/2005
#
# v0.01   - Initial release
#           by simon@grimsworld.org.uk on 15/06/2005
#
# This code may be used under the BSD license:
#   http://www.grimsworld.org.uk/licenses/bsd.txt

import sys
import getopt
import os
import signal
import time

# Defaults
stdoutFilename = "/dev/null"
stderrFilename = "/dev/null"
maxRestarts = 0
minTime = 0
daemonize = True
parentPIDFilename = None
childPIDFilename = None

def minTimeToStr(t):
    if t == 0:
        return "infinite"
    else:
        return str(t)

def usage():
    """Displays a usage message."""

    print """Usage:
  """+sys.argv[0]+""" [option] ... command [arg] ...
Description:
  Daemonizes and runs a user-specified command, optionally restarting it
  if it terminates. Note that the full pathname of the command must be
  specified, and the command will be started in the root directory with
  umask 000. Upon receiving a SIGHUP, the daemonize process will kill the
  currently executing child process (restarting it if appropriate). A
  maximum number of restarts in a period of time can be specified; if these
  limits are exceeded then daemonize will sleep() until the specified period
  has elapsed. If the period of time is unspecified, daemonize will terminate
  after the maximum number of restarts has been reached.
Options:
  -o file        File to use for standard output (default: """+stdoutFilename+""")
  -e file        File to use for standard error (default: """+stderrFilename+""")
  -n num         Specify maximum number of restarts (default: """+str(maxRestarts)+""")
  -t num         Specify minimum number of seconds those restarts may occur
                   in (default: """+minTimeToStr(minTime)+""")
  -d             Don't daemonize or redirect streams (useful for testing)
  -p file        PID file for parent (daemonize)
  -c file        PID file for child
Example:
  """+sys.argv[0]+""" -o ~/out -n 5 /bin/ls"""

[opts, command] = getopt.getopt(sys.argv[1:], "o:e:n:t:dhp:c:")

for [opt, arg] in opts:
    if opt == "-o" and arg != "":
        stdoutFilename = arg

    elif opt == "-e" and arg != "":
        stderrFilename = arg

    elif opt == "-n" and arg != "":
        maxRestarts = int(arg)

    elif opt == "-t" and arg != "":
        minTime = int(arg)

    elif opt == "-d":
        daemonize = False

    elif opt == "-h":
        usage()
        sys.exit(0)

    elif opt == "-p":
        parentPIDFilename = arg

    elif opt == "-c":
        childPIDFilename = arg

if len(command) < 1:
    usage()
    sys.exit(1)

if daemonize:
    # Fork our first child
    pid = os.fork()

    if pid != 0:
        os._exit(0)

    # Create a new session and process group, making us the leader
    os.setsid()

    # Ignore the SIGHUP when the first child terminates
    signal.signal(signal.SIGHUP, signal.SIG_IGN)

    # Fork our second child so that we are no longer the session leader and
    # can never acquire a terminal
    pid = os.fork()

    if pid != 0:
        os._exit(0)

os.chdir("/")

if parentPIDFilename is not None:
    parentPIDFile = file(parentPIDFilename, "w")
    parentPIDFile.write(str(os.getpid()))
    parentPIDFile.close()

os.umask(0)

class HangupException(Exception):
    """The exception we raise when a SIGHUP is received."""

    def __init__(self, value = "SIGHUP received"):
        self.value = value

    def __str__(self):
        return repr(self.value)

class TerminateException(Exception):
    """The exception we raise when a SIGTERM is received."""

    def __init__(self, value = "SIGTERM received"):
        self.value = value

    def __str__(self):
        return repr(self.value)

def hangupHandler(signum, frame):
    """The function that is called when a SIGHUP is received."""
    raise HangupException()

def terminateHandler(signum, frame):
    """The function that is called when a SIGTERM is received."""
    raise TerminateException()

signal.signal(signal.SIGHUP, hangupHandler)
signal.signal(signal.SIGTERM, terminateHandler)

starts = [] # timestamps from when each child was spawned
running = True

while running:
    if daemonize:
        # Close all open files
        try:
            maxfd = os.sysconf("SC_OPEN_MAX")
        except (AttributeError, ValueError):
            maxfd = 256

        for fd in range(0, maxfd):
            try:
                os.close(fd)
            except OSError:
                pass

        # Reopen std(in|out|err)
        os.open("/dev/null", os.O_RDONLY) # stdin
        os.open(stdoutFilename, os.O_WRONLY | os.O_APPEND | os.O_CREAT, 0644)
        os.open(stderrFilename, os.O_WRONLY | os.O_APPEND | os.O_CREAT, 0644)

    # Store time so we can limit number of restarts in a given period
    starts.append(time.time())

    # And away we go...
    os.write(2, "Spawning child process\n")
    childPid = os.spawnv(os.P_NOWAIT, command[0], command)

    if childPIDFilename is not None:
        childPIDFile = file(childPIDFilename, "w")
        childPIDFile.write(str(childPid))
        childPIDFile.close()

    # Wait for the child to finish
    try:
        os.waitpid(childPid, 0)

    except HangupException:
        os.write(2, "SIGHUP received - killing child\n")
        os.kill(childPid, signal.SIGTERM)
        os.waitpid(childPid, 0)

    except TerminateException:
        os.write(2, "SIGTERM received - killing child\n")
        os.kill(childPid, signal.SIGTERM)
        os.waitpid(childPid, 0)
        running = False

    if childPIDFilename is not None:
        os.unlink(childPIDFilename)

    # Enforce limits if necessary
    if len(starts) > maxRestarts:
        if minTime == 0:
            os.write(2, "Maximum number of restarts reached\n")
            running = False

        else:
            firstStart = starts.pop(0)
            nextStart = time.time()

            if (nextStart - firstStart) < minTime:
                delay = minTime - (nextStart - firstStart)
                os.write(2, "Restarted too often; sleeping for " \
                    +str(delay)+" seconds\n")

                try:
                    time.sleep(delay)

                except HangupException:
                    os.write(2, "SIGHUP received whilst sleeping\n")

                except TerminateException:
                    os.write(2, "SIGTERM received whilst sleeping\n")
                    running = False

os.write(2, "Exiting\n")

if parentPIDFilename is not None:
    os.unlink(parentPIDFilename)