Edit

[EDIT] FILE: 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()

CANCEL