22

✏️ Editing: cluster_integration.py

#!/opt/imunify360/venv/bin/python3
"""
Cluster integration script for Imunify360.

Reads active applications from the NATS JetStream KV bucket that the
im360-k8s-syncer publishes to, and renders the result in the
integration.conf format the panel expects.

This replaces the old sqlite-via-resident-agent path: the
registration_grpc plugin and its active_applications table are gone;
the KV bucket is the source of truth.

The read is resilient to transient NATS glitches. nats-py's ``kv.keys()``
walks a watcher and can return a *short* list with no error raised when a
key lands on the watcher's init boundary while the syncer is concurrently
republishing every entry. A short list, taken as authoritative, has made
the agent reject live apps ("not a registered cluster application") and strip
their feature permissions. To prevent that:

* every read is validated against the bucket's authoritative live-key
    count (``BucketStatus.values``) and re-listed before being trusted;
  * a confirmed-complete read is cached on local disk;
  * a read that cannot be confirmed complete (short count, failed key
    fetch, or unreachable broker) serves the last-good cached snapshot
    instead of a degraded list, and reports the degradation to Sentry.

Usage:
    cluster_integration.py users    - Returns list of users (app_ids)
    cluster_integration.py domains  - Returns domain-to-owner mapping
"""
import asyncio
import json
import os
import sys
import tempfile
import time
from configparser import ConfigParser

try:
    import nats
    from nats.js.errors import BucketNotFoundError, NoKeysError
except ImportError:  # pragma: no cover - runtime environment always has it
    nats = None
    BucketNotFoundError = Exception
    NoKeysError = Exception

INTEGRATION_CONF = "/etc/sysconfig/imunify360/integration.conf"
DEFAULT_BASEDIR = "/home"

# NATS connection — same addr + token files the resident-agent writes
# at startup, same convention the asyncclient NATSGatewayAPI uses.
NATS_ADDR_PATH = os.getenv(
    "I360_NATS_ADDR_PATH", "/var/run/imunify360/nats.addr"
)
NATS_TOKEN_PATH = os.getenv(
    "I360_NATS_TOKEN_PATH", "/var/run/imunify360/nats.token"
)
NATS_DEFAULT_ADDR = os.getenv(
    "I360_NATS_DEFAULT_ADDR",
    f"127.0.0.1:{os.getenv('I360_NATS_PORT', '44222')}",
)
REGISTRATION_BUCKET = os.getenv("I360_REGISTRATION_KV_BUCKET", "REGISTRATION")
CONNECT_TIMEOUT = 5
# Whole-read bound. CONNECT_TIMEOUT covers only the handshake: nats-py can
# park its reconnect machinery on the loop instead of raising, and an
# unbounded await there deadlocks first boot, since the migration that calls
# this runs before the resident agent that serves NATS exists.
FETCH_TIMEOUT = int(os.getenv("I360_REGISTRATION_FETCH_TIMEOUT", "30"))

# Local last-good snapshot. Lives on the persisted (PVC-backed) data dir,
# NOT tmpfs, so it survives pod restarts and migrations — the no-cache
# branch in _fetch_registrations then only fires on a genuine first run
# before NATS has ever been reachable. Overridable for tests.
CACHE_PATH = os.getenv(
    "I360_REGISTRATION_CACHE_PATH",
    "/var/imunify360/cluster_registration_cache.json",
)
# How many times to re-list within a single connection before giving up on
# a complete read and falling back to cache. The dominant failure is a
# transient short kv.keys(); a re-list almost always returns the full set.
LIST_ATTEMPTS = int(os.getenv("I360_REGISTRATION_LIST_ATTEMPTS", "3"))
# Minimum seconds between Sentry reports — the panel polls this script, so
# unthrottled degraded reports would spam a single Sentry issue.
SENTRY_COOLDOWN = int(os.getenv("I360_REGISTRATION_SENTRY_COOLDOWN", "300"))

class IncompleteReadError(Exception):
    """The live KV read could not be confirmed complete.

Raised for an unreachable broker, a failed listing/fetch, or a key
    count short of the bucket's authoritative entry count. The caller
    serves the cached snapshot instead of trusting a degraded read.
    """

def get_basedir():
    """Read basedir from integration.conf [malware] section."""
    cfg = ConfigParser()
    cfg.read(INTEGRATION_CONF)
    return cfg.get("malware", "basedir", fallback=DEFAULT_BASEDIR)

def _ok_response(data):
    """Wrap data with metadata expected by panel validation."""
    return {"data": data, "metadata": {"result": "ok"}}

def _error_response(message):
    """Return error response with metadata."""
    return {"data": None, "metadata": {"result": "error", "message": message}}

def _read_first_line(path):
    try:
        with open(path) as f:
            return f.read().strip()
    except OSError:
        return ""

async def _close(nc):
    try:
        await nc.drain()
    except Exception:
        pass
    try:
        await nc.close()
    except Exception:
        pass

async def _list_once(kv):
    """One listing pass, validated against the bucket's live-key count.

The count cross-check (``BucketStatus.values``) is what catches a
    silently truncated ``kv.keys()``: the syncer bucket is History=1 with
    no deletes, so the stream message count equals the number of live
    keys. A failed get/parse or a count short of that authoritative number
    both raise ``IncompleteReadError`` rather than yield a partial list.

(If the bucket config ever gained history>1 or delete tombstones,
    ``values`` would exceed the live-key count and we would fall back to
    cache more often — never publish a short list. The safe direction.)
    """
    try:
        status = await kv.status()
        expected = status.values
    except Exception as exc:
        raise IncompleteReadError("status: %s" % exc)

try:
        keys = await kv.keys()
    except NoKeysError:
        keys = []
    except Exception as exc:
        raise IncompleteReadError("keys: %s" % exc)

result = {}
    for key in keys:
        try:
            entry = await kv.get(key)
            payload = json.loads(entry.value)
        except Exception as exc:
            # A listed-but-unreadable key, or a corrupt value, means we
            # cannot represent this app — treat the snapshot as incomplete
            # rather than silently dropping the app from the list.
            raise IncompleteReadError("key %s: %s" % (key, exc))
        app_id = payload.get("app_id") or key
        result[app_id] = {"domains": payload.get("domains") or {}}

# Completeness is about how many live KEYS we fetched vs the bucket's
    # authoritative count — not how many distinct app_ids ended up in
    # ``result``. Keying ``result`` by app_id can collapse two keys into one
    # entry (e.g. a payload whose app_id duplicates another's), which would
    # make a full read look short and wedge us onto cache. Count keys.
    if len(keys) < expected:
        raise IncompleteReadError(
            "short read: %d of %d entries" % (len(keys), expected)
        )
    return result

async def _read_registrations_live():
    """Return {app_id: {"domains": {...}}} from a confirmed-complete read.

Raises ``IncompleteReadError`` when the snapshot cannot be confirmed
    complete, so the caller can fall back to the cached snapshot rather
    than publish a degraded list.
    """
    if nats is None:
        raise IncompleteReadError("nats-py not installed")

addr = _read_first_line(NATS_ADDR_PATH) or NATS_DEFAULT_ADDR
    token = _read_first_line(NATS_TOKEN_PATH)

connect_kwargs = {
        "servers": [f"nats://{addr}"],
        "connect_timeout": CONNECT_TIMEOUT,
        "max_reconnect_attempts": 0,
    }
    if token:
        connect_kwargs["token"] = token

# A connect failure fails fast to the cache — retrying a 5s connect
    # would just stall the agent's integration call during an outage.
    try:
        nc = await nats.connect(**connect_kwargs)
    except Exception as exc:
        raise IncompleteReadError("connect: %s" % exc)

try:
        js = nc.jetstream()
        try:
            kv = await js.key_value(REGISTRATION_BUCKET)
        except BucketNotFoundError as exc:
            # A missing bucket is degraded, not a confirmed-empty read: with a
            # good cache we serve it (a transient bucket gap must not strip
            # live apps); on a cold cache the caller seeds empty so a genuine
            # first run before the bucket exists still proceeds.
            raise IncompleteReadError("bucket not found: %s" % exc)

# Re-list on the same connection: a transient short kv.keys() is
        # the failure we are guarding against, and a fresh listing almost
        # always returns the full set.
        last_err = None
        for _ in range(LIST_ATTEMPTS):
            try:
                return await _list_once(kv)
            except IncompleteReadError as exc:
                last_err = exc
        raise last_err or IncompleteReadError("listing failed")
    finally:
        await _close(nc)

def _load_cache():
    try:
        with open(CACHE_PATH) as f:
            data = json.load(f)
    except (OSError, ValueError):
        return None
    if not isinstance(data, dict) or "registrations" not in data:
        return None
    return data

def _atomic_write(path, payload):
    directory = os.path.dirname(path) or "."
    os.makedirs(directory, exist_ok=True)
    fd, tmp = tempfile.mkstemp(dir=directory, prefix=".reg_cache.")
    try:
        with os.fdopen(fd, "w") as f:
            f.write(payload)
        os.replace(tmp, path)
    except Exception:
        try:
            os.unlink(tmp)
        except OSError:
            pass
        raise

def _store_cache(registrations, last_sentry_at=None):
    """Persist a confirmed-complete snapshot as the new last-good cache.

Resets the Sentry throttle: a fresh successful read ends the degraded
    episode, so the next degradation reports promptly.
    """
    _atomic_write(
        CACHE_PATH,
        json.dumps(
            {
                "updated_at": time.time(),
                "last_sentry_at": last_sentry_at,
                "registrations": registrations,
            }
        ),
    )

def _touch_sentry_time(cache, now):
    """Bump the throttle timestamp without disturbing the cached snapshot."""
    updated = dict(cache)
    updated["last_sentry_at"] = now
    _atomic_write(CACHE_PATH, json.dumps(updated))

def _report_sentry(message, level, fingerprint):
    """Best-effort Sentry report. Never let it break integration output."""
    try:
        from defence360agent import sentry

sentry.configure_sentry()
        sentry.log_message(
            message,
            level=level,
            fingerprint=fingerprint,
            component="cluster_integration",
        )
        sentry.flush_sentry()
    except Exception:
        pass

def _maybe_report_degraded(cache, reason):
    now = time.time()
    last = cache.get("last_sentry_at") or 0
    if now - last < SENTRY_COOLDOWN:
        return
    try:
        _touch_sentry_time(cache, now)
    except OSError:
        # No persisted stamp means no cooldown to rate-limit on, so the report
        # would repeat on every poll. Skip it rather than flood Sentry.
        return
    age = now - (cache.get("updated_at") or now)
    _report_sentry(
        "cluster_integration: degraded NATS read (%s); served %d cached "
        "registrations (age %.0fs)"
        % (reason, len(cache.get("registrations") or {}), age),
        level="warning",
        fingerprint="cluster-integration-degraded-read",
    )

async def _read_registrations_bounded():
    try:
        return await asyncio.wait_for(
            _read_registrations_live(), FETCH_TIMEOUT
        )
    except asyncio.TimeoutError:
        raise IncompleteReadError(
            "live read exceeded %ss" % FETCH_TIMEOUT
        ) from None

async def _fetch_registrations():
    """Return ({app_id: {"domains": {...}}}, err).

err == "" → authoritative result: either a confirmed-complete live
    read or, on a degraded read, the last-good cached snapshot.
    err != "" → no complete read and no cache to serve; the caller
    surfaces this as a hard error (non-zero exit).
    """
    try:
        registrations = await _read_registrations_bounded()
    except IncompleteReadError as exc:
        cache = _load_cache()
        if cache is not None:
            _maybe_report_degraded(cache, str(exc))
            return cache["registrations"], ""
        # No cache yet: this is a first run / migration before NATS is up.
        # The agent calls panel_users() during migration; erroring here
        # would break it. Seed an empty persisted cache and return empty
        # data (a valid "no apps yet" state) so migration proceeds. The
        # first successful read fills the cache in, and because it is
        # persisted we never take this branch again.
        try:
            # Seed the cache stamped as just-reported so the next poll takes
            # the throttled warm path instead of re-reporting.
            _store_cache({}, last_sentry_at=time.time())
            _report_sentry(
                "cluster_integration: NATS unreachable on a cold cache; "
                "seeding empty registrations: %s" % exc,
                level="warning",
                fingerprint="cluster-integration-no-cache",
            )
        except OSError:
            # Cannot persist the seed, so we cannot throttle either: skip the
            # report rather than flood Sentry on every panel poll.
            pass
        return {}, ""
    try:
        _store_cache(registrations)
    except OSError:
        # The live read is authoritative; a best-effort cache write failure
        # must not discard it and force callers to fall back.
        pass
    return registrations, ""

def _fetch():
    return asyncio.run(_fetch_registrations())

def get_users():
    """
    Return users in integration.conf format.

Format: {"data": [{"username": "app_id", "home": "/home"}, ...],
             "metadata": {"result": "ok"}}
    """
    apps, err = _fetch()
    if err:
        return _error_response(err)

basedir = get_basedir()
    users = [
        {"username": app_id, "home": basedir} for app_id in sorted(apps.keys())
    ]
    return _ok_response(users)

def get_domains():
    """
    Return domains in integration.conf format.

Format: {"data": {"domain.com": {"owner": "app_id",
                                     "document_root": "/home"}, ...},
             "metadata": {"result": "ok"}}
    """
    apps, err = _fetch()
    if err:
        return _error_response(err)

basedir = get_basedir()
    domains_data = {}
    for app_id, info in apps.items():
        for domain in info.get("domains", {}):
            domains_data[domain] = {
                "owner": app_id,
                "document_root": basedir,
            }
    return _ok_response(domains_data)

def main():
    if len(sys.argv) < 2:
        print("Usage: cluster_integration.py [users|domains]", file=sys.stderr)
        sys.exit(1)

command = sys.argv[1]

if command == "users":
        result = get_users()
    elif command == "domains":
        result = get_domains()
    else:
        print(f"Unknown command: {command}", file=sys.stderr)
        sys.exit(1)

print(json.dumps(result))

metadata = result.get("metadata", {})
    if metadata.get("result") != "ok":
        # Defensive net. The fetch path no longer errors on an unreachable
        # broker — it serves cache or seeds empty (ok) so migration is not
        # broken. A non-ok result here therefore means something genuinely
        # malformed, which must be a hard failure rather than an exit-0
        # body the agent could mistake for a benign empty result. The agent
        # runs us via check_run(): on a non-zero exit it raises
        # CheckRunError *before* parsing stdout and records the error from
        # our *stderr* (see generic/panel.py:_get_integration_data) — so the
        # detail has to go to stderr; stdout is never read once the exit
        # code is non-zero.
        print(
            metadata.get("message") or "integration script error",
            file=sys.stderr,
        )
        sys.exit(1)

if __name__ == "__main__":
    main()

MR SKB