diff --git a/src/fgai/graylog_source.py b/src/fgai/graylog_source.py index b1a2ef0..1da8fc8 100644 --- a/src/fgai/graylog_source.py +++ b/src/fgai/graylog_source.py @@ -17,9 +17,19 @@ DEFAULT_FIELD_MAP = { "action": ("action", "event_action", "disposition"), } +DEFAULT_FIELDS = ["timestamp", "source", "srcip", "dstip", "srcport", "dstport", "service", "action", "severity", "policyid", "subtype", "type", "sentbyte", "rcvdbyte", "hitcount"] + 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)) @@ -44,7 +54,13 @@ class GraylogStreamSource: def fetch(self) -> tuple[list[LogEvent], dict[str, object]]: status = self.client.probe() - arguments: dict[str, object] = {"query": self.query, "limit": 1000, "range_seconds": 300} + mapping_fields = [str(value) for value in self.mapping.values() if isinstance(value, str)] + arguments: dict[str, object] = { + "query": self.query, + "limit": 1000, + "range_seconds": 300, + "fields": list(dict.fromkeys([*DEFAULT_FIELDS, *mapping_fields])), + } if self.stream: arguments["streams"] = [self.stream] result = self.client.call_tool("search_messages", arguments) diff --git a/tests/test_graylog_source.py b/tests/test_graylog_source.py index a5b5056..e645090 100644 --- a/tests/test_graylog_source.py +++ b/tests/test_graylog_source.py @@ -12,7 +12,7 @@ class _Client: def call_tool(self, _name, arguments): self.arguments = arguments - return {"result": {"content": [{"type": "text", "text": '{"messages":[{"client":"10.0.0.1","server":"8.8.8.8","result":"deny"}]}' }]}} + return {"result": {"content": [{"type": "text", "text": '{"schema":[{"field":"client"},{"field":"server"},{"field":"result"}],"datarows":[["10.0.0.1","8.8.8.8","deny"],["10.0.0.2","8.8.8.8","accept"]]}' }]}} class GraylogSourceTests(unittest.TestCase): @@ -21,12 +21,13 @@ class GraylogSourceTests(unittest.TestCase): events, status = GraylogStreamSource( client, "vpn", field_mapping='{"srcip":"client","dstip":"server","action":"result"}' ).fetch() - self.assertEqual(status["events_fetched"], 1) + self.assertEqual(status["events_fetched"], 2) self.assertEqual(events[0].src_ip, "10.0.0.1") self.assertEqual(events[0].dst_ip, "8.8.8.8") self.assertEqual(events[0].action, "deny") self.assertEqual(client.arguments["streams"], ["vpn"]) self.assertEqual(client.arguments["range_seconds"], 300) + self.assertIn("client", client.arguments["fields"]) if __name__ == "__main__":