PacketStation/server/packetstation/cli.py

"""`packetstation` command line: app, worker, migrate, doctor, version, healthcheck."""

from __future__ import annotations

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

from . import __version__
from .settings import Settings

log = logging.getLogger("packetstation")

DEFAULT_BEAT_FILE = "/tmp/ps-worker.beat"  # noqa: S108 (contract default; containers have a private /tmp)
HEALTH_URL = "http://127.0.0.1:8000/healthz"


def _settings() -> Settings:
    from .settings import get_settings

    return get_settings()


def _setup_logging(level: str | None = None) -> None:
    from . import logs

    logs.setup(level or os.environ.get("LOG_LEVEL", "INFO").upper())


def cmd_app(_: argparse.Namespace) -> int:
    import uvicorn

    s = _settings()
    _setup_logging(s.log_level)
    uvicorn.run(
        "packetstation.main:app",
        host="0.0.0.0",
        port=8000,
        proxy_headers=False,
        forwarded_allow_ips="",
        server_header=False,
        access_log=False,
        log_config=None,
    )
    return 0


def cmd_worker(_: argparse.Namespace) -> int:
    from . import guard, worker

    s = _settings()
    _setup_logging(s.log_level)
    guard.startup_checks(s)
    beat = Path(os.environ.get("WORKER_BEAT_FILE") or s.worker_beat_file)
    return worker.run(beat, s.ps_build)


def cmd_migrate(a: argparse.Namespace) -> int:
    from . import db, migrate

    s = _settings()
    _setup_logging(s.log_level)
    try:
        if a.bootstrap_roles:
            for item in migrate.bootstrap_roles(s):
                log.info("bootstrap created %s", item)
        with db.connect_as("migrate", s, autocommit=True) as conn:
            done = migrate.run(conn, db.MIGRATIONS_DIR)
    except migrate.MigrationError as e:
        log.error("migrate refused: %s", e)
        return 1
    log.info("migrations applied: %s", ", ".join(done) if done else "none (at head)")
    return 0


def cmd_doctor(a: argparse.Namespace) -> int:
    from . import doctor

    _setup_logging("WARNING")
    checks = doctor.run(skip=doctor.parse_skip(a.skip), send_test=a.send_test)
    if a.json:
        print(json.dumps(doctor.as_json(checks), indent=2))
    else:
        print(doctor.report(checks))
    return 1 if doctor.failed(checks) else 0


def cmd_version(_: argparse.Namespace) -> int:
    print(f"packetstation {__version__} build {os.environ.get('PS_BUILD', 'dev')}")
    return 0


def cmd_healthcheck(a: argparse.Namespace) -> int:
    """Exit 0 when healthy. Reads the environment directly so it works without full settings."""
    if a.target == "app":
        import httpx

        try:
            r = httpx.get(HEALTH_URL, timeout=3.0)
            return 0 if r.status_code == 200 and r.json().get("ok") is True else 1
        except (httpx.HTTPError, ValueError):
            return 1
    from .worker import STALE_SECONDS, beat_age

    age = beat_age(Path(os.environ.get("WORKER_BEAT_FILE") or DEFAULT_BEAT_FILE))
    return 0 if age is not None and age < STALE_SECONDS else 1


def build_parser() -> argparse.ArgumentParser:
    ap = argparse.ArgumentParser(prog="packetstation", description="Packet Station server")
    sub = ap.add_subparsers(dest="cmd", required=True)
    sub.add_parser("app", help="Run the web app (uvicorn on 0.0.0.0:8000)").set_defaults(func=cmd_app)
    sub.add_parser("worker", help="Run the background worker").set_defaults(func=cmd_worker)
    m = sub.add_parser("migrate", help="Apply pending SQL migrations as the migrate role")
    m.add_argument(
        "--bootstrap-roles",
        action="store_true",
        help="dev/test/ci: create roles and database first, as PG_SUPERUSER_URL",
    )
    m.set_defaults(func=cmd_migrate)
    d = sub.add_parser("doctor", help="Check configuration and every dependency")
    d.add_argument("--json", action="store_true", help="Print the checks as a JSON list")
    d.add_argument("--skip", metavar="LIST", help="Comma list of: visual,mfiles,ldap,smtp,db,volume")
    d.add_argument("--send-test", metavar="ADDR", help="Also send one test message to ADDR")
    d.set_defaults(func=cmd_doctor)
    sub.add_parser("version", help="Print the version and build").set_defaults(func=cmd_version)
    h = sub.add_parser("healthcheck", help="Container healthcheck (exit 0 healthy, 1 not)")
    h.add_argument("target", choices=["app", "worker"])
    h.set_defaults(func=cmd_healthcheck)
    return ap


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    try:
        return int(args.func(args))
    except KeyboardInterrupt:
        return 130
    except Exception as e:
        from pydantic import ValidationError

        from .guard import StackGuardError
        from .settings import SecretFileError

        if isinstance(e, (ValidationError, StackGuardError, SecretFileError)):
            print(f"packetstation {args.cmd}: {e}", file=sys.stderr)
            return 2
        raise


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