From 2e9badf2a643134f50ece25439b219a373169d6c Mon Sep 17 00:00:00 2001 From: Dan Dye Date: Fri, 6 Nov 2026 20:34:52 +0000 Subject: [PATCH 1/2] fix(chronicle): preserve stats results in search_udm and nl_search Attach normalized stats results via process_stats_results along with too_many_results when :udmSearch returns a stats payload on Chronicle SIEM YARA-L 2.0 statistical queries. Plain event filter queries without a stats payload continue returning only events, total_events, and more_data_available. Fixes #227 --- src/secops/chronicle/client.py | 9 ++- src/secops/chronicle/nl_search.py | 3 +- src/secops/chronicle/search.py | 16 +++++- tests/chronicle/test_nl_search.py | 58 +++++++++++++++++++ tests/chronicle/test_search.py | 93 +++++++++++++++++++++++++++++++ 5 files changed, 172 insertions(+), 7 deletions(-) diff --git a/src/secops/chronicle/client.py b/src/secops/chronicle/client.py index 45f704c0..d61c2fd5 100644 --- a/src/secops/chronicle/client.py +++ b/src/secops/chronicle/client.py @@ -974,8 +974,9 @@ def search_udm( Returns: If as_list is True: List of Events. - If as_list is False: Dict with event list, total number of event and - flag to check if more data is available. + If as_list is False: Dict with events list, total_events, + more_data_available, and an optional stats dict when the + query returns statistical aggregation results. Raises: APIError: If the API request fails @@ -3360,7 +3361,9 @@ def nl_search( max_attempts: Maximum number of polling attempts Returns: - Dict containing the search results with events + Dict containing the search results with events, and an + optional stats dict when the translated query returns + aggregations Raises: APIError: If the API request fails diff --git a/src/secops/chronicle/nl_search.py b/src/secops/chronicle/nl_search.py index 66e70647..264d66a3 100644 --- a/src/secops/chronicle/nl_search.py +++ b/src/secops/chronicle/nl_search.py @@ -72,7 +72,8 @@ def nl_search( max_attempts: Maximum number of polling attempts Returns: - Dict containing the search results with events + Dict containing the search results with events, and an optional + stats dict when the translated query returns aggregations Raises: APIError: If the API request fails after retries diff --git a/src/secops/chronicle/search.py b/src/secops/chronicle/search.py index 32f0c7a6..54e670f0 100644 --- a/src/secops/chronicle/search.py +++ b/src/secops/chronicle/search.py @@ -18,6 +18,7 @@ from typing import Any, TYPE_CHECKING from secops.chronicle.models import APIVersion +from secops.chronicle.stats import process_stats_results from secops.chronicle.utils.request_utils import ( chronicle_request, ) @@ -55,8 +56,9 @@ def search_udm( Returns: If as_list is True: List of Events. - If as_list is False: Dict with event list, total number of event and - flag to check if more data is available. + If as_list is False: Dict with events list, total_events, + more_data_available, and an optional stats dict when the + query returns statistical aggregation results. Raises: APIError: If the API request fails @@ -93,8 +95,16 @@ def search_udm( return result.get("events", []) events = result.get("events", []) - return { + response: dict[str, Any] = { "events": events, "total_events": len(events), "more_data_available": result.get("moreDataAvailable", False), } + if result.get("stats") is not None: + raw_stats = result["stats"] + stats_summary = process_stats_results(raw_stats) + stats_summary["too_many_results"] = raw_stats.get( + "tooManyResults", False + ) + response["stats"] = stats_summary + return response diff --git a/tests/chronicle/test_nl_search.py b/tests/chronicle/test_nl_search.py index 59c1b891..e065265a 100644 --- a/tests/chronicle/test_nl_search.py +++ b/tests/chronicle/test_nl_search.py @@ -244,3 +244,61 @@ def test_nl_search_max_retries_exceeded( # Verify sleep was called 10 times assert mock_sleep.call_count == 10 + + +def test_nl_search_preserves_stats_from_outcome_query(mock_client): + """Test nl_search returns stats when translation emits an outcome: query.""" + from secops.chronicle.search import search_udm + + mock_client.search_udm.side_effect = lambda **kwargs: search_udm( + mock_client, **kwargs + ) + + translate_resp = MagicMock() + translate_resp.status_code = 200 + translate_resp.json.return_value = { + "query": ( + 'metadata.event_type = "USER_LOGIN"\n' + "match:\n target.user.userid\n" + "outcome:\n $login_count = count(metadata.id)" + ) + } + + udm_search_resp = MagicMock() + udm_search_resp.status_code = 200 + udm_search_resp.json.return_value = { + "stats": { + "results": [ + { + "column": "target.user.userid", + "values": [{"value": {"stringVal": "alice"}}], + }, + { + "column": "login_count", + "values": [{"value": {"int64Val": "15"}}], + }, + ], + "tooManyResults": False, + } + } + + mock_client.session.request.side_effect = [translate_resp, udm_search_resp] + + start_time = datetime.now(timezone.utc) - timedelta(hours=24) + end_time = datetime.now(timezone.utc) + + result = nl_search( + mock_client, + "count user logins by user in the last 24 hours", + start_time, + end_time, + ) + + assert result["events"] == [] + assert result["total_events"] == 0 + assert result["stats"] == { + "columns": ["target.user.userid", "login_count"], + "rows": [{"target.user.userid": "alice", "login_count": 15}], + "total_rows": 1, + "too_many_results": False, + } diff --git a/tests/chronicle/test_search.py b/tests/chronicle/test_search.py index ef094b3d..a06899e6 100644 --- a/tests/chronicle/test_search.py +++ b/tests/chronicle/test_search.py @@ -77,6 +77,7 @@ def test_search_udm_defaults_when_keys_missing( self.assertEqual(result["events"], []) self.assertEqual(result["total_events"], 0) self.assertFalse(result["more_data_available"]) + self.assertNotIn("stats", result) @mock.patch("secops.chronicle.search.chronicle_request") @mock.patch("builtins.print") @@ -194,6 +195,98 @@ def test_search_udm_empty_events_list( self.assertEqual(result["events"], []) self.assertEqual(result["total_events"], 0) self.assertFalse(result["more_data_available"]) + self.assertNotIn("stats", result) + + @mock.patch("secops.chronicle.search.chronicle_request") + def test_search_udm_includes_processed_stats_when_present( + self, mock_chronicle_request: mock.MagicMock + ) -> None: + """Test that stats payloads from outcome: queries are preserved.""" + mock_chronicle_request.return_value = { + "stats": { + "results": [ + { + "column": "principal.ip", + "values": [ + {"value": {"stringVal": "10.0.0.1"}}, + {"value": {"stringVal": "10.0.0.2"}}, + ], + }, + { + "column": "login_count", + "values": [ + {"value": {"int64Val": "42"}}, + {"value": {"int64Val": "7"}}, + ], + }, + ], + "tooManyResults": True, + } + } + + result = search_udm( + client=self.client, + query=( + 'metadata.event_type = "USER_LOGIN"\n' + "match:\n principal.ip\n" + "outcome:\n $login_count = count(metadata.id)" + ), + start_time=self.start_time, + end_time=self.end_time, + ) + + self.assertEqual(result["events"], []) + self.assertEqual(result["total_events"], 0) + self.assertFalse(result["more_data_available"]) + self.assertIn("stats", result) + self.assertEqual( + result["stats"], + { + "columns": ["principal.ip", "login_count"], + "rows": [ + {"principal.ip": "10.0.0.1", "login_count": 42}, + {"principal.ip": "10.0.0.2", "login_count": 7}, + ], + "total_rows": 2, + "too_many_results": True, + }, + ) + + @mock.patch("secops.chronicle.search.chronicle_request") + def test_search_udm_stats_defaults_too_many_results_false( + self, mock_chronicle_request: mock.MagicMock + ) -> None: + """Test too_many_results defaults to False when omitted in stats.""" + mock_chronicle_request.return_value = { + "stats": { + "results": [ + { + "column": "event_count", + "values": [{"value": {"int64Val": "5"}}], + } + ] + } + } + + result = search_udm( + client=self.client, + query=( + 'metadata.event_type = "USER_LOGIN"\n' + "outcome:\n $event_count = count(metadata.id)" + ), + start_time=self.start_time, + end_time=self.end_time, + ) + + self.assertEqual( + result["stats"], + { + "columns": ["event_count"], + "rows": [{"event_count": 5}], + "total_rows": 1, + "too_many_results": False, + }, + ) if __name__ == "__main__": From a0e2b49a593d0d90d2128e457773141b8e94d807 Mon Sep 17 00:00:00 2001 From: Dan Dye Date: Fri, 6 Nov 2026 21:40:24 +0000 Subject: [PATCH 2/2] fix(chronicle): align search_udm stats schema with get_stats Match ChronicleClient.get_stats() 3-key stats dict, fold tooManyResults into more_data_available, guard non-dict stats payloads, and clarify search_udm and nl_search docstrings. --- README.md | 2 + src/secops/chronicle/client.py | 22 +++++--- src/secops/chronicle/nl_search.py | 5 +- src/secops/chronicle/search.py | 32 ++++++----- tests/chronicle/test_nl_search.py | 5 +- tests/chronicle/test_search.py | 92 ++++++++++++++++++++++++++----- 6 files changed, 118 insertions(+), 40 deletions(-) diff --git a/README.md b/README.md index 0770a732..07c097d9 100644 --- a/README.md +++ b/README.md @@ -1265,6 +1265,8 @@ The natural language search feature supports various query patterns: - "Show login failures in the last hour" - "Display connections to IP address 192.168.1.100" +When a natural language prompt translates into a statistical aggregation query, the returned dictionary includes a `stats` key with `columns`, `rows`, and `total_rows` matching `get_stats()`, while `events` is empty. + If the natural language cannot be translated to a valid UDM query, an `APIError` will be raised with a message indicating that no valid query could be generated. ### Entity Summary diff --git a/src/secops/chronicle/client.py b/src/secops/chronicle/client.py index d61c2fd5..f4ca1d95 100644 --- a/src/secops/chronicle/client.py +++ b/src/secops/chronicle/client.py @@ -969,14 +969,18 @@ def search_udm( max_attempts: Maximum number of polling attempts (default: 30) timeout: Timeout in seconds for each API request (default: 30) debug: Print debug information during execution - as_list: If True, return a list of events instead of a dict - with events list and nextPageToken. + as_list: If True, return only the list of events; statistical + results in stats are omitted when as_list is True. Returns: - If as_list is True: List of Events. - If as_list is False: Dict with events list, total_events, - more_data_available, and an optional stats dict when the - query returns statistical aggregation results. + If as_list is True: List of matched UDM events, or an empty list + if the query returned only statistical results. + If as_list is False: Dict with events, total_events, + more_data_available, and an optional stats dict containing + columns, rows, and total_rows matching get_stats() when the + query returns statistical results. For dedicated statistical + queries with custom max_values or a 120s default timeout, + use get_stats(). Raises: APIError: If the API request fails @@ -3361,9 +3365,9 @@ def nl_search( max_attempts: Maximum number of polling attempts Returns: - Dict containing the search results with events, and an - optional stats dict when the translated query returns - aggregations + Dict with events, total_events, more_data_available, and an + optional stats dict containing columns, rows, and total_rows + when the translated query returns statistical results. Raises: APIError: If the API request fails diff --git a/src/secops/chronicle/nl_search.py b/src/secops/chronicle/nl_search.py index 264d66a3..45f8cce9 100644 --- a/src/secops/chronicle/nl_search.py +++ b/src/secops/chronicle/nl_search.py @@ -72,8 +72,9 @@ def nl_search( max_attempts: Maximum number of polling attempts Returns: - Dict containing the search results with events, and an optional - stats dict when the translated query returns aggregations + Dict with events, total_events, more_data_available, and an + optional stats dict containing columns, rows, and total_rows + when the translated query returns statistical results. Raises: APIError: If the API request fails after retries diff --git a/src/secops/chronicle/search.py b/src/secops/chronicle/search.py index 54e670f0..941ccf31 100644 --- a/src/secops/chronicle/search.py +++ b/src/secops/chronicle/search.py @@ -52,13 +52,18 @@ def search_udm( for backwards compatibility) timeout: Timeout in seconds for each API request (default: 30) debug: Print debug information during execution - as_list: Whether to return results as a list or dictionary + as_list: If True, return only the list of events; statistical + results in stats are omitted when as_list is True. Returns: - If as_list is True: List of Events. - If as_list is False: Dict with events list, total_events, - more_data_available, and an optional stats dict when the - query returns statistical aggregation results. + If as_list is True: List of matched UDM events, or an empty list + if the query returned only statistical results. + If as_list is False: Dict with events, total_events, + more_data_available, and an optional stats dict containing + columns, rows, and total_rows matching get_stats() when the + query returns statistical results. For dedicated statistical + queries with custom max_values or a 120s default timeout, + use get_stats(). Raises: APIError: If the API request fails @@ -95,16 +100,17 @@ def search_udm( return result.get("events", []) events = result.get("events", []) + raw_stats = result.get("stats") + has_stats_dict = isinstance(raw_stats, dict) + more_data_available = bool( + result.get("moreDataAvailable", False) + or (has_stats_dict and raw_stats.get("tooManyResults", False)) + ) response: dict[str, Any] = { "events": events, "total_events": len(events), - "more_data_available": result.get("moreDataAvailable", False), + "more_data_available": more_data_available, } - if result.get("stats") is not None: - raw_stats = result["stats"] - stats_summary = process_stats_results(raw_stats) - stats_summary["too_many_results"] = raw_stats.get( - "tooManyResults", False - ) - response["stats"] = stats_summary + if has_stats_dict: + response["stats"] = process_stats_results(raw_stats) return response diff --git a/tests/chronicle/test_nl_search.py b/tests/chronicle/test_nl_search.py index e065265a..3dd0528b 100644 --- a/tests/chronicle/test_nl_search.py +++ b/tests/chronicle/test_nl_search.py @@ -18,6 +18,7 @@ from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock, patch from secops.chronicle.nl_search import translate_nl_to_udm, nl_search +from secops.chronicle.search import search_udm from secops.exceptions import APIError import unittest.mock @@ -248,8 +249,6 @@ def test_nl_search_max_retries_exceeded( def test_nl_search_preserves_stats_from_outcome_query(mock_client): """Test nl_search returns stats when translation emits an outcome: query.""" - from secops.chronicle.search import search_udm - mock_client.search_udm.side_effect = lambda **kwargs: search_udm( mock_client, **kwargs ) @@ -296,9 +295,9 @@ def test_nl_search_preserves_stats_from_outcome_query(mock_client): assert result["events"] == [] assert result["total_events"] == 0 + assert result["more_data_available"] is False assert result["stats"] == { "columns": ["target.user.userid", "login_count"], "rows": [{"target.user.userid": "alice", "login_count": 15}], "total_rows": 1, - "too_many_results": False, } diff --git a/tests/chronicle/test_search.py b/tests/chronicle/test_search.py index a06899e6..e50ba248 100644 --- a/tests/chronicle/test_search.py +++ b/tests/chronicle/test_search.py @@ -237,7 +237,7 @@ def test_search_udm_includes_processed_stats_when_present( self.assertEqual(result["events"], []) self.assertEqual(result["total_events"], 0) - self.assertFalse(result["more_data_available"]) + self.assertTrue(result["more_data_available"]) self.assertIn("stats", result) self.assertEqual( result["stats"], @@ -248,15 +248,88 @@ def test_search_udm_includes_processed_stats_when_present( {"principal.ip": "10.0.0.2", "login_count": 7}, ], "total_rows": 2, - "too_many_results": True, }, ) @mock.patch("secops.chronicle.search.chronicle_request") - def test_search_udm_stats_defaults_too_many_results_false( + def test_search_udm_empty_stats_payload_preserved( + self, mock_chronicle_request: mock.MagicMock + ) -> None: + """Test zero-row outcome: stats payloads remain distinguishable.""" + mock_chronicle_request.return_value = {"stats": {}} + + empty_result = search_udm( + client=self.client, + query=( + 'metadata.event_type = "USER_LOGIN"\n' + "outcome:\n $event_count = count(metadata.id)" + ), + start_time=self.start_time, + end_time=self.end_time, + ) + + self.assertFalse(empty_result["more_data_available"]) + self.assertEqual( + empty_result["stats"], + {"columns": [], "rows": [], "total_rows": 0}, + ) + + mock_chronicle_request.return_value = { + "stats": { + "results": [ + {"column": "principal.ip"}, + {"column": "login_count"}, + ] + } + } + + header_only_result = search_udm( + client=self.client, + query=( + 'metadata.event_type = "USER_LOGIN"\n' + "match:\n principal.ip\n" + "outcome:\n $login_count = count(metadata.id)" + ), + start_time=self.start_time, + end_time=self.end_time, + ) + + self.assertFalse(header_only_result["more_data_available"]) + self.assertEqual( + header_only_result["stats"], + { + "columns": ["principal.ip", "login_count"], + "rows": [], + "total_rows": 0, + }, + ) + + @mock.patch("secops.chronicle.search.chronicle_request") + def test_search_udm_ignores_non_dict_stats_value( self, mock_chronicle_request: mock.MagicMock ) -> None: - """Test too_many_results defaults to False when omitted in stats.""" + """Test malformed non-dict stats payloads do not raise errors.""" + mock_chronicle_request.return_value = { + "events": [], + "stats": "unexpected", + } + + result = search_udm( + client=self.client, + query="test", + start_time=self.start_time, + end_time=self.end_time, + ) + + self.assertEqual(result["events"], []) + self.assertFalse(result["more_data_available"]) + self.assertNotIn("stats", result) + + @mock.patch("secops.chronicle.search.chronicle_request") + def test_search_udm_as_list_with_stats_returns_empty_events_list( + self, mock_chronicle_request: mock.MagicMock + ) -> None: + """Test as_list=True returns empty event list on stats responses.""" mock_chronicle_request.return_value = { "stats": { "results": [ @@ -276,17 +349,10 @@ def test_search_udm_stats_defaults_too_many_results_false( ), start_time=self.start_time, end_time=self.end_time, + as_list=True, ) - self.assertEqual( - result["stats"], - { - "columns": ["event_count"], - "rows": [{"event_count": 5}], - "total_rows": 1, - "too_many_results": False, - }, - ) + self.assertEqual(result, []) if __name__ == "__main__":