From 049d8ea92057515915fae3e4cb6a2f4d24a67e47 Mon Sep 17 00:00:00 2001
From: Neil Hanlon <neil@rockylinux.org>
Date: Fri, 20 Dec 2024 13:32:11 -0500
Subject: [PATCH] update run.py to enable a common entrypoint for the container

---
 run.py | 94 ++++++++++++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 81 insertions(+), 13 deletions(-)

diff --git a/run.py b/run.py
index 4bfdfa0..58b4c60 100644
--- a/run.py
+++ b/run.py
@@ -2,20 +2,88 @@ 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
 
-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}"
 
-application = create_app()
-application.debug = os.environ.get("MM2_DEBUG", False)
-application.config['SESSION_TYPE'] = "cachelib"
-application.config['SESSION_CACHELIB'] = FileSystemCache(cache_dir='/tmp/sessions', threshold=500)
-Session(application)
-application.wsgi_app = ProxyFix(application.wsgi_app, x_proto=1, x_host=1)
+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))