File: //opt/imunify360/venv/versions/imunify-core-8.12.0-2/defence360agent/api/server/send_message.py
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.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()
# 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."""
def _split_largest_field(item: dict):
"""Split a single sub-message on its largest list/dict field, chosen by
serialized byte size (not element count) so the heaviest field is the one
that shrinks. Returns two sub-messages, or None when nothing inside can be
split further (no list/dict field holds at least two elements)."""
field = None
largest = 0
for key, value in item.items():
if isinstance(value, (list, dict)) and len(value) > 1:
size = estimate_size(value)
if size > largest:
field, largest = key, size
if field is None:
return None
value = item[field]
if isinstance(value, list):
mid = len(value) // 2
return [{**item, field: value[:mid]}, {**item, field: value[mid:]}]
keys = list(value)
mid = len(keys) // 2
left = {k: value[k] for k in keys[:mid]}
right = {k: value[k] for k in keys[mid:]}
return [{**item, field: left}, {**item, field: right}]
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:]},
]
if isinstance(items, list) and len(items) == 1:
halves = _split_largest_field(items[0])
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:
# 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
if "method" not in message:
message["method"] = "INCIDENT_LIST"
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()}")
)
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
@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
method = loaded.pop("method", "UNKNOWN")
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"))]
)
fully_delivered = True
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.
self._oversized_dropped += 1
fully_delivered = False
logger.error(
"Dropping oversized NATS message: subject=%s"
" message_id=%s size=%d oversized_total=%d"
" error=%s preview=%r",
subject,
dedup_id,
len(payload),
self._oversized_dropped,
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 status report per logical message, not per fragment: a
# split message's fragments share the parent's reporter id, so
# reporting each would inflate the delivery-tracking cardinality.
# Skip the report when any fragment was dropped: the message did
# not fully reach NATS, so tracking it as sent overstates
# delivery. The drop is still counted and logged above.
if fully_delivered:
publisher.report(
{**loaded, "method": method},
_reporter_gen_nats,
stage="agent-nats-sending",
)
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()