104 lines
4.1 KiB
Python
104 lines
4.1 KiB
Python
import unittest
|
|
|
|
from fgai.llm import _json_object_from_text, compact_dashboard_analysis
|
|
|
|
|
|
class LlmTests(unittest.TestCase):
|
|
def test_json_object_from_text_accepts_markdown_wrapped_json(self):
|
|
payload = _json_object_from_text(
|
|
'```json\n{"profiles":[{"stream_id":"windows","entity_fields":["username"]}]}\n```'
|
|
)
|
|
|
|
self.assertEqual(payload["profiles"][0]["stream_id"], "windows")
|
|
|
|
def test_json_object_from_text_extracts_object_from_extra_text(self):
|
|
payload = _json_object_from_text(
|
|
'Here is the profile:\n{"profiles":[{"stream_id":"firewall"}]}\nDone.'
|
|
)
|
|
|
|
self.assertEqual(payload["profiles"][0]["stream_id"], "firewall")
|
|
|
|
def test_dashboard_compaction_removes_heavy_evidence_payloads(self):
|
|
status = {
|
|
"summary": {"total": 1000},
|
|
"capabilities": {
|
|
"graylog_mcp": {
|
|
"status": "connected",
|
|
"streams": [{"stream_id": f"stream-{index}", "events": index} for index in range(100)],
|
|
"aggregate_events": 1000,
|
|
"raw_events_fetched": 50,
|
|
}
|
|
},
|
|
"event_context": {
|
|
"source_profiles": [{"entity": f"10.0.0.{index}"} for index in range(20)],
|
|
"related_activity": [{"entity": f"10.0.0.{index}"} for index in range(40)],
|
|
},
|
|
"field_deviations": {
|
|
"alice": [
|
|
{
|
|
"stream_name": "Windows",
|
|
"detector": "new_relationship",
|
|
"field": "relationship:username->srcip",
|
|
"score": 90,
|
|
"reason": "new srcip value",
|
|
"sample_events": [{"message": "very large raw event", "graylog_query": "username:alice"}],
|
|
}
|
|
]
|
|
},
|
|
}
|
|
|
|
compact = compact_dashboard_analysis(status)
|
|
|
|
self.assertEqual(compact["capabilities"]["graylog_mcp"]["status"], "connected")
|
|
self.assertNotIn("streams", compact["capabilities"]["graylog_mcp"])
|
|
self.assertEqual(len(compact["event_context"]["source_profiles"]), 10)
|
|
self.assertEqual(len(compact["event_context"]["related_activity"]), 12)
|
|
self.assertEqual(compact["field_deviations"][0]["entity"], "alice")
|
|
self.assertNotIn("sample_events", compact["field_deviations"][0])
|
|
self.assertNotIn("graylog_query", str(compact))
|
|
|
|
def test_dashboard_compaction_strips_incident_and_correlation_samples(self):
|
|
status = {
|
|
"incidents": [
|
|
{
|
|
"entity": "10.0.0.5",
|
|
"score": 90,
|
|
"severity": "critical",
|
|
"correlated_streams": ["Firewall", "Proxy"],
|
|
"evidence": ["large burst of blocked traffic"],
|
|
"timeline": [{"message": "raw incident evidence", "graylog_query": "srcip:10.0.0.5"}],
|
|
}
|
|
],
|
|
"cross_source_correlations": [
|
|
{
|
|
"entity": "10.0.0.5",
|
|
"streams": ["Firewall", "Proxy"],
|
|
"events": 20,
|
|
"samples": [{"message": "raw correlation evidence", "graylog_query": "srcip:10.0.0.5"}],
|
|
}
|
|
],
|
|
"anomalies": [
|
|
{
|
|
"subject": "10.0.0.5",
|
|
"score": 90,
|
|
"severity": "critical",
|
|
"reasons": ["burst"],
|
|
"evidence": {"timed_events": 100, "sample_events": [{"message": "raw anomaly"}]},
|
|
}
|
|
],
|
|
}
|
|
|
|
compact = compact_dashboard_analysis(status)
|
|
compact_text = str(compact)
|
|
|
|
self.assertNotIn("timeline", compact_text)
|
|
self.assertNotIn("samples", compact_text)
|
|
self.assertNotIn("sample_events", compact_text)
|
|
self.assertNotIn("graylog_query", compact_text)
|
|
self.assertNotIn("raw incident evidence", compact_text)
|
|
self.assertEqual(compact["incidents"][0]["entity"], "10.0.0.5")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|