Add rate limiting and baseurl

pyproject.toml

@@ -7,7 +7,7 @@ name = "secops"
 version = "0.2.0"
 description = "Python SDK for wrapping the Google SecOps API for common use cases"
 readme = "README.md"
-requires-python = ">=3.7"
+requires-python = ">=3.7,<4.0"
 license = "Apache-2.0"
 authors = [
     { name = "Google SecOps Team", email = "[email protected]" }
@@ -28,6 +28,8 @@ dependencies = [
     "google-auth>=2.0.0",
     "google-auth-httplib2>=0.1.0",
     "google-api-python-client>=2.0.0",
+    "requests-ratelimiter (>=0.7.0,<0.8.0)",
+    "requests-toolbelt (>=1.0.0,<2.0.0)",
 ]
 
 [project.urls]

src/secops/auth.py

@@ -18,6 +18,9 @@
 from google.oauth2 import service_account
 import google.auth
 import google.auth.transport.requests
+from requests_ratelimiter import LimiterSession
+from requests_toolbelt.sessions import BaseUrlSession
+
 from secops.exceptions import AuthenticationError
 
 # Define default scopes needed for Chronicle API
@@ -33,7 +36,8 @@ def __init__(
         credentials: Optional[Credentials] = None,
         service_account_path: Optional[str] = None,
         service_account_info: Optional[Dict[str, Any]] = None,
-        scopes: Optional[List[str]] = None
+        scopes: Optional[List[str]] = None,
+        base_url: Optional[str] = None,
     ):
         """Initialize authentication for SecOps.
 
@@ -49,6 +53,7 @@ def __init__(
             service_account_path,
             service_account_info
         )
+        self.base_url = base_url
         self._session = None
 
     def _get_credentials(
@@ -87,8 +92,11 @@ def session(self):
         Returns:
             Authorized session for API requests
         """
-        if self._session is None:
-            self._session = google.auth.transport.requests.AuthorizedSession(
-                self.credentials
-            )
-        return self._session 
+        self._session |= google.auth.transport.requests.AuthorizedSession(
+            self.credentials,
+            session=BaseUrlSession(
+                base_url=self.base_url,
+                session=LimiterSession(per_minute=60, limit_statuses=[429]),
+            ),
+        )
+        return self._session

src/secops/chronicle/client.py

@@ -205,6 +205,7 @@ def __init__(
                         "https://www.googleapis.com/auth/chronicle-backstory",
                     ] + (extra_scopes or []),
                     credentials=credentials,
+                    base_url=self.base_url
                 )
 
             self._session = auth.session