main.py

"""LLMProxy — LLM Security Gateway. Entry point."""

import asyncio
import logging
import os
import yaml
from dotenv import load_dotenv

from proxy.rotator import ProxyOrchestrator
from core.metrics import start_metrics_server
from core.discovery_utils import get_tailscale_ip
from core.tracing import TraceManager

load_dotenv()

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("llmproxy")


async def main():
    # Supply chain integrity check (pre-startup)
    from scripts.verify_deps import verify_all
    if not verify_all(strict=False):
        logger.critical("Supply chain integrity check FAILED. Aborting startup.")
        return

    config_file = os.environ.get("CONFIG_FILE", "config.yaml")
    with open(config_file, 'r') as f:
        config = yaml.safe_load(f) or {}

    # Apply env-based endpoint + WAF overlays so startup validation sees the
    # same merged view the orchestrator will use (see ProxyOrchestrator._load_config).
    from core.env_endpoints import inject_env_endpoints
    inject_env_endpoints(config)
    _firewall_env = os.environ.get("LLM_PROXY_FIREWALL_ENABLED")
    if _firewall_env is not None:
        _enabled = _firewall_env.strip().lower() not in ("0", "false", "off", "no", "")
        config.setdefault("security", {}).setdefault("firewall", {})["enabled"] = _enabled

    # Validate configuration before proceeding
    from core.startup_checks import run_startup_checks
    run_startup_checks(config)

    # Bind to Tailscale interface if available
    ts_ip = get_tailscale_ip()
    if ts_ip != "0.0.0.0":  # nosec B104
        config["server"]["host"] = ts_ip
        logger.info(f"Binding to Tailscale interface: {ts_ip}")

    # Initialize store
    from store.factory import StorageFactory
    store = StorageFactory.get_repository(config_file)
    await store.init()

    # Initialize tracing (optional)
    obs_config = config.get("observability", {}).get("tracing", {})
    if obs_config.get("enabled"):
        sentry_dsn = None
        sentry_cfg = config.get("observability", {}).get("sentry", {})
        dsn_env = sentry_cfg.get("dsn_env")
        if dsn_env:
            sentry_dsn = os.environ.get(dsn_env)
        TraceManager.initialize(
            service_name=obs_config.get("service_name", "llmproxy"),
            otlp_endpoint=obs_config.get("otlp_endpoint"),
            console_export=obs_config.get("console_export", False),
            sentry_dsn=sentry_dsn,
        )

    # Start metrics server (optional)
    metrics_cfg = config.get("server", {}).get("metrics", {})
    if metrics_cfg.get("enabled"):
        start_metrics_server(port=metrics_cfg.get("port", 9091))

    # Launch the security gateway
    rotator = ProxyOrchestrator(store)
    port = config.get("server", {}).get("port", 8090)
    logger.info(f"LLMProxy Security Gateway starting on port {port}")

    await rotator.run(port=port)


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("Shutdown requested")