taegis-sdk-python 1.6.0__tar.gz → 1.6.1__tar.gz

{taegis_sdk_python-1.6.0 → taegis_sdk_python-1.6.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: taegis-sdk-python
-Version: 1.6.0
+Version: 1.6.1
 Summary: Taegis Python SDK
 Project-URL: Homepage, https://github.com/secureworks/taegis-sdk-python
 Project-URL: Bug Reports, https://github.com/secureworks/taegis-sdk-python/issues
@@ -194,6 +194,7 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Requires-Python: >=3.9
 Requires-Dist: aiohttp
+Requires-Dist: backoff
 Requires-Dist: dataclasses-json
 Requires-Dist: filelock>=3.12.4
 Requires-Dist: gql<4,>=3.5.0

{taegis_sdk_python-1.6.0 → taegis_sdk_python-1.6.1}/pyproject.toml

@@ -34,6 +34,7 @@ dependencies = [
     "requests-toolbelt",
     "websockets",
     "jinja2",
+    "backoff",
 ]
 
 [project.optional-dependencies]

{taegis_sdk_python-1.6.0 → taegis_sdk_python-1.6.1}/taegis_sdk_python/_version.py

@@ -1,3 +1,3 @@
 """Version idenitier."""
 
-__version__ = "1.6.0"
+__version__ = "1.6.1"

taegis_sdk_python-1.6.1/taegis_sdk_python/middlewares/__init__.py

@@ -0,0 +1,6 @@
+"""Taegis SDK Middlewares."""
+
+from taegis_sdk_python.middlewares.logging import headers_logging_middleware
+from taegis_sdk_python.middlewares.retry import retry_middleware
+
+__all__ = ["headers_logging_middleware", "retry_middleware"]

taegis_sdk_python-1.6.1/taegis_sdk_python/middlewares/logging/__init__.py

@@ -0,0 +1,5 @@
+"""Taegis SDK Logging Middlewares."""
+
+from ._default import headers_logging_middleware
+
+__all__ = ["headers_logging_middleware"]

taegis_sdk_python-1.6.1/taegis_sdk_python/middlewares/logging/_default.py

@@ -0,0 +1,31 @@
+"""Taegis SDK Logging Middlewares."""
+
+import logging
+
+from aiohttp import ClientHandlerType, ClientRequest, ClientResponse
+
+from taegis_sdk_python.middlewares.utils import scrub_dict
+
+log = logging.getLogger(__name__)
+
+
+async def headers_logging_middleware(
+    req: ClientRequest, handler: ClientHandlerType
+) -> ClientResponse:
+    """Debug Log AIOHTTP Request and Response Headers.  Remove access tokens from headers and values.
+
+    Parameters
+    ----------
+    req : ClientRequest
+    handler : ClientHandlerType
+
+    Returns
+    -------
+    ClientResponse
+    """
+    headers = scrub_dict(req.headers)
+    log.debug(f"Request Headers: {headers}")
+    resp = await handler(req)
+    headers = scrub_dict(resp.headers)
+    log.debug(f"Response Headers: {headers}")
+    return resp

taegis_sdk_python-1.6.1/taegis_sdk_python/middlewares/retry/__init__.py

@@ -0,0 +1,5 @@
+"""Taegis SDK Retry Middlewares."""
+
+from ._default import retry_middleware
+
+__all__ = ["retry_middleware"]

taegis_sdk_python-1.6.1/taegis_sdk_python/middlewares/retry/_default.py

@@ -0,0 +1,80 @@
+"""Taegis SDK Retry Middlewares."""
+
+from http import HTTPStatus
+from typing import Optional
+
+import backoff
+from aiohttp import ClientHandlerType, ClientRequest, ClientResponse
+
+from taegis_sdk_python.config import get_config
+
+SECTION = "backoff.retry"
+
+
+def config_max_time() -> Optional[int]:
+    """Provide Max Time elapsed for retry.
+
+    Returns
+    -------
+    int
+        Seconds.
+    """
+    config = get_config()
+
+    return config.getint(SECTION, "max_time", fallback=10)
+
+
+def config_max_tries() -> Optional[int]:
+    """Provide Max Tries for retry.
+
+    Returns
+    -------
+    Optional[int]
+        Max tries.
+    """
+    config = get_config()
+
+    return config.getint(SECTION, "max_tries", fallback=None)
+
+
+@backoff.on_predicate(
+    backoff.runtime,
+    predicate=lambda r: r.status == HTTPStatus.TOO_MANY_REQUESTS,
+    value=lambda r: int(r.headers.get("Retry-After")),
+    max_time=config_max_time,
+    max_tries=config_max_tries,
+)
+@backoff.on_predicate(
+    backoff.expo,
+    lambda r: r.status
+    in (
+        HTTPStatus.REQUEST_TIMEOUT,
+        HTTPStatus.INTERNAL_SERVER_ERROR,
+        HTTPStatus.BAD_GATEWAY,
+        HTTPStatus.SERVICE_UNAVAILABLE,
+        HTTPStatus.GATEWAY_TIMEOUT,
+    ),
+    max_time=config_max_time,
+    max_tries=config_max_tries,
+)
+async def retry_middleware(
+    req: ClientRequest, handler: ClientHandlerType
+) -> ClientResponse:
+    """Retry HTTP requests middleware.
+
+    Retry on exponential backoff for REQUEST_TIMEOUT (408), INTERNAL_SERVER_ERROR (500),
+    BAD_GATEWAY (502), SERVICE_UNAVAILABLE (503), GATEWAY_TIMEOUT (504).
+
+    If TOO_MANY_REQUESTS (429) is encounted, the value in the "Retry-After" header
+    will be used as the backoff timer.
+
+    Parameters
+    ----------
+    req : ClientRequest
+    handler : ClientHandlerType
+
+    Returns
+    -------
+    ClientResponse
+    """
+    return await handler(req)

taegis_sdk_python-1.6.1/taegis_sdk_python/middlewares/utils.py

@@ -0,0 +1,39 @@
+"""Taegis SDK Python Middlewares Utilities."""
+
+from collections.abc import Mapping
+from typing import List
+
+from taegis_sdk_python.config import get_config
+
+
+def get_access_tokens() -> List[str]:
+    """Retrieve all cached access tokens."""
+    config = get_config()
+
+    tokens = []
+    for section in config.sections():
+        if token := config.get(section, "access_token", fallback=None):
+            tokens.append(token)
+
+    return tokens
+
+
+def scrub_dict(d: Mapping):
+    """Remove sensitive data from dictionary keys and values."""
+    access_tokens = get_access_tokens()
+
+    if "authorization" in d:
+        token = d["Authorization"].replace("Bearer ")
+        access_tokens.append(token)
+
+    scrubbed = d.copy()
+    for k, v in d.items():
+        if isinstance(k, str):
+            for token in access_tokens:
+                k = k.replace(token, "<redacted>")
+        if isinstance(v, str):
+            for token in access_tokens:
+                v = v.replace(token, "<redacted>")
+        scrubbed[k] = v
+
+    return scrubbed

{taegis_sdk_python-1.6.0 → taegis_sdk_python-1.6.1}/taegis_sdk_python/service_core.py

@@ -209,7 +209,10 @@ class ServiceCore:
     @property
     def sync_client(self) -> Client:
         """GraphQL Synchronous Transport with Client."""
-        client_session_args = {"trust_env": self.service.trust_env}
+        client_session_args = {
+            "trust_env": self.service.trust_env,
+            "middlewares": self.service.middlewares,
+        }
 
         transport = AIOHTTPTransport(
             f"{self.sync_url}{self.gateway}",
@@ -246,7 +249,10 @@ class ServiceCore:
         if self.service.tenant_id:
             subprotocols.append(f"x-tenant-context-{self.service.tenant_id}")
 
-        client_session_args = {"trust_env": self.service.trust_env}
+        client_session_args = {
+            "trust_env": self.service.trust_env,
+            "middlewares": self.service.middlewares,
+        }
 
         transport = AIOHTTPWebsocketsTransport(
             f"{self.wss_url}{self.gateway}",
@@ -483,11 +489,13 @@ class ServiceCore:
         """
         extra_args = {}
         if self.service.proxy:
-            extra_args = {
-                "proxy": self.service.proxy,
-                "proxy_auth": self.service.proxy_auth,
-                "proxy_headers": self.service.proxy_headers,
-            }
+            extra_args.update(
+                {
+                    "proxy": self.service.proxy,
+                    "proxy_auth": self.service.proxy_auth,
+                    "proxy_headers": self.service.proxy_headers,
+                }
+            )
 
         query = gql(query_string)
         return await self.sync_client.execute_async(

{taegis_sdk_python-1.6.0 → taegis_sdk_python-1.6.1}/taegis_sdk_python/services/__init__.py

@@ -84,7 +84,7 @@ log = logging.getLogger(__name__)
 class GraphQLService:
     """Taegis GraphQL Service manager."""
 
-    def __init__(
+    def __init__(  # pylint: disable=too-many-locals
         self,
         *,
         environment: Optional[str] = None,
@@ -101,6 +101,7 @@ class GraphQLService:
         execute_timeout: Optional[Union[int, float]] = 300,
         max_message_size: int = 0,
         use_universal_authentication: bool = False,
+        middlewares=None,
     ):  # pylint: disable=too-many-statements
         """
         GraphQLService
@@ -175,6 +176,7 @@ class GraphQLService:
 
         self._execute_timeout = execute_timeout
         self._max_message_size = max_message_size
+        self._middlewares = middlewares or ()
 
         self._access_points = None
         self._agent = None
@@ -424,6 +426,11 @@ class GraphQLService:
         """GraphQL Execute Timeout."""
         return self._context_manager.get("execute_timeout", self._execute_timeout)
 
+    @property
+    def middlewares(self):
+        """Transport Middleware handlers."""
+        return self._context_manager.get("middlewares", self._middlewares)
+
     @property
     def access_points(self):
         """Access Points Service Endpoint."""

{taegis_sdk_python-1.6.0 → taegis_sdk_python-1.6.1}/taegis_sdk_python/templates/_jinja2.py

@@ -24,13 +24,16 @@ def escape_value(value):
 # Filters
 
 
-def filter_or(value: List[str], field_name: str, operator: str = "="):
+def filter_or(value: List[str], field_name: str = None, operator: str = "="):
     """Format values for Taegis QL OR statement."""
     if not isinstance(value, list):
         raise ValueError("Input is not list.")
 
     return " OR ".join(
-        [f"{field_name} {operator} {escape_value(item)}" for item in value]
+        [
+            f"{field_name} {operator} {escape_value(item)}" if field_name else value
+            for item in value
+        ]
     )