CrowdTLS-server/crowdtls/main.py

import argparse
import asyncio
import sys
from pathlib import Path
from types import FrameType

import uvicorn
import uvloop
from dotenv import load_dotenv

from crowdtls.logs import logger

app_fastapi = None
app_rocketry = None


class CrowdTLS(uvicorn.Server):
    def handle_exit(self, sig: int, frame: FrameType) -> None:
        logger.info("Shutting down CrowdTLS")
        return super().handle_exit(sig, frame)


async def start_server() -> None:
    logger.info("Starting CrowdTLS")
    server = CrowdTLS(config=uvicorn.Config(app=app_fastapi, workers=1, loop="uvloop"))

    fastapi = asyncio.create_task(server.serve())
    rocket = asyncio.create_task(app_rocketry.serve())
    app_rocketry.task

    await asyncio.wait([rocket, fastapi], return_when=asyncio.FIRST_COMPLETED)


def run(env: Path) -> None:
    global app_rocketry
    global app_fastapi

    if env.exists():
        try:
            load_dotenv(env)
        except Exception as e:
            logger.error(f"Could not load env file {env}: {e}")
            sys.exit(1)

    from crowdtls.scheduler import app as app_rocketry
    from crowdtls.webserver import app as app_fastapi
    from crowdtls.analytics.main import ANOMALY_SCANNERS

    logger.info(f"Anomaly scanners loaded: {', '.join([x.__name__ for x in ANOMALY_SCANNERS.keys()])}")

    if sys.version_info >= (3, 11):
        with asyncio.Runner(loop_factory=uvloop.new_event_loop) as runner:
            runner.run(start_server())
    else:
        uvloop.install()
        asyncio.run(start_server())


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="CrowdTLS Server")
    parser.add_argument(
        "--env",
        type=lambda p: Path(p).absolute(),
        default=Path(__file__).absolute().parent / ".env",
        required=False,
        help="Path to specific env file",
    )
    print(Path(__file__).absolute().parent / ".env")
    args = parser.parse_args()
    run(args.env)