144 lines
6.3 KiB
Python
144 lines
6.3 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import time
|
|
from collections.abc import Iterable
|
|
|
|
from .graylog_mcp import GraylogMcpClient
|
|
from .models import LogEvent
|
|
from .normalization import DEFAULT_SEARCH_FIELDS, normalize_fields
|
|
|
|
|
|
WINDOWS_DISCOVERY_FIELDS = (
|
|
"winlog.channel",
|
|
"winlog.provider_name",
|
|
"winlog.record_id",
|
|
"winlog.process.pid",
|
|
"winlog.process.thread.id",
|
|
"winlog.task",
|
|
"winlog.opcode",
|
|
"winlog.keywords",
|
|
"winlog.event_data.logontype",
|
|
"winlog.event_data.logon_type",
|
|
"winlog.event_data.authenticationpackagename",
|
|
"winlog.event_data.logonprocessname",
|
|
"winlog.event_data.failurereason",
|
|
"winlog.event_data.status",
|
|
"winlog.event_data.substatus",
|
|
"process.name",
|
|
"process.executable",
|
|
"process.pid",
|
|
"process.command_line",
|
|
"service.name",
|
|
)
|
|
|
|
DEFAULT_FIELDS = list(dict.fromkeys([*DEFAULT_SEARCH_FIELDS, *WINDOWS_DISCOVERY_FIELDS, "source", "eventid", "event_id", "winlog_event_id", "policyid", "policy_id", "policy", "rule", "rule_name", "sentbyte", "rcvdbyte", "bytes_in", "bytes_out", "hitcount", "elapsed"]))
|
|
|
|
|
|
def _records(value: object) -> Iterable[dict[str, object]]:
|
|
if isinstance(value, dict):
|
|
schema = value.get("schema")
|
|
datarows = value.get("datarows")
|
|
if isinstance(schema, list) and isinstance(datarows, list):
|
|
fields = [str(column.get("field", "")) for column in schema if isinstance(column, dict)]
|
|
for row in datarows:
|
|
if isinstance(row, list):
|
|
yield {field: row[index] for index, field in enumerate(fields) if field and index < len(row)}
|
|
return
|
|
for key in ("messages", "results", "events", "data"):
|
|
if isinstance(value.get(key), list):
|
|
yield from (item for item in value[key] if isinstance(item, dict))
|
|
return
|
|
if value:
|
|
yield value
|
|
elif isinstance(value, list):
|
|
yield from (item for item in value if isinstance(item, dict))
|
|
|
|
|
|
class GraylogStreamSource:
|
|
def __init__(self, client: GraylogMcpClient, stream: str, query: str = "*", field_mapping: str = "", stream_label: str = "", profile_fields: tuple[str, ...] = ()) -> None:
|
|
self.client = client
|
|
self.stream = stream
|
|
self.query = query or "*"
|
|
self.stream_label = stream_label or stream
|
|
self.profile_fields = tuple(str(field).lower() for field in profile_fields if field)
|
|
try:
|
|
self.mapping = json.loads(field_mapping) if field_mapping else {}
|
|
except json.JSONDecodeError as exc:
|
|
raise RuntimeError("invalid_graylog_field_mapping") from exc
|
|
if not isinstance(self.mapping, dict):
|
|
raise RuntimeError("invalid_graylog_field_mapping")
|
|
|
|
def fetch(self, *, max_events: int = 5_000, range_seconds: int = 300, probe_status: dict[str, object] | None = None, deadline_monotonic: float | None = None) -> tuple[list[LogEvent], dict[str, object]]:
|
|
if probe_status is not None:
|
|
status = dict(probe_status)
|
|
else:
|
|
try:
|
|
status = self.client.probe()
|
|
except RuntimeError as exc:
|
|
return [], {
|
|
"status": "error",
|
|
"source": "graylog_mcp",
|
|
"events_fetched": 0,
|
|
"pages": 0,
|
|
"partial": True,
|
|
"error": f"probe_error: {exc}",
|
|
"truncated": False,
|
|
"latest_event_time": "",
|
|
}
|
|
mapping_fields = [str(value) for value in self.mapping.values() if isinstance(value, str)]
|
|
arguments: dict[str, object] = {
|
|
"query": self.query,
|
|
"range_seconds": max(1, int(range_seconds)),
|
|
"fields": list(dict.fromkeys([*DEFAULT_FIELDS, *mapping_fields, *self.profile_fields])),
|
|
}
|
|
if self.stream:
|
|
arguments["streams"] = [self.stream]
|
|
events: list[LogEvent] = []
|
|
page_size = min(5_000, max(1, int(max_events)))
|
|
pages = 0
|
|
partial_error = ""
|
|
while len(events) < max_events:
|
|
if deadline_monotonic is not None and time.monotonic() >= deadline_monotonic:
|
|
partial_error = "poll_budget_exceeded"
|
|
break
|
|
request_limit = min(page_size, max_events - len(events))
|
|
try:
|
|
result = self.client.call_tool("search_messages", {**arguments, "limit": request_limit, "offset": len(events)})
|
|
except RuntimeError as exc:
|
|
partial_error = str(exc)
|
|
break
|
|
content = result.get("result", {}).get("content", []) if isinstance(result.get("result"), dict) else []
|
|
if isinstance(result.get("result"), dict) and result["result"].get("isError"):
|
|
detail = next((str(item.get("text")) for item in content if isinstance(item, dict) and item.get("type") == "text"), "Graylog search failed")
|
|
partial_error = f"graylog_search_error: {detail}"
|
|
break
|
|
records: list[dict[str, object]] = []
|
|
for item in content if isinstance(content, list) else []:
|
|
if isinstance(item, dict) and item.get("type") == "text":
|
|
try:
|
|
records.extend(_records(json.loads(str(item.get("text", "")))))
|
|
except json.JSONDecodeError:
|
|
continue
|
|
events.extend(self._event(record) for record in records)
|
|
pages += 1
|
|
if len(records) < request_limit:
|
|
break
|
|
latest = max((event.fields.get("eventtime", "") for event in events), default="")
|
|
status.update({
|
|
"source": "graylog_mcp",
|
|
"events_fetched": len(events),
|
|
"pages": pages,
|
|
"partial": bool(partial_error),
|
|
"error": partial_error,
|
|
"truncated": len(events) >= max_events,
|
|
"latest_event_time": latest,
|
|
})
|
|
return events, status
|
|
|
|
def _event(self, record: dict[str, object]) -> LogEvent:
|
|
fields = {str(key).lower(): str(value) for key, value in record.items() if value is not None}
|
|
fields["fgai_stream"] = self.stream_label
|
|
fields["fgai_stream_id"] = self.stream
|
|
return LogEvent(raw=json.dumps(record, sort_keys=True), fields=normalize_fields(fields, self.mapping))
|