drinks-manager/scripts/_bootstrap.py

#!./venv/bin/python3
# Copyright 2023 Julian Müller (ChaoticByte)

import os

from argparse import ArgumentParser
from atexit import register as register_exithandler
from pathlib import Path
from signal import SIGINT
from subprocess import Popen
from sys import stdout, stderr
from time import sleep

from yaml import safe_load


banner = """ ___        _        _
|   \  _ _ (_) _ _  | |__ ___  ___
| |) || '_|| || ' \ | / /(_-< |___|
|___/ |_|  |_||_||_||_\_\/__/
 __  __                        Version {version}
|  \/  | __ _  _ _   __ _  __ _  ___  _ _
| |\/| |/ _` || ' \ / _` |/ _` |/ -_)| '_|
|_|  |_|\__,_||_||_|\__,_|\__, |\___||_|
                          |___/
"""

base_directory = Path(__file__).parent.parent
data_directory = base_directory / "data"
logfile_directory = data_directory / "logs"
configuration_file = data_directory / "config.yml"
caddyfile = data_directory / "Caddyfile"
logfile_caddy = logfile_directory / "caddy.log"
logfile_app = logfile_directory / "app.log"
logfile_sessioncleanup = logfile_directory / "session-cleanup.log"


class MonitoredSubprocess:
    def __init__(
        self,
        name: str,
        commandline: list,
        logfile: Path,
        environment: dict = os.environ,
        max_tries: int = 5,
    ):
        self.name = name
        self.commandline = commandline
        self.logfile = logfile
        self.environment = environment
        self.max_tries = max_tries
        self.s = None # the subprocess object
        self._tries = 0
        self._stopped = False

    def try_start(self):
        if self._tries < self.max_tries:
            self._tries += 1
            print(f"Starting {self.name}...")
            if self.logfile is None:
                self.s = Popen(
                    self.commandline,
                    stdout=stdout.buffer,
                    stderr=stderr.buffer,
                    env=self.environment)
            else:
                with self.logfile.open("ab") as l:
                    self.s = Popen(
                        self.commandline,
                        stdout=l,
                        stderr=l,
                        env=self.environment)
            return True
        else:
            print(f"Max. tries exceeded ({self.name})!")
            # the process must already be stopped at this
            # point, so we can set the variable accordingly
            self._stopped = True
            return False

    def stop(self):
        if not self._stopped:
            print(f"Stopping {self.name}...")
            self.s.terminate()
            self._stopped = True


def cleanup_procs(processes):
    for p in processes:
        p.stop()


def start_and_monitor(monitored_subprocesses: list):
    # display banner
    print(banner.format(version=os.environ["APP_VERSION"]))
    # start processes
    for p in monitored_subprocesses:
        p.try_start()
    register_exithandler(cleanup_procs, monitored_subprocesses)
    # monitor processes
    try:
        while True:
            sleep(1)
            for p in monitored_subprocesses:
                returncode = p.s.poll()
                if returncode is None:
                    continue
                else:
                    print(f"{p.name} stopped with exit code {returncode}.")
                    if p.try_start() is False:
                        # stop everything if the process
                        # has exceeded max. tries
                        exit()
    except KeyboardInterrupt:
        print("Received KeyboardInterrupt, exiting...")
        exit()


if __name__ == "__main__":
    argp = ArgumentParser()
    argp.add_argument("--devel", help="Start development server", action="store_true")
    args = argp.parse_args()
    # Load configuration
    with configuration_file.open("r") as f:
        config = safe_load(f)
    # Prepare
    os.chdir(str(base_directory))
    Popen(
        ["./venv/bin/python3", "./manage.py", "collectstatic", "--noinput"], env=os.environ).wait()
    Popen(
        ["./venv/bin/python3", "./manage.py", "migrate", "--noinput"], env=os.environ).wait()
    # Caddy configuration via env
    environment_caddy = os.environ
    environment_caddy["DATADIR"] = str(data_directory.absolute())
    environment_caddy["CADDY_HOSTS"] = ", ".join(config["caddy"]["hosts"])
    environment_caddy["HTTP_PORT"] = str(config["caddy"]["http_port"])
    environment_caddy["HTTPS_PORT"] = str(config["caddy"]["https_port"])
    environment_caddy["APPLICATION_PORT"] = str(config["app"]["application_port"])
    environment_caddy["ACCESS_LOG"] = config["logs"]["http_access"]
    # Start
    if args.devel:
        procs = [
            MonitoredSubprocess(
                "Caddy Webserver",
                ["caddy", "run", "--config", str(caddyfile)],
                None,
                environment=environment_caddy),
            MonitoredSubprocess(
                "Django Development Server",
                ["./venv/bin/python3", "./manage.py", "runserver", str(config["app"]["application_port"])],
                None),
            MonitoredSubprocess(
                "Session Autocleaner",
                ["./scripts/_session-autocleaner.py", str(config["app"]["session_clear_interval"])],
                None)
        ]
        start_and_monitor(procs)
    else:
        # Application configuration via env
        environment_app = os.environ
        environment_app["APP_PROD"] = "1"
        print("\nRunning in production mode.\n")
        # define processes
        procs = [
            MonitoredSubprocess(
                "Caddy Webserver",
                ["caddy", "run", "--config", str(caddyfile)],
                logfile_caddy,
                environment=environment_caddy),
            MonitoredSubprocess(
                "Drinks-Manager",
                [
                    "./venv/bin/python3",
                    "-m",
                    "uvicorn",
                    "--host",
                    "127.0.0.1",
                    "--port",
                    str(config["app"]["application_port"]),
                    "project.asgi:application",
                ],
                logfile_app,
                environment=environment_app),
            MonitoredSubprocess(
                "Session Autocleaner",
                ["./scripts/_session-autocleaner.py", str(config["app"]["session_clear_interval"])],
                logfile_sessioncleanup)
        ]
        start_and_monitor(procs)