chore: minor refactoring and improvements

Author: mihirvala08

src/secops/chronicle/client.py

@@ -1626,21 +1626,26 @@ def list_feeds(
         page_size: int = 100,
         page_token: str = None,
         api_version: APIVersion | None = None,
-    ) -> list[dict[str, Any]]:
+        as_list: bool = True,
+    ) -> dict[str, Any] | list[dict[str, Any]]:
         """List feeds.
 
         Args:
             page_size: The maximum number of feeds to return
             page_token: A page token, received from a previous ListFeeds call
             api_version: (Optional) Preferred API version to use.
+            as_list: If True, return only the list of feeds.
+                If False, return dict with metadata and pagination tokens.
+                Defaults to True for backward compatibility.
 
         Returns:
-            List of feed dictionaries
+            If as_list is True: List of feed dictionaries.
+            If as_list is False: Dict with feeds list and pagination metadata.
 
         Raises:
             APIError: If the API request fails
         """
-        return _list_feeds(self, page_size, page_token, api_version)
+        return _list_feeds(self, page_size, page_token, api_version, as_list)
 
     def get_feed(
         self, feed_id: str, api_version: APIVersion | None = None
@@ -2085,7 +2090,8 @@ def list_rules(
         page_size: int | None = None,
         page_token: str | None = None,
         api_version: APIVersion | None = APIVersion.V1,
-    ) -> dict[str, Any]:
+        as_list: bool = False,
+    ) -> dict[str, Any] | list[Any]:
         """Gets a list of rules.
 
         Args:
@@ -2099,9 +2105,12 @@ def list_rules(
             page_size: Maximum number of rules to return per page.
             page_token: Token for the next page of results, if available.
             api_version: (Optional) Preferred API version to use.
+            as_list: If True, return only the list of rules.
+                If False, return dict with metadata and pagination tokens.
 
         Returns:
-            Dictionary containing information about rules
+            If as_list is True: List of rules.
+            If as_list is False: Dict with rules list and pagination metadata.
 
         Raises:
             APIError: If the API request fails
@@ -2112,6 +2121,7 @@ def list_rules(
             page_size=page_size,
             page_token=page_token,
             api_version=api_version,
+            as_list=as_list,
         )
 
     def update_rule(
@@ -4852,17 +4862,22 @@ def list_rule_deployments(
         page_token: str | None = None,
         filter_query: str | None = None,
         api_version: APIVersion | None = APIVersion.V1,
-    ) -> dict[str, Any]:
+        as_list: bool = False,
+    ) -> dict[str, Any] | list[Any]:
         """List rule deployments for the instance.
 
         Args:
             page_size: Maximum number of deployments to return per page
             page_token: Token for the next page of results, if available
             filter_query: Optional filter query to restrict results
             api_version: (Optional) Preferred API version to use.
+            as_list: If True, return only the list of rule deployments.
+                If False, return dict with metadata and pagination tokens.
 
         Returns:
-            Dictionary containing rule deployments and pagination info
+            If as_list is True: List of rule deployments.
+            If as_list is False: Dict with ruleDeployments list and
+                pagination metadata.
 
         Raises:
             APIError: If the API request fails
@@ -4873,6 +4888,7 @@ def list_rule_deployments(
             page_token=page_token,
             filter_query=filter_query,
             api_version=api_version,
+            as_list=as_list,
         )
 
     def set_rule_alerting(

src/secops/chronicle/rule.py

@@ -319,7 +319,7 @@ def search_rules(
     except re.error as e:
         raise SecOpsError(f"Invalid regular expression: {query}") from e
 
-    rules = list_rules(client, api_version=api_version)
+    rules = list_rules(client, api_version=api_version, as_list=False)
     results = {"rules": []}
     for rule in rules["rules"]:
         rule_text = rule.get("text", "")

src/secops/cli/commands/feed.py

@@ -16,6 +16,7 @@
 
 import sys
 
+from secops.cli.utils.common_args import add_as_list_arg, add_pagination_args
 from secops.cli.utils.formatters import output_formatter
 
 
@@ -29,6 +30,8 @@ def setup_feed_command(subparsers):
 
     # List feeds command
     list_parser = feed_subparsers.add_parser("list", help="List feeds")
+    add_as_list_arg(list_parser)
+    add_pagination_args(list_parser)
     list_parser.set_defaults(func=handle_feed_list_command)
 
     # Get feed command
@@ -87,7 +90,11 @@ def setup_feed_command(subparsers):
 def handle_feed_list_command(args, chronicle):
     """Handle feed list command."""
     try:
-        result = chronicle.list_feeds()
+        result = chronicle.list_feeds(
+            page_size=args.page_size,
+            page_token=args.page_token,
+            as_list=args.as_list,
+        )
         output_formatter(result, args.output)
     except Exception as e:  # pylint: disable=broad-exception-caught
         print(f"Error: {e}", file=sys.stderr)

src/secops/cli/commands/rule.py

@@ -14,10 +14,12 @@
 #
 """Google SecOps CLI rule commands"""
 
+from ast import arg
 import json
 import sys
 
 from secops.cli.utils.common_args import (
+    add_as_list_arg,
     add_pagination_args,
     add_time_range_args,
 )
@@ -37,6 +39,7 @@ def setup_rule_command(subparsers):
     # List rules command
     list_parser = rule_subparsers.add_parser("list", help="List rules")
     add_pagination_args(list_parser)
+    add_as_list_arg(list_parser)
     list_parser.add_argument(
         "--view",
         type=str,
@@ -152,6 +155,7 @@ def setup_rule_command(subparsers):
         "list-deployments", help="List rule deployments"
     )
     add_pagination_args(list_dep_parser)
+    add_as_list_arg(list_dep_parser)
     list_dep_parser.add_argument(
         "--filter",
         dest="filter_query",
@@ -232,7 +236,10 @@ def handle_rule_list_command(args, chronicle):
     """Handle rule list command."""
     try:
         result = chronicle.list_rules(
-            view=args.view, page_size=args.page_size, page_token=args.page_token
+            view=args.view,
+            page_size=args.page_size,
+            page_token=args.page_token,
+            as_list=args.as_list,
         )
         output_formatter(result, args.output)
     except Exception as e:  # pylint: disable=broad-exception-caught
@@ -397,6 +404,7 @@ def handle_rule_list_deployments_command(args, chronicle):
             filter_query=(
                 args.filter_query if hasattr(args, "filter_query") else None
             ),
+            as_list=args.as_list,
         )
         output_formatter(result, args.output)
     except Exception as e:  # pylint: disable=broad-exception-caught

tests/chronicle/test_case.py

@@ -808,3 +808,24 @@ def test_patch_case_json_parse_error(chronicle_client):
     ):
         with pytest.raises(APIError):
             patch_case(chronicle_client, "12345", {"status": "CLOSED"})
+
+
+def test_list_cases_as_list(chronicle_client):
+    """Test list_cases function returning a list of cases."""
+    mock_response = Mock()
+    mock_response.status_code = 200
+    mock_response.json.return_value = {
+        "cases": [
+            {"id": "case-1", "displayName": "First Case"},
+            {"id": "case-2", "displayName": "Second Case"},
+        ]
+    }
+    with patch.object(
+        chronicle_client.session, "request", return_value=mock_response
+    ):
+        result = list_cases(chronicle_client, as_list=True)
+
+        assert isinstance(result, list)
+        assert len(result) == 2
+        assert result[0]["id"] == "case-1"
+        assert result[1]["id"] == "case-2"

tests/chronicle/test_client.py

@@ -13,12 +13,14 @@
 # limitations under the License.
 #
 """Tests for Chronicle API client."""
-from datetime import datetime, timezone, timedelta
-import pytest
+from datetime import datetime, timedelta, timezone
 from unittest.mock import Mock, patch
+
+import pytest
+
 from secops.chronicle.client import ChronicleClient
-from secops.chronicle.models import CaseList
-from secops.exceptions import APIError, SecOpsError
+from secops.chronicle.models import APIVersion, CaseList
+from secops.exceptions import APIError
 
 
 @pytest.fixture
@@ -723,3 +725,20 @@ def test_fix_json_formatting(chronicle_client):
     json_without_trailing_commas = '{"a": [1, 2], "b": {"c": 3, "d": 4}}'
     fixed = chronicle_client._fix_json_formatting(json_without_trailing_commas)
     assert fixed == json_without_trailing_commas
+
+@patch("secops.chronicle.client._list_rules")
+def test_client_list_rules_as_list(mock_list_rules, chronicle_client):
+    """Test that ChronicleClient.list_rules passes the as_list parameter."""
+    mock_list_rules.return_value = [{"name": "rule1"}]
+    
+    result = chronicle_client.list_rules(as_list=True)
+    
+    mock_list_rules.assert_called_once_with(
+        chronicle_client,
+        view="FULL",
+        page_size=None,
+        page_token=None,
+        api_version=APIVersion.V1,
+        as_list=True
+    )
+    assert isinstance(result, list)

tests/chronicle/test_rule.py

@@ -772,3 +772,39 @@ def test_list_rule_deployments_with_filter(chronicle_client, mock_response):
             ]
         }
         assert len(result["ruleDeployments"]) == 1
+
+
+def test_list_rules_as_list(chronicle_client):
+    """Test list_rules function with as_list=True."""
+    mock_response = Mock()
+    mock_response.status_code = 200
+    mock_response.json.return_value = {
+        "rules": [{"name": "rule1"}, {"name": "rule2"}]
+    }
+
+    with patch.object(
+        chronicle_client.session, "request", return_value=mock_response
+    ) as mock_get:
+        # Act
+        result = list_rules(chronicle_client, as_list=True)
+
+        # Assert
+        mock_get.assert_called_once()
+        assert isinstance(result, list)
+        assert len(result) == 2
+        assert result[0]["name"] == "rule1"
+        assert result[1]["name"] == "rule2"
+
+
+def test_list_rules_empty_as_list(chronicle_client):
+    """Test list_rules function with as_list=True when no rules exist."""
+    mock_response = Mock()
+    mock_response.status_code = 200
+    mock_response.json.return_value = {}
+
+    with patch.object(
+        chronicle_client.session, "request", return_value=mock_response
+    ):
+        result = list_rules(chronicle_client, as_list=True)
+        assert isinstance(result, list)
+        assert len(result) == 0

tests/chronicle/test_rule_detection.py

@@ -194,3 +194,20 @@ def test_list_errors_api_error(chronicle_client, response_mock):
 
     with pytest.raises(APIError, match="Failed to list rule errors"):
         rule_detection.list_errors(chronicle_client, rule_id="ru_missing")
+
+
+def test_list_detections_as_list(chronicle_client, response_mock):
+    """Test list_detections returning a list when as_list=True."""
+    chronicle_client.session.request.return_value = response_mock
+
+    result = rule_detection.list_detections(
+        chronicle_client,
+        rule_id="ru_12345678-1234-1234-1234-1234567890ab",
+        page_size=1,
+        as_list=True,
+    )
+
+    chronicle_client.session.request.assert_called_once()
+    assert isinstance(result, list)
+    assert len(result) == 1
+    assert result[0]["id"] == "de_12345678-1234-1234-1234-1234567890ab"

tests/chronicle/test_watchlist.py

@@ -455,3 +455,24 @@ def test_update_watchlist_error(chronicle_client):
             )
 
         assert "Failed to update watchlist" in str(exc_info.value)
+
+
+def test_list_watchlists_as_list(chronicle_client):
+    """Test list_watchlists function returning a list."""
+    mock_response = Mock()
+    mock_response.status_code = 200
+    mock_response.json.return_value = {
+        "watchlists": [
+            {"name": "watchlist-1", "displayName": "VIPs"},
+            {"name": "watchlist-2", "displayName": "Suspicious IPs"},
+        ]
+    }
+    with patch.object(
+        chronicle_client.session, "request", return_value=mock_response
+    ):
+        result = list_watchlists(chronicle_client, as_list=True)
+
+        assert isinstance(result, list)
+        assert len(result) == 2
+        assert result[0]["name"] == "watchlist-1"
+        assert result[1]["displayName"] == "Suspicious IPs"