bureaucrat

#!/usr/bin/env python
""" Bureaucrat - The Procfile & Deployfile process manager for Python Virtual Environments """

import argparse
import logging
import os
import re
import subprocess
import signal
import sys
import time


# Config
EXIT_FATAL = 128
EXIT_OK = 0
PROCESS_POLL_INTERVAL = 5


# create logger
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)

# create console handler and set level to debug
ch = logging.StreamHandler(sys.stdout)
ch.setLevel(logging.DEBUG)
# create formatter
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
ch.setFormatter(formatter)  # add formatter to ch
logger.addHandler(ch)  # add ch to logger


class ProcessLine(object):
    """ Represents a single process instance, eg a line of a Procfile """

    def __init__(self, name, cmd):
        self.name = name
        self.cmd = cmd


class Process(object):
    """ Represents a single process manager """

    def __init__(self, name, cmd):
        self.pl = ProcessLine(name, cmd)
        self.sub_process = None
        self.pid_file = None
        self.log_file = None
        self.ended = False

    @property
    def name(self):
        return self.pl.name

    @property
    def cmd(self):
        return self.pl.cmd

    @property
    def pid(self):
        if self.pid_file:
            try:
                # open pid file
                with open(self.pid_file, 'r') as f:
                    pid = f.read()
            except IOError:
                return None
            else:
                return int(pid)
        else:
            return self.sub_process.pid

    def kill(self):
        # kill process
        if hasattr(self.sub_process, 'returncode') and self.sub_process.returncode is not None:
            self.ended = True
            return
        else:
            logger.info("Stopping %s: %s" % (self.name, self.cmd))
            try:
                os.kill(self.pid, signal.SIGTERM)
            except OSError:
                logger.warning("Proccess %s (pid: %s) not found." % (self.name, self.pid))
            finally:
                if self.pid_file:
                    self.rm_pid()
                self.ended = True

    def rm_pid(self):
        # remove pid
        try:
            os.remove(self.pid_file)
        except OSError:
            logger.warning("Couldn't remove %s" % self.pid_file)
            sys.exit(EXIT_FATAL)

    def expanded_cmd(self):
        # Replace all environment vars in the cmd path and split
        return os.path.expandvars(self.pl.cmd).split()

    def execute(self, cwd, background=True):
        if self.pid_file and os.path.exists(self.pid_file):
            logger.error("%s: %s exists" % (self.name, self.pid_file))
            return False
        else:
            logger.info("%s: %s" % (self.name, self.cmd))
        cmd = self.expanded_cmd()

        # On windows if shell is not set to true, virtualenv python is not used.
        # known bug for subprocess in windows https://bugs.python.org/issue18069
        # Also if the subprocess forks and returns, then we won't be able to track
        # the pid for successful start/stop actions. In this case you'll need to use a
        # wrapper script that invokes a shell that waits for the process to exit.
        if sys.platform == 'win32':
            shell = True
        else:
            shell = False

        if not background:
            log_to = subprocess.PIPE
        elif self.log_file is None:
            log_to = None  # This will go to stdout and still allow backgrounding
        else:
            log_to = open(self.log_file, 'w')

        try:
            self.sub_process = subprocess.Popen(cmd, cwd=cwd, stdout=log_to, stderr=subprocess.STDOUT, shell=shell)
            if not background:
                # blocks until the next line-feed
                while True:
                    line = self.sub_process.stdout.readline()
                    if not line:
                        break
                    logger.info(line.rstrip())

        except OSError as e:
            logger.error("%s. Command %s not found?" % (e, cmd[:1]))
            return False
        else:
            if self.pid_file:
                # write pid
                with open(self.pid_file, 'w') as pid:
                    pid.write(str(self.sub_process.pid))
            return self.sub_process


class ProcessManager(object):
    """ Manage all processes in a given process file  """

    def __init__(self, process_file, path_additions, env_file, log_path, pid_path, create_pids, named_procs=None,
                 debug=False):
        self.debug = debug
        self.create_pids = create_pids
        self.log_file = None

        # Init
        self._set_environment(env_file)
        self.processes = self._parse_process_file(process_file, log_path, pid_path)
        if named_procs is not None and len(named_procs) > 0:  # truncate processes
            self.processes = [p for p in self.processes if p.name in named_procs]
        # Set Path
        for p in path_additions:
            os.environ['PATH'] = p + os.pathsep + os.environ['PATH']

    def _set_environment(self, env_file):
        """ Set environment variables from .env file """
        reg = re.compile('(?P<name>\w+)(\=(?P<value>.+))')
        for line in open(env_file):
            m = reg.match(line)
            if m:
                name = m.group('name')
                value = ''
                if m.group('value'):
                    value = m.group('value')
                os.environ[name] = value
                if self.debug:
                    print("set %s %s" % (name, value))

    def _parse_process_file(self, process_file, log_path, pid_path):
        """ Parse Processfile (eg Procfile, Deployfile)
        :return: list of process objects
        """
        processes = []
        with open(process_file) as f:
            for line in f.readlines():
                match = re.search(r'([a-zA-Z0-9_-]+):(.*)', line)  # TODO: Exclude lines starting with '#'
                if not match:
                    raise Exception('Bad Process file line')
                name = match.group(1).strip()
                cmd = match.group(2).strip()
                p = Process(name, cmd)
                if pid_path is not None and self.create_pids:
                    p.pid_file = os.path.join(pid_path, '%s.pid' % name)  # Assign pid file
                else:
                    p.pid_file = None
                if log_path != '-':  # not stdout
                    p.log_file = os.path.join(log_path, '%s.log' % name)  # assign log file
                processes.append(p)
        return processes


class Bureaucrat(object):
    """ Main Class """

    def __init__(self, env_file, virtual_env, app_path, log_path, pid_path, debug=False):
        self.debug = debug
        self.virtual_env = virtual_env
        self.app_path = app_path
        self.env_file = env_file
        self.log_path = log_path
        self.pid_path = pid_path
        # Listen for shutdown
        signal.signal(signal.SIGTERM, self._sigterm_handler)

    def load_processfile(self, process_file, create_pids=True, named_procs=None):
        virtual_env_bin_path = '%s/bin' % self.virtual_env
        if sys.platform == 'win32':
            virtual_env_bin_path = '%s/Scripts' % self.virtual_env
        path_additions = [virtual_env_bin_path, self.app_path]
        self.process_manager = ProcessManager(process_file,
                                              path_additions,
                                              self.env_file,
                                              self.log_path,
                                              self.pid_path,
                                              create_pids,
                                              named_procs)

    def start(self):
        """  Start processes defined in Procfile """
        logger.info('Spawning Process Task(s)')
        for p in self.process_manager.processes:
            p.execute(cwd=self.app_path, background=True)

    def stop(self):
        """  Terminate processes defined in Procfile """
        for p in self.process_manager.processes:
            p.kill()

    def deploy(self):
        """  Run all tasks in Deployfile """
        logger.info('Running Deployment Task(s)')
        for p in self.process_manager.processes:
            p.execute(cwd=self.app_path, background=False)

    def _sigterm_handler(self, signal, frame):
        """ SIGTERM, eg kill handler """
        print('Got SIGTERM. Shutting down.')
        self.stop()
        sys.exit(EXIT_OK)

    @staticmethod
    def _check_running(processes):
        running = False
        for p in processes:
            if p.sub_process is False:
                logger.error('Failed to spawn process: %s' % p.name)
                sys.exit(EXIT_FATAL)
            elif p.sub_process is not None and p.ended is False:
                p.sub_process.poll()  # update returncode
                if p.sub_process.returncode is None:
                    running = True
                elif p.sub_process.returncode is 0:
                    logger.info("Spawned process ended: %s (pid: %s exit: %s)" %
                                (p.name, p.sub_process.pid, p.sub_process.returncode))
                    p.ended = True  # Stop monitoring
                    if p.pid_file:
                        p.rm_pid()
                elif p.sub_process.returncode is not 0:
                    logger.error("Spawned process exited with error: %s (pid: %s exit: %s)" %
                                 (p.name, p.sub_process.pid, p.sub_process.returncode))
                    sys.exit(EXIT_FATAL)  # bail out
        return running

    def monitor(self):
        """ Monitor processes defined in Procfile """
        try:
            while True:
                # Monitor all processes
                if self._check_running(self.process_manager.processes) is False:
                    logger.info('All spawned processes have ended.')
                    return False  # Exit Ok
                time.sleep(PROCESS_POLL_INTERVAL)
        except KeyboardInterrupt:
            logger.info('Shutting down.')
            self.stop()
            sys.exit(EXIT_OK)


def bureaucrat_init(args, process_file='Procfile', create_pids=True):
    # init
    virtual_env = args.venv or os.environ.get('VIRTUAL_ENV', os.getcwd())  # read env var from venv or param or cwd
    app_path = args.app or virtual_env
    env_file = args.envfile or os.path.join(app_path, '.env')

    if process_file == 'Procfile':
        process_file = args.procfile or os.path.join(app_path, 'Procfile')
        # Only truncate ProcFile processes, all DeployFile steps always run
        if hasattr(args, 'process'):
            named_procs = set(args.process) or set(os.environ.get('PROCFILE_TASKS', '').split()) or None
        else:
            named_procs = set(os.environ.get('PROCFILE_TASKS', '').split()) or None
    elif process_file == 'Deployfile':
        process_file = args.deployfile or os.path.join(app_path, 'Deployfile')
        named_procs = None

    if hasattr(args, 'pidpath'):
        pid_path = args.pidpath or ''
    else:
        pid_path = ''

    if hasattr(args, 'no_pid'):
        if args.no_pid:
            create_pids = False

    log_path = args.logpath or ''
    b = Bureaucrat(env_file, virtual_env, app_path, log_path, pid_path)
    b.load_processfile(process_file, create_pids, named_procs)
    return b


def stop(args):
    b = bureaucrat_init(args)
    return b.stop()


def start(args):
    b = bureaucrat_init(args)
    return b.start()


def restart(args):
    b = bureaucrat_init(args)
    b.stop()
    b.start()


def deploy(args):
    b = bureaucrat_init(args, process_file='Deployfile', create_pids=False)
    b.deploy()


def init(args):
    deploy(args)
    b = bureaucrat_init(args, create_pids=False)
    b.start()
    b.monitor()


if __name__ == "__main__":

    # Parse command line arguments
    parser = argparse.ArgumentParser(prog='Bureaucrat',
                         description='Bureaucrat - the Procfile & Deployfile manager for Python Virtual Environments')
    # Venv root
    venv = argparse.ArgumentParser(add_help=False)
    venv.add_argument('--venv', type=str, help='Virtualenv root')
    # Project / app root
    app = argparse.ArgumentParser(add_help=False)
    app.add_argument('--app', type=str, help='Application root')
    # Procfile
    procfile = argparse.ArgumentParser(add_help=False)
    procfile.add_argument('--procfile', type=str, help='Procfile path')
    # Deployfile
    deployfile = argparse.ArgumentParser(add_help=False)
    deployfile.add_argument('--deployfile', type=str, help='Deployfile path')
    # .env
    envfile = argparse.ArgumentParser(add_help=False)
    envfile.add_argument('--envfile', type=str, help='.env file path')
    # logpath
    logpath = argparse.ArgumentParser(add_help=False)
    logpath.add_argument('--logpath', type=str, help='log file path')
    # pidpath
    pidpath = argparse.ArgumentParser(add_help=False)
    pidpath.add_argument('--pidpath', type=str, help='pid file path')
    # don't create pid
    no_pid = argparse.ArgumentParser(add_help=False)
    no_pid.add_argument('--no-create-pid', action='store_true', help='Don\'t create PID files')
    # process name
    process_name = argparse.ArgumentParser(add_help=False)
    process_name.add_argument('process', nargs='*', type=str, help='Procfile Process Name')
    # Start
    sp = parser.add_subparsers()
    sp_start = sp.add_parser('start', parents=[venv, app, procfile, envfile, logpath, pidpath, process_name],
                             help='Starts Procfile processes')
    sp_start.set_defaults(func=start)
    # Stop
    sp_stop = sp.add_parser('stop', parents=[venv, app, procfile, envfile, logpath, pidpath, process_name],
                            help='Stops Procfile processes')
    sp_stop.set_defaults(func=stop)
    # Restart
    sp_restart = sp.add_parser('restart', parents=[venv, app, procfile, envfile, logpath, pidpath, process_name],
                               help='Restarts Procfile processes')
    sp_restart.set_defaults(func=restart)
    # Deploy
    sp_deploy = sp.add_parser('deploy', parents=[venv, app, deployfile, envfile, logpath],
                              help='Run tasks in Deployfile')
    sp_deploy.set_defaults(func=deploy)
    # Init
    sp_init = sp.add_parser('init', parents=[venv, app, deployfile, procfile, envfile, logpath, pidpath, no_pid,
                                             process_name],
                               help='Run Deployfile tasks and then start Procfile processes in foreground')
    sp_init.set_defaults(func=init)

    # Init parser if args passed
    if len(sys.argv) > 1:
        args = parser.parse_args()
        args.func(args)