Viewing File: /opt/imunify360/venv/versions/imunify-core-8.14.0-1/defence360agent/subsys/nats_discovery.py

"""Shared NATS address + token discovery for asyncclient.

All NATS clients in the asyncclient (and the K8s integration script)
need to find the embedded broker on the same convention: addr written
by the resident-agent at startup, env override, then a localhost
fallback. Originally each caller carried its own copy of this logic
(`NATSGatewayAPI._read_addr` in api/server/send_message.py,
`imav.malwarelib.cleanup.storage` module constants, `scripts/cluster_integration.py`,
plus the agent watchdog). New callers should import these helpers
instead of growing yet another copy; pre-existing copies will be
migrated incrementally so that a change to the discovery chain (e.g.
adding a new env variable) is a single edit.
"""

import os

DEFAULT_NATS_ADDR_PATH = "/var/run/imunify360/nats.addr"
DEFAULT_NATS_TOKEN_PATH = "/var/run/imunify360/nats.token"
DEFAULT_NATS_PORT = "44222"


def read_nats_addr() -> str:
    """Resolve the embedded NATS broker address.

    Discovery chain (shared with all asyncclient NATS clients so they
    all land on the same broker):
      1. addr file written by the resident-agent at startup
         (path overridable via ``I360_NATS_ADDR_PATH``)
      2. ``I360_NATS_DEFAULT_ADDR`` (full ``host:port`` — for cross-pod K8s)
      3. ``127.0.0.1:I360_NATS_PORT`` (default port 44222)
    """
    addr_path = os.getenv("I360_NATS_ADDR_PATH", DEFAULT_NATS_ADDR_PATH)
    try:
        with open(addr_path) as f:
            addr = f.read().strip()
        if addr:
            return addr
    except OSError:
        pass
    default_addr = os.getenv("I360_NATS_DEFAULT_ADDR")
    if default_addr:
        return default_addr
    port = os.getenv("I360_NATS_PORT", DEFAULT_NATS_PORT)
    return f"127.0.0.1:{port}"


def nats_token_path() -> str:
    """Where the resident-agent writes the broker's auth token.

    Exposed separately because the watchdog probes the file's existence to
    decide whether the agent started NATS at all, rather than reading it.
    """
    return os.getenv("I360_NATS_TOKEN_PATH", DEFAULT_NATS_TOKEN_PATH)


def read_nats_token() -> str:
    """Read the embedded NATS broker auth token.

    Path is the token file written by the resident-agent at startup,
    overridable via ``I360_NATS_TOKEN_PATH``. Raises ``OSError`` on
    missing / unreadable file — callers should let it propagate so the
    failure is observable rather than silently degraded.
    """
    with open(nats_token_path()) as f:
        return f.read().strip()