From d414d1a0f4205e371eaf5a3fbd37cf928e2ce6a7 Mon Sep 17 00:00:00 2001 From: Chronicle Team Date: Wed, 23 Aug 2023 12:34:44 -0700 Subject: [PATCH] Add python example for making requests to `entityRiskScores:query` endpoint PiperOrigin-RevId: 559509949 --- v1alpha/entity_risk_scores/__init__.py | 14 +++ v1alpha/entity_risk_scores/query.py | 149 +++++++++++++++++++++++ v1alpha/entity_risk_scores/query_test.py | 94 ++++++++++++++ 3 files changed, 257 insertions(+) create mode 100644 v1alpha/entity_risk_scores/__init__.py create mode 100644 v1alpha/entity_risk_scores/query.py create mode 100644 v1alpha/entity_risk_scores/query_test.py diff --git a/v1alpha/entity_risk_scores/__init__.py b/v1alpha/entity_risk_scores/__init__.py new file mode 100644 index 0000000..af1d471 --- /dev/null +++ b/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. +# diff --git a/v1alpha/entity_risk_scores/query.py b/v1alpha/entity_risk_scores/query.py new file mode 100644 index 0000000..8a83cff --- /dev/null +++ b/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, + ) + ) diff --git a/v1alpha/entity_risk_scores/query_test.py b/v1alpha/entity_risk_scores/query_test.py new file mode 100644 index 0000000..85280ce --- /dev/null +++ b/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()