Add python example for making requests to entityRiskScores:query endpoint

v1alpha/entity_risk_scores/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#

v1alpha/entity_risk_scores/query.py

@@ -0,0 +1,149 @@
+#!/usr/bin/env python3
+
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""Executable and reusable sample for retrieving a list of entity risk scores."""
+
+import argparse
+import json
+from typing import Any, Mapping
+
+from google.auth.transport import requests
+
+from common import chronicle_auth
+from common import regions
+
+
+CHRONICLE_API_BASE_URL = "https://chronicle.googleapis.com"
+
+
+def list_entity_risk_scores(
+    http_session: requests.AuthorizedSession,
+    region: str,
+    project: str,
+    instance: str,
+    filters: str = "",
+    order_by: str = "",
+    page_size: int = 0,
+    page_token: str = "",
+) -> Mapping[str, Any]:
+  """Retrieves a list of entity risk scores.
+
+  Args:
+    http_session: Authorized session for HTTP requests.
+    region: region in which the target project is located.
+    project: GCP project id or number which the target instance belongs to.
+    instance: uuid of the instance whose feeds are being listed (with dashes).
+    filters: Filters to be applied. Please see API definition for usage.
+    order_by: Ordering of Entity Risk Scores.
+    page_size: The maximum number of Entity Risk Scores to return.
+    page_token: A page token, received from a previous call. Used to retrieve
+      subsequent pages.
+
+  Returns:
+    JSON Object containing entity_risk_scores as well as other information about
+    the scores.
+  Raises:
+    requests.exceptions.HTTPError: HTTP request resulted in an error
+      (response.status_code >= 400).
+  """
+  url_region = "us" if region == "govcloud-us" else region
+  url = f"{CHRONICLE_API_BASE_URL}/v1alpha/projects/{project}/locations/{url_region}/instances/{instance}/entityRiskScores:query"
+  params_list = [
+      ("filter", filters),
+      ("order_by", order_by),
+      ("page_size", page_size),
+      ("page_token", page_token),
+  ]
+  params = {k: v for k, v in params_list if v}
+
+  response = http_session.request("GET", url, params=params)
+  if response.status_code >= 400:
+    print(response.text)
+  response.raise_for_status()
+  return response.json()
+
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  chronicle_auth.add_argument_credentials_file(parser)
+  regions.add_argument_region(parser)
+  parser.add_argument(
+      "-i",
+      "--instance",
+      type=str,
+      required=True,
+      help="instance to list feeds for",
+  )
+  parser.add_argument(
+      "-p",
+      "--project",
+      type=str,
+      required=True,
+      help="project to list feeds for",
+  )
+  parser.add_argument(
+      "-f",
+      "--filter",
+      type=str,
+      required=False,
+      help=(
+          "Filter to be applied over multiple entity risk score fields. Please"
+          " see API definition for usage"
+      ),
+  )
+  parser.add_argument(
+      "-o",
+      "--order_by",
+      type=str,
+      required=False,
+      help="Ordering of Entity Risk Scores",
+  )
+  parser.add_argument(
+      "-s",
+      "--page_size",
+      type=int,
+      required=False,
+      help="maximum number of results to return",
+  )
+  parser.add_argument(
+      "-t",
+      "--page_token",
+      type=str,
+      required=False,
+      help="page token from a previous call used for pagination",
+  )
+
+  args = parser.parse_args()
+  CHRONICLE_API_BASE_URL = regions.url(CHRONICLE_API_BASE_URL, args.region)
+  session = chronicle_auth.initialize_http_session(
+      args.credentials_file,
+      scopes=["https://www.googleapis.com/auth/cloud-platform"],
+  )
+  print(
+      json.dumps(
+          list_entity_risk_scores(
+              session,
+              args.region,
+              args.project,
+              args.instance,
+              args.filter,
+              args.order_by,
+              args.page_size,
+              args.page_token,
+          ),
+          indent=2,
+      )
+  )

v1alpha/entity_risk_scores/query_test.py

@@ -0,0 +1,94 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""Unit tests for the "query" module."""
+
+import unittest
+from unittest import mock
+
+from google.auth.transport import requests
+
+from . import query
+
+
+class QueryTest(unittest.TestCase):
+
+  # Class variables for ease.
+  arguments = {
+      "region": "us",
+      "project": "123456789101",
+      "instance": "test4bb9-878b-11e7-8455-10604b7cb5c1",
+  }
+
+  @mock.patch.object(requests, "AuthorizedSession", autospec=True)
+  @mock.patch.object(requests.requests, "Response", autospec=True)
+  def test_http_error(self, mock_response, mock_session):
+    mock_session.request.return_value = mock_response
+    type(mock_response).status_code = mock.PropertyMock(return_value=400)
+    mock_response.raise_for_status.side_effect = (
+        requests.requests.exceptions.HTTPError()
+    )
+
+    with self.assertRaises(requests.requests.exceptions.HTTPError):
+      query.list_entity_risk_scores(mock_session, **self.arguments)
+
+  @mock.patch.object(requests, "AuthorizedSession", autospec=True)
+  @mock.patch.object(requests.requests, "Response", autospec=True)
+  def test_happy_path_without_query_params(self, mock_response, mock_session):
+    mock_session.request.return_value = mock_response
+    type(mock_response).status_code = mock.PropertyMock(return_value=200)
+    expected_response = {
+        "entityRiskScores": [
+            {
+                "entity": {"metadata": {"entityType": "ASSET"}},
+                "riskWindow": {
+                    "startTime": "2023-08-16T17:45:59.870653094Z",
+                    "endTime": "2023-08-23T17:45:59.870654674Z",
+                },
+                "riskScore": 1000,
+                "riskDelta": {"previousRiskScore": 1000},
+                "detectionsCount": 3089795,
+                "firstDetectionTime": "2023-08-16T17:00:00.232292Z",
+                "lastDetectionTime": "2023-08-23T15:59:54.973741Z",
+                "entityIndicator": {"assetIpAddress": "127.0.0.1"},
+            },
+        ],
+        "entityCountDistributions": [
+            {
+                "dailyTimeBucket": {
+                    "startTime": "2023-08-16T07:00:00Z",
+                    "endTime": "2023-08-17T07:00:00Z",
+                },
+                "assetEntityCount": 9412,
+                "userEntityCount": 2481,
+            },
+        ],
+    }
+    mock_response.json.return_value = expected_response
+
+    response = query.list_entity_risk_scores(mock_session, **self.arguments)
+    self.assertIn("entityRiskScores", response)
+    self.assertIn("entity", response.get("entityRiskScores", [])[0])
+    self.assertIn("entityCountDistributions", response)
+    self.assertIn(
+        "assetEntityCount", response.get("entityCountDistributions", [])[0]
+    )
+    self.assertIn(
+        "userEntityCount", response.get("entityCountDistributions", [])[0]
+    )
+
+
+if __name__ == "__main__":
+  unittest.main()
+  googletest.main()