xml-pipeline/deploy/entrypoint.py

"""
AgentOS container entrypoint.

Validates config, generates keys if needed, applies security lockdowns,
and boots the organism with dual-port servers (agent + management).

Usage:
    python -m deploy.entrypoint /config/organism.yaml
    python -m deploy.entrypoint --dry-run /config/organism.yaml
"""

from __future__ import annotations

import argparse
import asyncio
import logging
import os
import sys
from pathlib import Path

logger = logging.getLogger("agentos.entrypoint")


def detect_mode() -> str:
    """Detect organism mode from config or environment."""
    return os.environ.get("ORGANISM_MODE", "container")


def validate_config(config_path: Path) -> bool:
    """Validate organism config file exists and is parseable."""
    if not config_path.exists():
        logger.error(f"Config not found: {config_path}")
        return False

    try:
        import yaml

        with open(config_path) as f:
            raw = yaml.safe_load(f)

        if not isinstance(raw, dict):
            logger.error("Config must be a YAML mapping")
            return False

        org = raw.get("organism", {})
        if not org.get("name"):
            logger.error("organism.name is required")
            return False

        logger.info(f"Config valid: {org['name']}")

        # Count listeners
        listeners = raw.get("listeners", [])
        if isinstance(listeners, list):
            logger.info(f"  Listeners: {len(listeners)}")

        return True
    except Exception as e:
        logger.error(f"Config parse error: {e}")
        return False


def ensure_identity_key(config_path: Path) -> None:
    """Generate Ed25519 identity key if not present."""
    import yaml

    with open(config_path) as f:
        raw = yaml.safe_load(f)

    identity_path = raw.get("organism", {}).get("identity")
    if not identity_path:
        return

    identity_file = Path(identity_path)
    if identity_file.exists():
        logger.info(f"Identity key found: {identity_file}")
        return

    try:
        from xml_pipeline.crypto import generate_identity

        identity_file.parent.mkdir(parents=True, exist_ok=True)
        identity = generate_identity()
        public_path = identity_file.with_suffix(".pub")
        identity.save(identity_file, public_path)
        logger.info(f"Generated identity key: {identity_file}")
    except Exception as e:
        logger.warning(f"Could not generate identity key: {e}")


def apply_container_lockdowns(mode: str) -> None:
    """Apply security lockdowns based on organism mode."""
    if mode != "container":
        logger.info(f"Mode '{mode}' — skipping container lockdowns")
        return

    logger.info("Applying container security lockdowns")

    from xml_pipeline.security.defaults import apply_container_defaults

    apply_container_defaults()


async def boot_organism(config_path: Path, mode: str) -> None:
    """Bootstrap and run the organism with dual-port servers."""
    from xml_pipeline.message_bus import bootstrap

    # Bootstrap the pump
    pump = await bootstrap(str(config_path))

    # Determine ports from environment
    agent_port = int(os.environ.get("AGENT_PORT", "8080"))
    management_port = int(os.environ.get("MANAGEMENT_PORT", "9090"))
    host = os.environ.get("BIND_HOST", "0.0.0.0")

    try:
        import uvicorn
    except ImportError:
        logger.error("uvicorn not installed. Install with: pip install xml-pipeline[server]")
        sys.exit(1)

    # Create agent-facing app (minimal: /health, /inject, /ws)
    from xml_pipeline.server.agent_app import create_agent_app

    agent_app = create_agent_app(pump)

    # Create management app (full API, dashboard, audit)
    from xml_pipeline.server.management import create_management_app

    management_app = create_management_app(pump)

    # Configure uvicorn servers
    agent_config = uvicorn.Config(
        agent_app,
        host=host,
        port=agent_port,
        log_level="info",
        access_log=False,
    )
    management_config = uvicorn.Config(
        management_app,
        host="127.0.0.1" if mode == "container" else host,
        port=management_port,
        log_level="info",
    )

    agent_server = uvicorn.Server(agent_config)
    management_server = uvicorn.Server(management_config)

    # Run pump + both servers concurrently
    pump_task = asyncio.create_task(pump.run())

    logger.info(f"Agent bus:   http://{host}:{agent_port}")
    logger.info(f"Management:  http://127.0.0.1:{management_port}")
    logger.info(f"Dashboard:   http://127.0.0.1:{management_port}/dashboard/")

    try:
        await asyncio.gather(
            agent_server.serve(),
            management_server.serve(),
        )
    finally:
        await pump.shutdown()
        pump_task.cancel()
        try:
            await pump_task
        except asyncio.CancelledError:
            pass


def main() -> int:
    """Entrypoint for container boot."""
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
    )

    parser = argparse.ArgumentParser(description="AgentOS entrypoint")
    parser.add_argument("config", nargs="?", default="/config/organism.yaml", help="Config path")
    parser.add_argument("--dry-run", action="store_true", help="Validate config and exit")
    parser.add_argument("--mode", help="Override organism mode (container/development)")
    args = parser.parse_args()

    config_path = Path(args.config)
    mode = args.mode or detect_mode()

    logger.info(f"AgentOS starting (mode={mode})")

    # Validate config
    if not validate_config(config_path):
        return 1

    if args.dry_run:
        logger.info("Dry run complete — config is valid")
        return 0

    # Generate identity key if needed
    ensure_identity_key(config_path)

    # Apply security lockdowns
    apply_container_lockdowns(mode)

    # Boot the organism
    try:
        asyncio.run(boot_organism(config_path, mode))
        return 0
    except KeyboardInterrupt:
        logger.info("Shutdown requested")
        return 0
    except Exception as e:
        logger.error(f"Boot failed: {e}")
        return 1


if __name__ == "__main__":
    sys.exit(main())