Update chronicle_auth to support alternative OAuth scopes. Add a library to

strasburg · copybara-github · commit 848f28e4ef20 · 2021-07-14T11:12:42.000-07:00
specify the Chronicle customer region and to hold region-specific
configuration.

PiperOrigin-RevId: 384741287
diff --git a/common/chronicle_auth.py b/common/chronicle_auth.py
@@ -27,7 +27,7 @@
 
 import argparse
 import pathlib
-from typing import Optional, Union
+from typing import Optional, Sequence, Union
 
 from google.auth.transport import requests
 from google.oauth2 import service_account
@@ -38,15 +38,18 @@
 
 
 def initialize_http_session(
-    credentials_file_path: Optional[Union[str, pathlib.Path]]
-) -> requests.AuthorizedSession:
+    credentials_file_path: Optional[Union[str, pathlib.Path]],
+    scopes: Optional[Sequence[str]] = None) -> requests.AuthorizedSession:
   """Initializes an authorized HTTP session, based on the given credentials.
 
   Args:
     credentials_file_path: Absolute or relative path to a JSON file containing
       the private OAuth 2.0 credentials of a Google Cloud Platform service
       account. Optional - the default is ".chronicle_credentials.json" in the
       user's home directory. Keep it secret, keep it safe.
+    scopes: A list of OAuth scopes (https://oauth.net/2/scope/) that are
+      associated with the end points to be accessed. The default is the
+      Chronicle API scope.
 
   Returns:
     HTTP session object to send authorized requests and receive responses.
@@ -56,10 +59,9 @@ def initialize_http_session(
       (https://docs.python.org/library/exceptions.html#os-exceptions).
     ValueError: Invalid file contents.
   """
-  if not credentials_file_path:
-    credentials_file_path = DEFAULT_CREDENTIALS_FILE
   credentials = service_account.Credentials.from_service_account_file(
-      str(credentials_file_path), scopes=AUTHORIZATION_SCOPES)
+      str(credentials_file_path or DEFAULT_CREDENTIALS_FILE),
+      scopes=scopes or AUTHORIZATION_SCOPES)
   return requests.AuthorizedSession(credentials)
 
 
diff --git a/common/chronicle_auth_test.py b/common/chronicle_auth_test.py
@@ -16,7 +16,10 @@
 
 import os
 import tempfile
+
 import unittest
+from unittest import mock
+from google.oauth2 import service_account
 
 from . import chronicle_auth
 
@@ -65,8 +68,27 @@ def setUp(self):
     os.write(fd, fake_json_credentials.strip() + fake_private_key + b'"\n}\n')
     os.close(fd)
 
-  def test_initialize_http_session_with_custom_json_credentials(self):
+  @mock.patch.object(service_account.Credentials, "from_service_account_file")
+  def test_initialize_http_session(self, mock_from_service_account_file):
+    chronicle_auth.initialize_http_session("")
+    mock_from_service_account_file.assert_called_once_with(
+        str(chronicle_auth.DEFAULT_CREDENTIALS_FILE),
+        scopes=chronicle_auth.AUTHORIZATION_SCOPES)
+
+  @mock.patch.object(service_account.Credentials, "from_service_account_file")
+  def test_initialize_http_session_with_custom_json_credentials(
+      self, mock_from_service_account_file):
     chronicle_auth.initialize_http_session(self.path)
+    mock_from_service_account_file.assert_called_once_with(
+        self.path, scopes=chronicle_auth.AUTHORIZATION_SCOPES)
+
+  @mock.patch.object(service_account.Credentials, "from_service_account_file")
+  def test_initialize_http_session_with_custom_creds_and_scopes(
+      self, mock_from_service_account_file):
+    scopes = ["https://www.googleapis.com/auth/malachite-ingestion"]
+    chronicle_auth.initialize_http_session(self.path, scopes=scopes)
+    mock_from_service_account_file.assert_called_once_with(
+        self.path, scopes=scopes)
 
   def tearDown(self):
     os.remove(self.path)
