"""Send WordPress incidents to the correlation server."""
import hashlib
import itertools
import json
import logging
from datetime import datetime
from functools import partial
from time import monotonic
from types import MappingProxyType
from typing import Any, Mapping
from defence360agent.contracts.messages import SensorWordpressIncidentList
from defence360agent.contracts.plugins import MessageSink
from defence360agent.internals import delivery_ack
from defence360agent.model.wordpress_incident import (
get_unsent_wordpress_incidents,
settle_wordpress_incidents_reported,
)
logger = logging.getLogger(__name__)
class IncidentSender:
"""
Send WordPress incidents to the correlation server.
WordPress incidents are already in the Incident table (visible to UI).
This class sends them to correlation via Reportable messages, which are
handled by the SendToServer/SendToServerNATS/SendToServerFGW plugins.
Those plugins queue a message rather than deliver it, and a send round
can lose its batch or be force-cancelled mid-publish while the agent
shuts down. Each incident therefore keeps a count of the occurrences the
transport has not acknowledged, and every collection cycle sends whatever
is still outstanding.
"""
# rows one cycle takes on, so a backlog is paged instead of read whole;
# the byte budget, not this, is what bounds a single message
MAX_INCIDENTS_PER_CYCLE = 1000
# how long to wait for an acknowledgement before assuming the round was
# lost; also paces retries, since a batch nobody acknowledges is re-sent
# at most once per timeout
ACK_TIMEOUT = 300
def __init__(self) -> None:
# message_id -> ({incident id: occurrences it reported}, deadline)
self._inflight: dict[str, tuple[dict[int, int], float]] = {}
def _prepare_incident_for_correlation(
self, incident: dict
) -> dict[str, Any]:
"""
Prepare an incident for sending to the correlation server.
WordPress incidents use extra_info JSON field to store plugin-specific data.
Args:
incident: WordpressIncident dictionary (with extra_info populated)
Returns:
Dictionary formatted for correlation server
"""
logger.info("Preparing incident for correlation: %s", incident)
# JSONField automatically deserializes to dict, fallback to empty dict if None
extra: dict = incident.get("extra_info") or {}
# Convert timestamp to int and format date
timestamp_value: float = float(incident.get("timestamp") or 0)
timestamp = int(timestamp_value)
dt = (
datetime.fromtimestamp(timestamp_value).strftime("%Y-%m-%d")
if timestamp_value
else ""
)
return {
"timestamp": timestamp,
"dt": dt,
"plugin_id": incident.get("plugin"),
"rule": incident.get("rule") or "unknown",
"name": incident.get("name"),
"message": incident.get("description"),
"severity": incident.get("severity"),
"attackers_ip": incident.get("abuser") or "",
"domain": incident.get("domain") or "",
# the occurrences this message reports, not the row's running
# total: re-reporting occurrences correlation already counted
# inflates the heuristics that consume this field
"retries": incident.get("unsent_retries") or 1,
"uri": extra.get("request_uri") or "",
"user_agent": extra.get("http_user_agent") or "",
"http_method": extra.get("request_method") or "",
"user_logged_in": extra.get("user_logged_in") == "true"
if extra.get("user_logged_in")
else None,
"file_path": extra.get("site_path") or "",
"user": extra.get("username") or "",
"tag": self._build_tags(extra),
# Include WordPress-specific fields
"target": extra.get("target") or "",
"slug": extra.get("slug") or "",
"version": extra.get("version") or "",
"mode": extra.get("mode") or "",
"details": extra,
}
def _build_tags(self, extra: dict) -> list[str]:
tags = ["wordpress", "cve"]
if extra.get("cve"):
tags.append(extra["cve"])
if extra.get("target"):
tags.append(f"target_{extra['target']}")
if extra.get("mode"):
tags.append(f"mode_{extra['mode']}")
return tags
async def send_pending_incidents(self, sink: MessageSink | None) -> int:
"""Send every incident the server has not acknowledged.
Runs once per collection cycle, after the freshly collected
incidents were stored, so one pass covers both the new rows and the
ones whose earlier message never left the agent.
"""
if sink is None:
logger.warning("No sink provided, skipping incident sending")
return 0
self._expire_inflight()
incidents = get_unsent_wordpress_incidents(
limit=self.MAX_INCIDENTS_PER_CYCLE,
exclude_ids=self._inflight_ids(),
)
return await self.send_incidents(sink, incidents)
async def send_incidents(
self, sink: MessageSink | None, incidents: list[dict]
) -> int:
"""
Send WordPress incidents to the correlation server.
Since incidents are already in the WordpressIncident table (visible to UI),
we just need to send them to correlation.
Args:
incidents: List of incidents to send
Returns:
Number of incidents sent
"""
if sink is None:
logger.warning("No sink provided, skipping incident sending")
return 0
if len(incidents) == 0:
logger.debug("No incidents to send, skipping")
return 0
logger.info(
"Sending %d incidents to correlation server", len(incidents)
)
correlation_batch = [
self._prepare_incident_for_correlation(incident)
for incident in incidents
]
pending = iter(
[
(incident.get("id"), incident.get("unsent_retries") or 1)
for incident in incidents
]
)
for batch in SensorWordpressIncidentList.batched(correlation_batch):
await self._send_batch(
sink,
batch,
{
incident_id: occurrences
for incident_id, occurrences in itertools.islice(
pending, len(batch)
)
if incident_id is not None
},
)
return len(correlation_batch)
async def _send_batch(
self,
sink: MessageSink,
correlation_batch: list[dict],
reported: Mapping[int, int] = MappingProxyType({}),
):
"""Send a batch of incidents to correlation server.
Uses SensorIncidentList Reportable message which is sent to
correlation via the SendToServer/SendToServerNATS/SendToServerFGW
plugins.
"""
logger.info(
"Sending batch of %d incidents to correlation server",
len(correlation_batch),
)
logger.info(
"Correlation batch json: %s",
json.dumps(correlation_batch, indent=2),
)
message = SensorWordpressIncidentList(correlation_batch)
# Pin the id the send path would otherwise generate itself, so its
# acknowledgement can be tied back to these rows. Deriving it from the
# rows also makes a re-send carry the id of the message it repeats,
# which lets the transport's de-duplication window collapse the two.
message_id = hashlib.sha1(
json.dumps(
[sorted(reported.items()), correlation_batch], sort_keys=True
).encode()
).hexdigest()
message["message_id"] = message_id
self._watch(message_id, reported)
try:
await sink.process_message(message)
logger.info(
"Queued %d wordpress incident(s) for correlation server"
" (message %s)",
len(correlation_batch),
message_id,
)
except Exception as e:
# the message never made it into the queue, so stop waiting for
# an acknowledgement and let the next cycle send it again
self._unwatch(message_id)
logger.error(
"Failed to queue incident batch: %s",
e,
)
raise
def _watch(self, message_id: str, reported: Mapping[int, int]) -> None:
if not reported:
return
reported = dict(reported)
self._inflight[message_id] = (
reported,
monotonic() + self.ACK_TIMEOUT,
)
delivery_ack.registry.watch(
message_id,
partial(self._on_delivered, message_id, reported),
)
def _unwatch(self, message_id: str) -> None:
self._inflight.pop(message_id, None)
delivery_ack.registry.unwatch(message_id)
def _on_delivered(self, message_id: str, reported: dict[int, int]) -> None:
# written here rather than buffered for the next cycle because a
# send round usually lands during shutdown, and a count kept in
# memory until then would not survive the restart
self._inflight.pop(message_id, None)
settled = settle_wordpress_incidents_reported(reported)
logger.info(
"Discounted %d wordpress incident(s) delivered to correlation",
settled,
)
def _expire_inflight(self) -> None:
"""Give up on batches the transport never acknowledged, so their
incidents become eligible to be sent again."""
now = monotonic()
expired = [
message_id
for message_id, (_, deadline) in self._inflight.items()
if deadline <= now
]
for message_id in expired:
reported, _ = self._inflight.pop(message_id)
delivery_ack.registry.unwatch(message_id)
logger.warning(
"No delivery confirmation for %d wordpress incident(s)"
" (message %s) in %ds, sending them again",
len(reported),
message_id,
self.ACK_TIMEOUT,
)
def _inflight_ids(self) -> set[int]:
return {
incident_id
for reported, _ in self._inflight.values()
for incident_id in reported
}
|