taegis-sdk-python 1.2.4__tar.gz → 1.2.6__tar.gz

{taegis-sdk-python-1.2.4/taegis_sdk_python.egg-info → taegis_sdk_python-1.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: taegis-sdk-python
-Version: 1.2.4
+Version: 1.2.6
 Summary: Taegis Python SDK
 Home-page: https://github.com/secureworks/taegis-sdk-python
 Author: Secureworks
@@ -21,6 +21,7 @@ Requires-Dist: dataclasses-json
 Requires-Dist: pyjwt>=2.8.0
 Requires-Dist: requests-toolbelt
 Requires-Dist: websockets
+Requires-Dist: filelock>=3.12.4
 
 # Taegis SDK for Python
 

taegis-sdk-python-1.2.4/taegis_sdk_python.egg-info/requires.txt → taegis_sdk_python-1.2.6/requirements.txt

@@ -8,3 +8,4 @@ dataclasses-json
 pyjwt>=2.8.0
 requests-toolbelt
 websockets
+filelock>=3.12.4

taegis_sdk_python-1.2.6/taegis_sdk_python/_version.py

@@ -0,0 +1 @@
+__version__ = "v1.2.6"

{taegis-sdk-python-1.2.4 → taegis_sdk_python-1.2.6}/taegis_sdk_python/authentication.py

@@ -8,6 +8,7 @@ import os
 from getpass import getpass
 from time import time
 from typing import Any, Dict, Tuple, Union, Optional
+import threading
 
 from oauthlib.oauth2 import BackendApplicationClient
 from requests import HTTPError, post
@@ -26,6 +27,8 @@ except ImportError:  # pragma: no cover
 
 logger = logging.getLogger(__name__)
 
+LOCK = threading.RLock()
+
 
 def check_username(
     request_url: str, username: str
@@ -103,23 +106,31 @@ def get_token(environment: str, request_url: str) -> str:  # pragma: no cover
     access_token = get_cached_token(environment)
 
     if not access_token:
-        client_id, client_secret = get_oauth_from_env(environment)
-        if client_id and client_secret:
-            access_token = get_token_by_oauth(request_url, client_id, client_secret)
-        else:
-            username = input("Username: ")
-            response = check_username(request_url, username)
-
-            if response.get("login_type") == "username-password":
-                access_token = get_token_by_password_grant(request_url, username)
-            elif response.get("login_type") == "sso":
-                access_token = get_token_by_sso_device_code(request_url)
-            else:
-                raise InvalidAuthenticationMethod(
-                    message="No known authenticition method for user"
-                )
-
-        write_to_config(environment, "access_token", access_token)
+        with LOCK:
+            access_token = get_cached_token(environment)
+
+            if not access_token:
+                client_id, client_secret = get_oauth_from_env(environment)
+                if client_id and client_secret:
+                    access_token = get_token_by_oauth(
+                        request_url, client_id, client_secret
+                    )
+                else:
+                    username = input("Username: ")
+                    response = check_username(request_url, username)
+
+                    if response.get("login_type") == "username-password":
+                        access_token = get_token_by_password_grant(
+                            request_url, username
+                        )
+                    elif response.get("login_type") == "sso":
+                        access_token = get_token_by_sso_device_code(request_url)
+                    else:
+                        raise InvalidAuthenticationMethod(
+                            message="No known authentication method for user"
+                        )
+
+                write_to_config(environment, "access_token", access_token)
 
     return access_token
 
@@ -130,7 +141,7 @@ def get_cached_token(env: str) -> Union[str, None]:  # pragma: no cover
 
     # check for token and expiry in config
     token = str(config.get(env, "access_token", fallback=""))
-    if token and get_token_exp(token) >= int(time()):
+    if token and get_token_exp(token) >= int(time()) + 15:
         return token
 
     return None

{taegis-sdk-python-1.2.4 → taegis_sdk_python-1.2.6}/taegis_sdk_python/config.py

@@ -2,8 +2,13 @@
 
 Taegis SDK Configuration management.
 """
+
 from configparser import ConfigParser
 from pathlib import Path
+import threading
+from filelock import FileLock
+
+LOCK = threading.RLock()
 
 
 def get_config_file() -> Path:
@@ -44,9 +49,14 @@ def write_config(config: ConfigParser):
         Config object
     """
     config_fp = get_config_file()
+    lock_file = config_fp.with_suffix(".lock")
+    file_lock = FileLock(lock_file)
+
+    with file_lock:
+        with config_fp.open(mode="w") as f:  # pylint: disable=invalid-name
+            config.write(f)
 
-    with config_fp.open(mode="w") as f:  # pylint: disable=invalid-name
-        config.write(f)
+    lock_file.unlink(missing_ok=True)
 
 
 def write_to_config(section: str, key: str, value: str):

{taegis-sdk-python-1.2.4 → taegis_sdk_python-1.2.6}/taegis_sdk_python/service_core.py

@@ -8,8 +8,10 @@ from __future__ import annotations
 import logging
 from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+import threading
 
 from gql import Client, gql
+from gql.transport import Transport
 from gql.transport.requests import RequestsHTTPTransport
 from gql.transport.websockets import WebsocketsTransport
 from graphql import GraphQLError, GraphQLField, GraphQLSchema
@@ -25,57 +27,135 @@ log = logging.getLogger(__name__)
 
 
 class SchemaCache:
-    """GraphQL Schema Cache."""
+    """Taegis GraphQL Schema Cache."""
 
-    def __init__(self, schema: Optional[GraphQLSchema] = None, expires: int = 5):
-        """Initialize SchemaCache.
+    _lock = threading.RLock()
 
-        Parameters
-        ----------
-        schema : Optional[GraphQLSchema], optional
-            GraphQL schema object, by default None
-        expires : int, optional
-            Schema expires in minutes, by default 5
-        """
+    def __init__(self, expires: int, schema: Optional[GraphQLSchema] = None):
+        self._expires = expires
         self._schema = schema
         self._inserted_at = datetime.now()
-        self.expires = expires
+
+    def is_expired(self) -> bool:
+        """Returns if the schema is expired.
+
+        Returns
+        -------
+        bool
+            Is schema expired?
+        """
+        return self._inserted_at + timedelta(minutes=self._expires) <= datetime.now()
 
     @property
-    def schema(self) -> Union[GraphQLSchema, None]:
-        """Return GraphQLSchema if it exists and is not expired.
+    def schema(self) -> Optional[GraphQLSchema]:
+        """Get the schema.
 
         Returns
         -------
-        Union[GraphQLSchema, None]
-            GraphQL Schema
+        Optional[GraphQLSchema]
         """
-        if self._inserted_at + timedelta(minutes=self.expires) < datetime.now():
-            self.schema = None
+        log.debug(f"Schema is expired: {self.is_expired()}")
+
         return self._schema
 
-    @schema.setter
-    def schema(self, schema: Optional[GraphQLSchema]):
-        """Set schema and update inserted_at.
+
+class SchemaCacheMap:
+    """Taegis Schema Cache Map.
+
+    This is a singleton class that will cache the introspection schema for a given
+    GraphQL server.  The schema will be cached for a given amount of time and will
+    be retrieved from the transport if expired.
+    """
+
+    _instance = None
+    _lock = threading.RLock()
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super(SchemaCacheMap, cls).__new__(cls)
+                    cls._instance.__post_init__(*args, **kwargs)
+        return cls._instance
+
+    def __post_init__(self, expires: int = 5):
+        self._map = {}
+        self._expires = expires
+
+    def get(
+        self, transport: Transport, introspection_args: Dict[str, Any]
+    ) -> GraphQLSchema:
+        """Get the schema from the map.
+
+        Retrieves new schema if expired, thread locks to prevent another service
+        from retrieving the schema at the same time.
+
+        Parameters
+        ----------
+        transport : Transport
+            gql Transport
+        introspection_args : Dict[str, Any]
+            Introspection Query arguments
+
+        Returns
+        -------
+        GraphQLSchema
+        """
+        if transport.url not in self._map or self._map[transport.url].is_expired():
+            log.debug("Schema not in map or expired...")
+            with self._lock:
+                if (
+                    transport.url not in self._map
+                    or self._map[transport.url].is_expired()
+                ):
+                    self._map[transport.url] = SchemaCache(
+                        expires=self._expires,
+                        schema=self.retrieve_schema(transport, introspection_args),
+                    )
+
+        return self._map.get(transport.url).schema
+
+    def clear(self, transport: Transport):
+        """Clear the schema from the map.
 
         Parameters
         ----------
-        schema :
-            _description_
+        transport : Transport
+            gql Transport
         """
-        if self._schema is not schema:
-            self._schema = schema
-            self._inserted_at = datetime.now()
+        if transport.url in self._map:
+            with self._lock:
+                if transport.url in self._map:
+                    del self._map[transport.url]
+
+    @staticmethod
+    def retrieve_schema(
+        transport: Transport, introspection_args: Dict[str, Any]
+    ) -> GraphQLSchema:
+        """Retrieve the schema from a gql Client with the provided transport and
+        introspection query arguments.
 
-    def should_fetch_schema(self) -> bool:
-        """Return True if schema is should be fetched in transport.
+        Parameters
+        ----------
+        transport : Transport
+        introspection_args : Dict[str, Any]
 
         Returns
         -------
-        bool
-            True if schema is should be fetched in transport.
+        GraphQLSchema
         """
-        return not self.schema
+        log.debug(f"Retrieving schema for {transport.url}...")
+
+        client = Client(
+            transport=transport,
+            fetch_schema_from_transport=True,
+            introspection_args=introspection_args,
+        )
+
+        with client:
+            schema = client.schema
+
+        return schema
 
 
 class ServiceCore:
@@ -88,7 +168,7 @@ class ServiceCore:
         self._gateway = self.service._gateway
         self._input_value_deprecation = True
 
-        self._schema = SchemaCache(schema=None, expires=self.service.schema_expiry)
+        self._cache_map = SchemaCacheMap(expires=self.service.schema_expiry)
 
         self._queries = None
         self._mutations = None
@@ -134,20 +214,18 @@ class ServiceCore:
 
         client = Client(
             transport=transport,
-            schema=self._schema.schema,
-            fetch_schema_from_transport=self._schema.should_fetch_schema(),
-            introspection_args={
-                "input_value_deprecation": (
-                    bool(self.service.input_value_deprecation)
-                    if self.service.input_value_deprecation is not None
-                    else self._input_value_deprecation
-                )
-            },
+            schema=self._cache_map.get(
+                transport,
+                introspection_args={
+                    "input_value_deprecation": (
+                        bool(self.service.input_value_deprecation)
+                        if self.service.input_value_deprecation is not None
+                        else self._input_value_deprecation
+                    )
+                },
+            ),
         )
 
-        with client:
-            self._schema.schema = client.schema
-
         return client
 
     @property
@@ -158,12 +236,6 @@ class ServiceCore:
             f"access-token-{self.service.access_token}",
         ]
 
-        # we cannot build the async client and use it to fetch the schema
-        # due to async issues with the transport.  This will grab the schema
-        # from a sync client and use it to build the async client.
-        if not self._schema.schema:
-            self.sync_client  # pylint: disable=pointless-statement
-
         if self.service.tenant_id:
             subprotocols.append(f"x-tenant-context-{self.service.tenant_id}")
 
@@ -176,15 +248,7 @@ class ServiceCore:
 
         client = Client(
             transport=transport,
-            schema=self._schema.schema,
-            fetch_schema_from_transport=self._schema.should_fetch_schema(),
-            introspection_args={
-                "input_value_deprecation": (
-                    bool(self.service.input_value_deprecation)
-                    if self.service.input_value_deprecation is not None
-                    else self._input_value_deprecation
-                )
-            },
+            schema=self.schema,
         )
 
         return client
@@ -197,13 +261,11 @@ class ServiceCore:
         -------
         GraphQLSchema
         """
-        if not self._schema.schema:
-            self.sync_client  # pylint: disable=pointless-statement
-        return self._schema.schema
+        return self.sync_client.schema
 
     def clear_schema(self):
         """Clears the introspection schema."""
-        self._schema.schema = None
+        self._cache_map.clear(self.sync_client.transport)
 
     def get_sync_schema(self) -> Union[GraphQLSchema, None]:
         """Retrieves introspection schema from Synchronous endpoint.
@@ -212,10 +274,7 @@ class ServiceCore:
         -------
         GraphQLSchema
         """
-        client = self.sync_client
-        with client:
-            schema = client.schema
-        return schema
+        return self.schema
 
     @async_block
     async def get_ws_schema(self) -> Union[GraphQLSchema, None]:

{taegis-sdk-python-1.2.4 → taegis_sdk_python-1.2.6}/taegis_sdk_python/services/__init__.py

@@ -3,6 +3,7 @@ This needs to be a generated file.  Need to make jinja template.
 """
 
 from typing import Dict, Optional, Any
+import threading
 
 from taegis_sdk_python._consts import TAEGIS_ENVIRONMENT_URLS
 from taegis_sdk_python._version import __version__
@@ -20,6 +21,7 @@ from taegis_sdk_python.services.byoti import ByotiService
 from taegis_sdk_python.services.clients import ClientsService
 from taegis_sdk_python.services.collector import CollectorService
 from taegis_sdk_python.services.comments import CommentsService
+from taegis_sdk_python.services.contracted_endpoint import ContractedEndpointService
 from taegis_sdk_python.services.datasources import DatasourcesService
 from taegis_sdk_python.services.detector_registry import DetectorRegistryService
 from taegis_sdk_python.services.endpoint_command_manager import (
@@ -33,6 +35,7 @@ from taegis_sdk_python.services.events import EventsService
 from taegis_sdk_python.services.exports import ExportsService
 from taegis_sdk_python.services.fast_ioc import FastIocService
 from taegis_sdk_python.services.file_info import FileInfoService
+from taegis_sdk_python.services.ingest_stats import IngestStatsService
 from taegis_sdk_python.services.investigations import InvestigationsService
 from taegis_sdk_python.services.investigations2 import Investigations2Service
 from taegis_sdk_python.services.mitre_attack_info import MitreAttackInfoService
@@ -100,8 +103,8 @@ class GraphQLService:
 
         self._tenant_id = tenant_id
         self._gateway = gateway or "/graphql"
-        self._context_manager = {}
-        self._context_kwargs = []
+        self._thread_id = threading.get_ident()
+        self._context_kwargs = {}
         if not extra_headers:
             self._extra_headers = {}
         else:
@@ -120,6 +123,7 @@ class GraphQLService:
         self._clients = None
         self._collector = None
         self._comments = None
+        self._contracted_endpoint = None
         self._datasources = None
         self._detector_registry = None
         self._endpoint_command_manager = None
@@ -129,6 +133,7 @@ class GraphQLService:
         self._exports = None
         self._fast_ioc = None
         self._file_info = None
+        self._ingest_stats = None
         self._investigations = None
         self._investigations2 = None
         self._mitre_attack_info = None
@@ -151,22 +156,34 @@ class GraphQLService:
         self._users = None
 
     def __call__(self, **kwargs):
-        self._context_kwargs.append(kwargs)
+        if threading.get_ident() not in self._context_kwargs:
+            self._context_kwargs[threading.get_ident()] = []
+
+        self._context_kwargs[threading.get_ident()].append(kwargs)
         return self
 
     def __enter__(self):
-        for kwarg in self._context_kwargs:
-            self._context_manager.update(kwarg)
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self._context_manager.clear()
+        if self._context_kwargs[threading.get_ident()]:
+            self._context_kwargs[threading.get_ident()].pop()
+
+        else:
+            del self._context_kwargs[threading.get_ident()]
+
+    @property
+    def _context_manager(self):
+        """Internal Context Manager property."""
+        temp_context = {}
 
-        if self._context_kwargs:
-            self._context_kwargs.pop()
+        for kwarg in self._context_kwargs.get(self._thread_id, []):
+            temp_context.update(kwarg)
+        if self._thread_id != threading.get_ident():
+            for kwarg in self._context_kwargs.get(threading.get_ident(), []):
+                temp_context.update(kwarg)
 
-            for kwarg in self._context_kwargs:
-                self._context_manager.update(kwarg)
+        return temp_context
 
     @property
     def environment(self):
@@ -325,6 +342,13 @@ class GraphQLService:
             self._comments = CommentsService(self)
         return self._comments
 
+    @property
+    def contracted_endpoint(self):
+        """ContractedEndpoint Service Endpoint."""
+        if not self._contracted_endpoint:
+            self._contracted_endpoint = ContractedEndpointService(self)
+        return self._contracted_endpoint
+
     @property
     def datasources(self):
         """Datasources Service Endpoint."""
@@ -388,6 +412,13 @@ class GraphQLService:
             self._file_info = FileInfoService(self)
         return self._file_info
 
+    @property
+    def ingest_stats(self):
+        """IngestStats Service Endpoint."""
+        if not self._ingest_stats:
+            self._ingest_stats = IngestStatsService(self)
+        return self._ingest_stats
+
     @property
     def investigations(self):
         """Investigations Service Endpoint."""

{taegis-sdk-python-1.2.4 → taegis_sdk_python-1.2.6}/taegis_sdk_python/services/alerts/types.py

@@ -1350,6 +1350,35 @@ class UpdateThreatScoreResponse:
     )
 
 
+@dataclass_json
+@dataclass(order=True, eq=True, frozen=True)
+class MatchedYaraRule:
+    """MatchedYaraRule."""
+
+    rule_name: Optional[str] = field(
+        default=None, metadata=config(field_name="rule_name")
+    )
+    rule_description: Optional[str] = field(
+        default=None, metadata=config(field_name="rule_description")
+    )
+    classification: Optional[str] = field(
+        default=None, metadata=config(field_name="classification")
+    )
+    confidence: Optional[float] = field(
+        default=None, metadata=config(field_name="confidence")
+    )
+    severity: Optional[float] = field(
+        default=None, metadata=config(field_name="severity")
+    )
+    attack_technique_ids: Optional[List[str]] = field(
+        default=None, metadata=config(field_name="attack_technique_ids")
+    )
+    vid: Optional[str] = field(default=None, metadata=config(field_name="vid"))
+    rule_created_date: Optional[Timestamp] = field(
+        default=None, metadata=config(field_name="rule_created_date")
+    )
+
+
 @dataclass_json
 @dataclass(order=True, eq=True, frozen=True)
 class DDosIpCount:
@@ -2008,6 +2037,16 @@ class AlertsInvestigationInfo:
     )
 
 
+@dataclass_json
+@dataclass(order=True, eq=True, frozen=True)
+class FileAnalysisDetail:
+    """FileAnalysisDetail."""
+
+    matched_yara_rule: Optional[List[MatchedYaraRule]] = field(
+        default=None, metadata=config(field_name="matched_yara_rule")
+    )
+
+
 @dataclass_json
 @dataclass(order=True, eq=True, frozen=True)
 class HandsOnKeyboardDetails:
@@ -2178,6 +2217,9 @@ class EnrichmentDetail:
     ] = field(
         default=None, metadata=config(field_name="account_compromise_detector_detail")
     )
+    file_analysis_detail: Optional[FileAnalysisDetail] = field(
+        default=None, metadata=config(field_name="file_analysis_detail")
+    )
     generic: Optional[GenericDetail] = field(
         default=None, metadata=config(field_name="generic")
     )