import base64
import collections
import hashlib
import json
import os
import time
import urllib.error
try:
import nats
import nats.errors
import nats.js.errors
_has_nats = True
_NATSMaxPayloadError = nats.errors.MaxPayloadError
_NATSAPIError = nats.js.errors.APIError
except ImportError:
_has_nats = False
class _NATSMaxPayloadError(Exception):
pass
class _NATSAPIError(Exception):
err_code = None
import urllib.request
from abc import ABC, abstractmethod
from logging import getLogger
from typing import Optional
import asyncio
import uuid
from defence360agent.api.server import (
API,
APIError,
APITokenError,
FGWSendMessgeException,
NATSSendMessageException,
)
from defence360agent.contracts.config import Core
from defence360agent.contracts.messages import estimate_size, Message
from defence360agent.internals import delivery_ack
from defence360agent.internals.feature_flags import (
MESSAGE_LOSS_OBSERVABILITY_FLAG,
is_enabled,
)
from defence360agent.internals.global_scope import g
from defence360agent.internals.iaid import (
IndependentAgentIDAPI,
IAIDTokenError,
)
from defence360agent.internals.message_status_publisher import Gen, publisher
from defence360agent.utils.async_utils import AsyncIterate
from defence360agent.utils.json import ServerJSONEncoder
logger = getLogger(__name__)
_reporter_gen_fgw = Gen()
_reporter_gen_nats = Gen()
_method_missing_dropped_total = 0
_method_missing_dropped_delta = 0
def pop_method_missing_dropped() -> int:
"""Method-less drops since the last call, then reset (delta)."""
global _method_missing_dropped_delta
value, _method_missing_dropped_delta = _method_missing_dropped_delta, 0
return value
def _drop_method_less(message: dict, sink: str) -> None:
"""Count and log a message no sink can route."""
# Key names only, except plugin_id: it names the producer, not the payload.
global _method_missing_dropped_total, _method_missing_dropped_delta
_method_missing_dropped_total += 1
_method_missing_dropped_delta += 1
logger.error(
"Dropping message without a method: sink=%s message_id=%s"
" plugin_id=%s timestamp=%s keys=%s dropped_total=%d",
sink,
message.get("message_id"),
message.get("plugin_id"),
message.get("timestamp"),
sorted(message),
_method_missing_dropped_total,
)
# Returned for both "maximum messages exceeded" and "maximum bytes exceeded"
# once the stream is full; reaches the client only because the stream discards
# new rather than old messages.
_JS_ERR_STREAM_FULL = 10077
# The stream's own MaxMsgSize rejection, distinct from the client-side
# MaxPayloadError checked against the server's max_payload. Both caps are 10MB
# today, so this only fires if MaxMsgSize is lowered below max_payload.
_JS_ERR_MSG_TOO_LARGE = 10054
class _StreamFull(Exception):
"""Stream is at capacity: re-queue the rest, the connection is healthy."""
# Keys the receiving side indexes a fragment on. A field carrying one of them
# identifies the message instead of holding its payload: bisecting it strands
# fragments the store path cannot key, so it is copied into every fragment.
_INDEX_KEYS = frozenset({"scanid", "scan_id"})
# Single-element descents before a message is called irreducible: json.loads
# accepts nesting deeper than this mutual recursion can safely walk.
_MAX_SPLIT_DEPTH = 32
def _element_ids(container) -> set:
"""Strings a sibling map could be keyed on: a dict's own keys, or the
scalar values carried by a list's elements."""
if isinstance(container, dict):
return {key for key in container if isinstance(key, str)}
ids = set()
for element in container:
if isinstance(element, str):
ids.add(element)
elif isinstance(element, dict):
ids.update(
value for value in element.values() if isinstance(value, str)
)
return ids
def _paired_maps(fields: dict) -> dict:
"""Sibling dicts keyed entirely by a field's elements, which have to follow
those records rather than be bisected away from them."""
ids = {name: _element_ids(value) for name, value in fields.items()}
return {
name: {
other
for other, value in fields.items()
if other != name
and isinstance(value, dict)
and value
and set(value) <= ids[name]
}
for name in fields
}
def _carries_index_key(value) -> bool:
return isinstance(value, dict) and not _INDEX_KEYS.isdisjoint(value)
def _split_element(element, depth):
if depth <= 0:
return None
if isinstance(element, dict):
return _split_largest_field(element, depth - 1)
if isinstance(element, list):
return _split_container(element, depth - 1)
return None
def _split_container(value, depth):
"""Bisect a list or dict, descending into a lone element so a single
oversized record can still be split within itself."""
if isinstance(value, dict):
keys = list(value)
if len(keys) > 1:
mid = len(keys) // 2
return [
{key: value[key] for key in keys[:mid]},
{key: value[key] for key in keys[mid:]},
]
if not keys:
return None
inner = _split_element(value[keys[0]], depth)
if inner is None:
return None
return [{keys[0]: half} for half in inner]
if len(value) > 1:
mid = len(value) // 2
return [value[:mid], value[mid:]]
if not value:
return None
inner = _split_element(value[0], depth)
if inner is None:
return None
return [[half] for half in inner]
def _split_largest_field(item: dict, depth=_MAX_SPLIT_DEPTH):
"""Bisect the heaviest payload field, while index-carrying fields and maps
paired with another field's records ride along instead of being split."""
# ponytail: pairing is inferred from depth-1 scalars, not read from the
# producer's declared batch field, so a nested or renamed join key quietly
# degrades to duplicating records; byte-bound the producers to retire this.
fields = {
name: value
for name, value in item.items()
if isinstance(value, (list, dict))
}
paired = _paired_maps(fields)
followers = {name for maps in paired.values() for name in maps}
candidates = [
name
for name, value in fields.items()
if name not in followers and not _carries_index_key(value)
]
def weight(name):
return estimate_size(fields[name]) + sum(
estimate_size(fields[other]) for other in paired[name]
)
# Heaviest first, but fall through to the next candidate when the heaviest
# cannot be bisected: sorted() is stable, so a tie keeps insertion order
# and a re-split reproduces the same fragments and dedup ids.
for field in sorted(candidates, key=weight, reverse=True):
halves = _split_container(fields[field], depth)
if halves is not None:
break
else:
return None
parts = []
for half in halves:
kept = _element_ids(half)
part = {**item, field: half}
for name in paired[field]:
part[name] = {
key: value
for key, value in fields[name].items()
if key in kept
}
parts.append(part)
return parts
def _split_oversized(loaded: dict):
"""Split an oversized message into smaller parts. Returns a list of parts,
or None when the message carries a single irreducible record."""
items = loaded.get("items")
if isinstance(items, list) and len(items) > 1:
mid = len(items) // 2
return [
{**loaded, "items": items[:mid]},
{**loaded, "items": items[mid:]},
]
single = items[0] if isinstance(items, list) and len(items) == 1 else None
if isinstance(single, dict):
try:
halves = _split_largest_field(single)
except RecursionError:
# estimate_size walks the item too, and json.loads accepts nesting
# deeper than it can measure. Unsplittable beats re-queueing the
# batch forever on a message no depth cap of ours can save.
return None
if halves is not None:
return [{**loaded, "items": [half]} for half in halves]
return None
async def _nats_error_cb(ex: Exception) -> None:
"""Downgrade nats-py internal errors to DEBUG.
Transient errors (ConnectionRefused, AuthorizationViolation) are
expected during agent restarts. Our code already logs a WARNING
with context, so the nats-py default ERROR + traceback is noise.
"""
logger.debug("nats: %s", ex)
class BaseSendMessageAPI(API, ABC):
URL = "/api/v2/send-message/{method}"
@abstractmethod
async def _send_request(self, message_method, headers, post_data) -> dict:
pass # pragma: no cover
def check_response(self, result: dict) -> None:
if "status" not in result:
raise APIError("unexpected server response: {!r}".format(result))
if result["status"] != "ok":
raise APIError("server error: {}".format(result.get("msg")))
async def send_data(self, method: str, post_data: bytes) -> None:
try:
token = await IndependentAgentIDAPI.get_token()
except IAIDTokenError as e:
raise APITokenError(f"IAID token error occurred {e}")
headers = {
"Content-Type": "application/json",
"X-Auth": token,
}
result = await self._send_request(method, headers, post_data)
self.check_response(result)
class SendMessageAPI(BaseSendMessageAPI):
_SOCKET_TIMEOUT = Core.DEFAULT_SOCKET_TIMEOUT
def __init__(self, rpm_ver: str, base_url: str = None, executor=None):
self._executor = executor
self.rpm_ver = rpm_ver
self.product_name = ""
self.server_id = None # type: Optional[str]
self.license = {} # type: dict
if base_url:
self.base_url = base_url
else:
self.base_url = self._BASE_URL
def set_product_name(self, product_name: str) -> None:
self.product_name = product_name
def set_server_id(self, server_id: Optional[str]) -> None:
self.server_id = server_id
def set_license(self, license: dict) -> None:
self.license = license
async def _send_request(self, message_method, headers, post_data):
request = urllib.request.Request(
self.base_url + self.URL.format(method=message_method),
data=post_data,
headers=headers,
method="POST",
)
return await self.async_request(request, executor=self._executor)
async def send_message(self, message: Message) -> None:
# Return, don't raise: raising re-queues the entry at the backlog head.
if not message.get("method"):
_drop_method_less(message, "http")
return
# add message handling time if it does not exist, so that
# the server does not depend on the time it was received
if "timestamp" not in message:
message["timestamp"] = time.time()
if "message_id" not in message:
message["message_id"] = uuid.uuid4().hex
data2send = {
"payload": message.payload,
"rpm_ver": self.rpm_ver,
"message_id": message.message_id,
"server_id": self.server_id,
"name": self.product_name,
}
post_data = json.dumps(data2send, cls=ServerJSONEncoder).encode()
await self.send_data(message.method, post_data)
class FileBasedGatewayAPI(SendMessageAPI):
async def _prepare_message(self, message, semaphore) -> dict:
async with semaphore:
loaded = await asyncio.to_thread(json.loads, message)
return {
"method": loaded["method"],
"data": {k: v for k, v in loaded.items() if k != "method"},
}
async def send_messages(self, messages: list[tuple[float, bytes]]) -> None:
max_threads = 5
semaphore = asyncio.Semaphore(max_threads)
tasks = [
self._prepare_message(msg, semaphore)
async for _, msg in AsyncIterate(messages)
]
prepared_messages = await asyncio.gather(*tasks)
for msg in prepared_messages:
flat = {**msg.get("data", {}), "method": msg.get("method", "")}
publisher.report(
flat, _reporter_gen_fgw, stage="agent-fgw-sending"
)
dumped_messages = await asyncio.to_thread(
json.dumps, prepared_messages
)
bin_file_path = os.getenv(
"I360_MESSAGE_GATEWAY_BIN_PATH", "/usr/libexec/"
)
bin_file = os.path.join(bin_file_path, "imunify-message-gateway")
command = [
bin_file,
"send-many",
"--producer=i360-agent-non-resident",
]
process = await asyncio.create_subprocess_exec(
*command,
stdin=asyncio.subprocess.PIPE,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
)
b64data = base64.b64encode(dumped_messages.encode())
stdout, stderr = await process.communicate(input=b64data)
if g.get("DEBUG"):
logger.info(
"Message sent to fgw: %s %s %s", len(messages), stdout, stderr
)
if process.returncode != 0:
logger.error(f"Error sending message: {stderr.decode()}")
raise FGWSendMessgeException(
str(f"Error sending message: {stderr.decode()}")
)
for msg in prepared_messages:
delivery_ack.registry.confirm(msg["data"].get("message_id"))
class NATSGatewayAPI:
"""Publishes messages to the embedded NATS server via localhost TCP.
Connects to nats://127.0.0.1:<port> with an auth token read from
a file written by the resident-agent on startup.
"""
NATS_SUBJECT_PREFIX = "imunify.api."
DEFAULT_PORT = 44222
DEFAULT_TOKEN_PATH = "/var/run/imunify360/nats.token"
DEFAULT_ADDR_PATH = "/var/run/imunify360/nats.addr"
CONNECT_TIMEOUT = 5
MIN_RECONNECT_INTERVAL = 5
def __init__(self):
self._nc = None
self._last_connect_attempt = 0
self._oversized_dropped = 0
self._oversized_rejected = 0
def pop_oversized_rejected(self) -> int:
"""Oversized rejections since the last call, then reset (delta)."""
value, self._oversized_rejected = self._oversized_rejected, 0
return value
@staticmethod
def _read_addr():
"""Read NATS listen address from addr file, fall back to env/default."""
addr_path = os.getenv(
"I360_NATS_ADDR_PATH", NATSGatewayAPI.DEFAULT_ADDR_PATH
)
try:
with open(addr_path) as f:
addr = f.read().strip()
if addr:
return addr
except OSError:
pass
# Fallback: env var / hardcoded default (for upgrades where
# the resident-agent hasn't written the addr file yet)
port = int(
os.getenv("I360_NATS_PORT", str(NATSGatewayAPI.DEFAULT_PORT))
)
return f"127.0.0.1:{port}"
async def _ensure_connected(self):
if self._nc is not None and self._nc.is_connected:
return
if not _has_nats:
raise NATSSendMessageException("nats-py is not installed")
now = time.monotonic()
since_last = now - self._last_connect_attempt
if since_last < self.MIN_RECONNECT_INTERVAL:
raise NATSSendMessageException(
"NATS reconnect backoff"
f" ({self.MIN_RECONNECT_INTERVAL - since_last:.1f}s remaining)"
)
self._last_connect_attempt = now
# Clean up stale connection before reconnecting
await self._close()
addr = self._read_addr()
token_path = os.getenv("I360_NATS_TOKEN_PATH", self.DEFAULT_TOKEN_PATH)
try:
with open(token_path) as f:
token = f.read().strip()
self._nc = await nats.connect(
f"nats://{addr}",
token=token,
connect_timeout=self.CONNECT_TIMEOUT,
max_reconnect_attempts=0,
error_cb=_nats_error_cb,
)
logger.info("Connected to NATS at %s", addr)
except Exception as e:
raise NATSSendMessageException(
f"Failed to connect to NATS: {e}"
) from e
async def send_messages(self, messages: list[tuple[float, bytes]]) -> None:
await self._ensure_connected()
published = 0
try:
js = self._nc.jetstream()
for _, msg_bytes in messages:
try:
loaded = json.loads(msg_bytes)
except (json.JSONDecodeError, UnicodeDecodeError) as e:
logger.warning("Skipping malformed message: %s", e)
published += 1 # count as handled, not re-queued
continue
if not loaded.get("method"):
_drop_method_less(loaded, "nats")
published += 1 # count as handled, not re-queued
continue
method = loaded.pop("method")
subject = self.NATS_SUBJECT_PREFIX + method
# (part, dedup_id) pairs. dedup_id pins each fragment's
# Nats-Msg-Id deterministically: when a message is split and a
# later fragment fails with a non-payload error, the wrapper
# re-queues the whole original; re-splitting reproduces the
# same fragments and ids, so JetStream de-duplicates the ones
# already delivered instead of duplicating them.
pending = collections.deque(
[(loaded, loaded.get("message_id"))]
)
# (id, size, error, preview) per irreducible part. str(e), not
# the exception: its traceback would pin this frame's payload
# and part until the message is done.
dropped = []
while pending:
part, dedup_id = pending.popleft()
payload = json.dumps(part).encode()
headers = {"Nats-Msg-Id": dedup_id} if dedup_id else None
try:
ack = await js.publish(
subject, payload, headers=headers
)
except (_NATSMaxPayloadError, _NATSAPIError) as e:
if isinstance(e, _NATSAPIError):
if e.err_code == _JS_ERR_STREAM_FULL:
# Abort the batch so this message and the rest
# are re-queued whole; already-published
# fragments carry deterministic ids and are
# de-duplicated on retry.
raise _StreamFull(
f"subject={subject}"
f" message_id={dedup_id}: {e}"
) from e
if e.err_code != _JS_ERR_MSG_TOO_LARGE:
raise
halves = _split_oversized(part)
if halves is None:
# A single record that alone exceeds the limit
# cannot be delivered over NATS. The other
# fragments of this message are still published;
# only this irreducible record is dropped (loudly,
# with a counter). It is intentionally counted as
# handled rather than re-queued, otherwise it would
# block the head of the queue forever.
dropped.append(
(dedup_id, len(payload), str(e), payload[:200])
)
continue
base_key = (
dedup_id or hashlib.sha1(payload).hexdigest()
)
for index, half in enumerate(halves):
child_key = f"{base_key}.{index}"
half["message_id"] = child_key
pending.append((half, child_key))
logger.warning(
"Splitting oversized NATS message: subject=%s"
" message_id=%s size=%d parts=%d",
subject,
dedup_id,
len(payload),
len(halves),
)
continue
if g.get("DEBUG"):
logger.debug(
"Published to %s, stream=%s, seq=%s",
subject,
ack.stream,
ack.seq,
)
# One drop and one status report per logical message, not per
# fragment: a split message's fragments share the parent's
# reporter id, and counting each recursive attempt inflates
# both the delivery-tracking cardinality and the drop metric.
# A message with any dropped fragment is not reported at all:
# it did not fully reach NATS, so tracking it as sent would
# overstate delivery.
if dropped:
self._oversized_dropped += 1
if is_enabled(MESSAGE_LOSS_OBSERVABILITY_FLAG):
self._oversized_rejected += 1
_, _, error, preview = dropped[0]
logger.error(
# size= is the bytes dropped across every part and is
# parsed by the rpm-test that guards this path; keep
# the field name if the format changes again.
"Dropping oversized NATS message: subject=%s"
" message_id=%s parts=%d size=%d fragments=%s"
" oversized_total=%d error=%s preview=%r",
subject,
loaded.get("message_id"),
len(dropped),
sum(size for _, size, _, _ in dropped),
[frag for frag, _, _, _ in dropped],
self._oversized_dropped,
error,
preview,
)
else:
publisher.report(
{**loaded, "method": method},
_reporter_gen_nats,
stage="agent-nats-sending",
)
delivery_ack.registry.confirm(loaded.get("message_id"))
published += 1
except _StreamFull as e:
# Keep the connection: it is healthy, and closing it would make
# recovery wait out MIN_RECONNECT_INTERVAL as well.
logger.warning(
"NATS stream full, %d/%d messages published, %d re-queued: %s",
published,
len(messages),
len(messages) - published,
e,
)
raise NATSSendMessageException(
f"Stream at capacity: {e}",
published=published,
) from e
except Exception as e:
await self._close()
logger.warning(
"NATS publish failed after %d/%d messages: %s",
published,
len(messages),
e,
)
raise NATSSendMessageException(
f"Failed to publish messages: {e}",
published=published,
) from e
async def _close(self):
if self._nc is not None:
try:
# close(), not drain(): drain PINGs the server we already
# consider broken, stalls the send path on the flush timeout,
# and on that timeout leaks the client with its read loop
# alive. Unacked messages are re-queued, so nothing is lost.
await self._nc.close()
except Exception:
pass
finally:
self._nc = None
async def close(self):
await self._close()