diff --git a/README.md b/README.md index 0770a73..07c097d 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 45f704c..f4ca1d9 100644 --- a/src/secops/chronicle/client.py +++ b/src/secops/chronicle/client.py @@ -969,13 +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 event list, total number of event and - flag to check if more data is available. + 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 @@ -3360,7 +3365,9 @@ def nl_search( max_attempts: Maximum number of polling attempts Returns: - Dict containing the search results with events + 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 66e7064..45f8cce 100644 --- a/src/secops/chronicle/nl_search.py +++ b/src/secops/chronicle/nl_search.py @@ -72,7 +72,9 @@ def nl_search( max_attempts: Maximum number of polling attempts Returns: - Dict containing the search results with events + 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 32f0c7a..941ccf3 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, ) @@ -51,12 +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 event list, total number of event and - flag to check if more data is available. + 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 @@ -93,8 +100,17 @@ def search_udm( return result.get("events", []) events = result.get("events", []) - return { + 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 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 59c1b89..3dd0528 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 @@ -244,3 +245,59 @@ 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.""" + 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["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, + } diff --git a/tests/chronicle/test_search.py b/tests/chronicle/test_search.py index ef094b3..e50ba24 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,164 @@ 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.assertTrue(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, + }, + ) + + @mock.patch("secops.chronicle.search.chronicle_request") + 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 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": [ + { + "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, + as_list=True, + ) + + self.assertEqual(result, []) if __name__ == "__main__":