mirrormanager-rocky/run.py

from werkzeug.middleware.proxy_fix import ProxyFix
from mirrormanager2.app import create_app
from flask_session import Session
from cachelib.file import FileSystemCache
from flask import Flask
import os
import sys
import subprocess
from threading import Thread


def setup_env():
    password = os.environ.get('MM2_DATABASE_PASSWORD')
    if password:
        user = os.environ.get('DB_USER')
        host = os.environ.get('DB_HOST')
        port = os.environ.get('DB_PORT')
        name = os.environ.get('DB_NAME')
        os.environ["MM2_SQLALCHEMY_DATABASE_URI"] = f"postgresql://{user}:{password}@{host}:{port}/{name}"


def mirrormanager_wsgi() -> Flask:
    application = create_app()
    application.debug = os.environ.get("MM2_DEBUG", False)
    application.config['SESSION_TYPE'] = "cachelib"
    application.config['SESSION_CACHELIB'] = FileSystemCache(cache_dir='/mnt/efs/fs/0/tmp/sessions', threshold=500)
    Session(application)
    application.wsgi_app = ProxyFix(application.wsgi_app, x_proto=1, x_host=1)


def execute(func: str) -> int:
    def stream_output(pipe, output_stream):
        while True:
            data = pipe.read(1024)  # Read in chunks of 1024 bytes
            if not data:
                break
            output_stream.write(data)
            output_stream.flush()
        pipe.close()

    argv = func.split()
    cmd = argv.pop(0)
    args = argv if argv else []

    print(f'running: {cmd} with args: {args}')

    process = subprocess.Popen([cmd] + args,
                               stdout=subprocess.PIPE,
                               stderr=subprocess.PIPE,
                               bufsize=0)  # Unbuffered binary mode

    stdout_thread = Thread(target=stream_output, args=(process.stdout, sys.stdout.buffer))
    stderr_thread = Thread(target=stream_output, args=(process.stderr, sys.stderr.buffer))

    stdout_thread.start()
    stderr_thread.start()

    process.wait()

    stdout_thread.join()
    stderr_thread.join()

    return process.returncode


if __name__ == "__main__":
    setup_env()
    # map cmd name to script/function
    cmd_map = {
        "check-propagation": "/usr/bin/echo undefined",
        "crawl-mirrors": "/usr/bin/echo undefined",
        "generate-worldmaps": "/opt/app-root/bin/mm2_generate-worldmap",
        "scan-primary-mirror": "/usr/local/bin/scan-primary-mirror",
        "update-clouds": "/usr/bin/echo undefined",
        "update-geoip": "/usr/bin/echo undefined",
        "generate-mirrorlist-cache": "/usr/local/bin/generate-mirrorlist-cache",
        "update-netblocks": "/opt/app-root/bin/mm2_get-netblocks",
        "update-EC2-netblocks": "/opt/app-root/bin/update-EC2-netblocks",
        "mirrorlist-server": "/usr/local/bin/mirrorlist-server",
        "mirrormanager-wsgi": mirrormanager_wsgi
    }

    arg = sys.argv[1] if len(sys.argv) > 1 else "mirrormanager-wsgi"
    if arg and (cmd := cmd_map.get(arg, None)):
        if not isinstance(cmd, str):
            application = cmd()
            exit(0)

        sys.exit(execute(cmd))