add mcp earch

src/fgai/graylog_source.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import json
+from collections.abc import Iterable
+
+from .graylog_mcp import GraylogMcpClient
+from .models import LogEvent
+
+
+DEFAULT_FIELD_MAP = {
+    "srcip": ("srcip", "src_ip", "source_ip", "client_ip"),
+    "dstip": ("dstip", "dst_ip", "destination_ip", "server_ip"),
+    "dstport": ("dstport", "dst_port", "destination_port"),
+    "srcport": ("srcport", "src_port", "source_port"),
+    "eventtime": ("eventtime", "timestamp", "time"),
+    "severity": ("severity", "level"),
+    "action": ("action", "event_action", "disposition"),
+}
+
+
+def _records(value: object) -> Iterable[dict[str, object]]:
+    if isinstance(value, dict):
+        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 = "") -> None:
+        self.client = client
+        self.stream = stream
+        self.query = query or "*"
+        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) -> tuple[list[LogEvent], dict[str, object]]:
+        status = self.client.probe()
+        result = self.client.call_tool("search_messages", {"query": self.query, "stream": self.stream, "limit": 1000})
+        content = result.get("result", {}).get("content", []) if isinstance(result.get("result"), dict) else []
+        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 = [self._event(record) for record in records]
+        status.update({"source": "graylog_mcp", "events_fetched": len(events)})
+        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}
+        for canonical, candidates in DEFAULT_FIELD_MAP.items():
+            mapped = self.mapping.get(canonical)
+            candidates = (str(mapped),) if mapped else candidates
+            for candidate in candidates:
+                if candidate.lower() in fields:
+                    fields[canonical] = fields[candidate.lower()]
+                    break
+        return LogEvent(raw=json.dumps(record, sort_keys=True), fields=fields)