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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: taegis-sdk-python
-Version: 1.5.4
+Version: 1.6.0
 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
@@ -198,6 +198,7 @@ Requires-Dist: dataclasses-json
 Requires-Dist: filelock>=3.12.4
 Requires-Dist: gql<4,>=3.5.0
 Requires-Dist: graphql-core
+Requires-Dist: jinja2
 Requires-Dist: oauthlib
 Requires-Dist: pyjwt>=2.8.0
 Requires-Dist: python-dateutil

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

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

{taegis_sdk_python-1.5.4 → taegis_sdk_python-1.6.0}/taegis_sdk_python/__init__.py

@@ -3,23 +3,22 @@ Commonly-used functions and data types from this package.
 """
 
 from taegis_sdk_python.errors import (
-    ServiceCoreException,
-    InvalidAuthenticationMethod,
     AccessTokenException,
-    MissingAccessTokenError,
-    InvalidAccessTokenError,
     GraphQLNoRowsInResultSetError,
     InvalidAccessTokenClaims,
+    InvalidAccessTokenError,
+    InvalidAuthenticationMethod,
     InvalidGraphQLEndpoint,
+    MissingAccessTokenError,
+    ServiceCoreException,
 )
-
 from taegis_sdk_python.services import GraphQLService
 from taegis_sdk_python.utils import (
     build_output_string,
+    build_output_string_from_introspection,
+    parse_union_result,
     prepare_input,
     prepare_variables,
-    parse_union_result,
-    build_output_string_from_introspection,
 )
 
 # pylint: disable=duplicate-code

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

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

taegis_sdk_python-1.6.0/taegis_sdk_python/commons/alerts/__init__.py

@@ -0,0 +1 @@
+"""Taegis Common Alerts Service Implementations."""

taegis_sdk_python-1.6.0/taegis_sdk_python/commons/alerts/federated_search.py

@@ -0,0 +1,163 @@
+"""Taegis Commons Federated Alerts Search implementation."""
+
+import logging
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List, Optional
+
+from dataclasses_json import config, dataclass_json
+
+from taegis_sdk_python import (
+    GraphQLNoRowsInResultSetError,
+    GraphQLService,
+    build_output_string,
+    prepare_input,
+)
+from taegis_sdk_python.services.alerts.types import (
+    Alert2,
+    AlertsList,
+    AlertsResponse,
+    AuxiliaryEvent,
+    PollRequestInput,
+    SearchRequestInput,
+)
+
+log = logging.getLogger(__name__)
+
+
+@dataclass_json
+@dataclass(order=True, eq=True, frozen=True)
+class TaegisCommonsAuxiliaryEvent(AuxiliaryEvent):
+    """My TaegisCommons Auxiliary Event - Extends Auxiliary Event with event_data
+    to take advantage of GQL federated services.
+    """
+
+    event_data: Optional[Dict[str, Any]] = field(
+        default=None, metadata=config(field_name="event_data")
+    )
+
+
+@dataclass_json
+@dataclass(order=True, eq=True, frozen=True)
+class TaegisCommonsAlert2(Alert2):
+    """My TaegisCommons Alert2."""
+
+    event_ids: Optional[List[TaegisCommonsAuxiliaryEvent]] = field(
+        default=None, metadata=config(field_name="event_ids")
+    )
+
+
+@dataclass_json
+@dataclass(order=True, eq=True, frozen=True)
+class TaegisCommonsAlertsList(AlertsList):
+    """My TaegisCommons AlertsList."""
+
+    list: Optional[List[TaegisCommonsAlert2]] = field(
+        default=None, metadata=config(field_name="list")
+    )
+
+
+@dataclass_json
+@dataclass(order=True, eq=True, frozen=True)
+class TaegisCommonsAlertsResponse(AlertsResponse):
+    """My TaegisCommons AlertsResponse."""
+
+    alerts: Optional[TaegisCommonsAlertsList] = field(
+        default=None, metadata=config(field_name="alerts")
+    )
+
+
+def alerts_service_search_with_events(
+    service: GraphQLService, in_: SearchRequestInput
+) -> TaegisCommonsAlertsResponse:
+    """Query Taegis Alerts with corresponding Events attached."""
+    endpoint = "alertsServiceSearch"
+    result = service.alerts.execute_query(
+        endpoint=endpoint,
+        variables={
+            "in": prepare_input(in_),
+        },
+        output=build_output_string(TaegisCommonsAlertsResponse),
+    )
+    if result is not None:
+        return TaegisCommonsAlertsResponse.from_dict(  # pylint: disable=no-member
+            result.get(endpoint)
+        )
+    raise GraphQLNoRowsInResultSetError("for query alertsServiceSearch")
+
+
+def alerts_service_poll_with_events(
+    service: GraphQLService, in_: PollRequestInput
+) -> TaegisCommonsAlertsResponse:
+    """Federated alerts_service_poll with event_data."""
+    endpoint = "alertsServicePoll"
+    result = service.alerts.execute_query(
+        endpoint=endpoint,
+        variables={
+            "in": prepare_input(in_),
+        },
+        output=build_output_string(TaegisCommonsAlertsResponse),
+    )
+    if result is not None:
+        return TaegisCommonsAlertsResponse.from_dict(  # pylint: disable=no-member
+            result.get(endpoint)
+        )
+    raise GraphQLNoRowsInResultSetError("for query alertsServicePoll")
+
+
+def alerts_federated_search(
+    service: GraphQLService,
+    query: str,
+    *,
+    limit: int = 10000,
+    caller_name: str = "Taegis SDK Commons",
+    federated_call: Callable,
+    federated_poll_call: Callable,
+) -> List[TaegisCommonsAlertsResponse]:
+    """
+    Search Taegis Alerts service.
+    """
+    if "aggregate" in query:
+        limit = 1
+
+    result = federated_call(
+        service,
+        SearchRequestInput(
+            cql_query=query,
+            offset=0,
+            limit=limit,
+            metadata={"callerName": caller_name},
+        ),
+    )
+
+    poll_responses = [result]
+    search_id = result.search_id
+    total_parts = result.alerts.total_parts
+
+    if search_id:
+        for part in range(2, total_parts + 1):
+            response = None
+            try:
+                log.debug(f"Submitting page {part}...")
+                response = federated_poll_call(
+                    service,
+                    PollRequestInput(
+                        search_id=search_id,
+                        part_id=part,
+                    ),
+                )
+            except Exception as exc:  # pylint: disable=broad-exception-caught
+                log.error(
+                    f"Cannot retrieve results for search_id:{search_id}:{part}::{exc}"
+                )
+                if "not found" in str(exc):
+                    break
+
+            if isinstance(response, AlertsResponse) and response.alerts is not None:
+                poll_responses.append(response)
+                # CX-92571 work around
+                if sum(len(response.alerts.list) for response in poll_responses) >= int(
+                    limit
+                ):
+                    break
+
+    return poll_responses

taegis_sdk_python-1.6.0/taegis_sdk_python/commons/alerts/search.py

@@ -0,0 +1,70 @@
+"""Taegis Commons Alerts Search implementations."""
+
+import logging
+from typing import List
+
+from taegis_sdk_python import GraphQLService
+from taegis_sdk_python.services.alerts.types import (
+    AlertsResponse,
+    PollRequestInput,
+    SearchRequestInput,
+)
+
+log = logging.getLogger(__name__)
+
+
+def alerts_search(
+    service: GraphQLService,
+    query: str,
+    *,
+    limit: int = 10000,
+    caller_name: str = "Taegis SDK Commons",
+) -> List[AlertsResponse]:
+    """
+    Search Taegis Alerts service.
+    """
+    if "aggregate" in query:
+        limit = 1
+
+    result = service.alerts.query.alerts_service_search(
+        service,
+        SearchRequestInput(
+            cql_query=query,
+            offset=0,
+            limit=limit,
+            metadata={"callerName": caller_name},
+        ),
+    )
+
+    poll_responses = [result]
+    search_id = result.search_id
+    total_parts = result.alerts.total_parts
+
+    if search_id:
+        for part in range(2, total_parts + 1):
+            response = None
+            try:
+                log.debug(f"Submitting page {part}...")
+                response = service.alerts.query.alerts_service_poll(
+                    service,
+                    PollRequestInput(
+                        search_id=search_id,
+                        part_id=part,
+                    ),
+                )
+            except Exception as exc:  # pylint: disable=broad-exception-caught
+                log.error(
+                    f"Cannot retrieve results for search_id:{search_id}:{part}::{exc}"
+                )
+                if "not found" in str(exc):
+                    break
+
+            if isinstance(response, AlertsResponse) and response.alerts is not None:
+                poll_responses.append(response)
+                # CX-92571 work around
+                if sum(len(response.alerts.list) for response in poll_responses) >= int(
+                    limit
+                ):
+                    break
+
+    return poll_responses

taegis_sdk_python-1.6.0/taegis_sdk_python/commons/cases/__init__.py

@@ -0,0 +1 @@
+"""Taegis Common Cases Service Implementations."""

taegis_sdk_python-1.6.0/taegis_sdk_python/commons/cases/federated_search.py

@@ -0,0 +1,179 @@
+"""Taegis Commons Cases implementations."""
+
+import logging
+import re
+from dataclasses import dataclass, field
+from typing import Callable, List, Optional
+
+from dataclasses_json import config, dataclass_json
+from taegis_magic.core.utils import remove_output_node
+
+from taegis_sdk_python import (
+    GraphQLNoRowsInResultSetError,
+    GraphQLService,
+    build_output_string,
+    prepare_input,
+)
+from taegis_sdk_python.services.investigations2.types import (
+    CreateInvestigationInput,
+    InvestigationsV2,
+    InvestigationsV2Arguments,
+    InvestigationV2,
+)
+from taegis_sdk_python.services.subjects.types import Subject as FederatedSubject
+
+log = logging.getLogger(__name__)
+
+
+@dataclass_json
+@dataclass(order=True, eq=True, frozen=True)
+class TaegisCommonsInvestigationV2(InvestigationV2):
+    """TaegisCommonsInvestigationV2"""
+
+    contributor_subjects: Optional[List[FederatedSubject]] = field(
+        default=None, metadata=config(field_name="contributorSubjects")
+    )
+    assignee_subject: Optional[FederatedSubject] = field(
+        default=None, metadata=config(field_name="assigneeSubject")
+    )
+    created_by_subject: Optional[FederatedSubject] = field(
+        default=None, metadata=config(field_name="createdBySubject")
+    )
+    updated_by_subject: Optional[FederatedSubject] = field(
+        default=None, metadata=config(field_name="updatedBySubject")
+    )
+
+
+@dataclass_json
+@dataclass(order=True, eq=True, frozen=True)
+class TaegisCommonsInvestigationsV2(InvestigationsV2):
+    """TaegisCommonsInvestigationsV2"""
+
+    investigations: List[TaegisCommonsInvestigationV2] = field(
+        default_factory=list, metadata=config(field_name="investigations")
+    )
+
+
+def investigations_create_with_subjects(
+    service: GraphQLService, input_: CreateInvestigationInput
+) -> TaegisCommonsInvestigationV2:
+    """createInvestigationV2 creates new investigation with the provided arguments."""
+    endpoint = "createInvestigationV2"
+
+    result = service.investigations2.execute_mutation(
+        endpoint=endpoint,
+        variables={
+            "input": prepare_input(input_),
+        },
+        output=build_output_string(TaegisCommonsInvestigationV2),
+    )
+    if result.get(endpoint) is not None:
+        return TaegisCommonsInvestigationV2.from_dict(  # pylint: disable=no-member
+            result.get(endpoint)
+        )
+    raise GraphQLNoRowsInResultSetError("for mutation createInvestigationV2")
+
+
+def investigations_search_with_subjects(
+    service, arguments: InvestigationsV2Arguments
+) -> TaegisCommonsInvestigationsV2:
+    """investigationsV2 returns a list of investigations matching the provided arguments."""
+    endpoint = "investigationsV2"
+
+    result = service.investigations2.execute_query(
+        endpoint=endpoint,
+        variables={
+            "arguments": prepare_input(arguments),
+        },
+        output=build_output_string(TaegisCommonsInvestigationsV2),
+    )
+    if result.get(endpoint) is not None:
+        return TaegisCommonsInvestigationsV2.from_dict(  # pylint: disable=no-member
+            result.get(endpoint)
+        )
+    raise GraphQLNoRowsInResultSetError("for query investigationsV2")
+
+
+def cases_federated_search(
+    service: GraphQLService,
+    query: str,
+    *,
+    limit: int = 10000,
+    federated_call: Callable,
+) -> List[InvestigationsV2]:
+    """Taegis Cases search with Subject federation."""
+
+    page = 1
+    per_page = 100
+
+    results = []
+
+    # fix for CX-99036
+    pattern = r"\|\s*(head|tail)\s*([0-9]+)"
+    match = re.search(pattern, query)
+
+    if not limit:
+        if match and match.group(1) == "tail":  # pragma: no cover
+            log.warning(
+                "tail is not currently supported, it will be used as the limit..."
+            )
+
+        if match:
+            limit = int(match.group(2))
+    elif match:  # pragma: no cover
+        log.warning(
+            f"limit and {match.group(1)} both provided, only limit will be honored..."
+        )
+
+    query = re.sub(pattern, "", query)
+
+    if limit and limit < per_page:
+        per_page = limit
+    # endfix
+
+    # fix for CX-103490
+    output = build_output_string(TaegisCommonsInvestigationsV2)
+
+    output = remove_output_node(output, "metric")
+    output = remove_output_node(output, "metrics")
+    # endfix
+
+    with service(output=output):
+        investigations_results = federated_call(
+            service=service,
+            arguments=InvestigationsV2Arguments(
+                page=page,
+                per_page=per_page,
+                cql=query,
+            ),
+        )
+
+    results.append(investigations_results)
+
+    # fix for CX-99036
+    if not limit or investigations_results.total_count < limit:
+        limit = investigations_results.total_count
+    # endfix
+
+    while (
+        sum_results := sum(len(result.investigations) for result in results)
+    ) < limit:
+        page += 1
+
+        # fix for CX-99036
+        if (per_page * page) > limit:
+            per_page = limit - sum_results
+        # endfix
+
+        with service(output=output):
+            investigations_results = federated_call(
+                service=service,
+                arguments=InvestigationsV2Arguments(
+                    page=page,
+                    per_page=per_page,
+                    cql=query,
+                ),
+            )
+        results.append(investigations_results)
+
+    return results

taegis_sdk_python-1.6.0/taegis_sdk_python/commons/cases/search.py

@@ -0,0 +1,98 @@
+"""Taegis Commons Cases implementations."""
+
+import logging
+import re
+from typing import List
+
+from taegis_sdk_python import GraphQLService, build_output_string
+from taegis_sdk_python.services.investigations2.types import (
+    InvestigationsV2,
+    InvestigationsV2Arguments,
+)
+from taegis_sdk_python.utils import remove_output_node
+
+log = logging.getLogger(__name__)
+
+
+def cases_search(
+    service: GraphQLService,
+    query: str,
+    *,
+    limit: int = 10000,
+) -> List[InvestigationsV2]:
+    """Taegis Cases search."""
+
+    page = 1
+    per_page = 100
+
+    results = []
+
+    # fix for CX-99036
+    pattern = r"\|\s*(head|tail)\s*([0-9]+)"
+    match = re.search(pattern, query)
+
+    if not limit:
+        if match and match.group(1) == "tail":  # pragma: no cover
+            log.warning(
+                "tail is not currently supported, it will be used as the limit..."
+            )
+
+        if match:
+            limit = int(match.group(2))
+    elif match:  # pragma: no cover
+        log.warning(
+            f"limit and {match.group(1)} both provided, only limit will be honored..."
+        )
+
+    query = re.sub(pattern, "", query)
+
+    if limit and limit < per_page:
+        per_page = limit
+    # endfix
+
+    # fix for CX-103490
+    output = build_output_string(InvestigationsV2)
+
+    output = remove_output_node(output, "metric")
+    output = remove_output_node(output, "metrics")
+    # endfix
+
+    with service(output=output):
+        investigations_results = service.investigations2.query.investigations_v2(
+            service=service,
+            arguments=InvestigationsV2Arguments(
+                page=page,
+                per_page=per_page,
+                cql=query,
+            ),
+        )
+
+    results.append(investigations_results)
+
+    # fix for CX-99036
+    if not limit or investigations_results.total_count < limit:
+        limit = investigations_results.total_count
+    # endfix
+
+    while (
+        sum_results := sum(len(result.investigations) for result in results)
+    ) < limit:
+        page += 1
+
+        # fix for CX-99036
+        if (per_page * page) > limit:
+            per_page = limit - sum_results
+        # endfix
+
+        with service(output=output):
+            investigations_results = service.investigations2.query.investigations_v2(
+                service=service,
+                arguments=InvestigationsV2Arguments(
+                    page=page,
+                    per_page=per_page,
+                    cql=query,
+                ),
+            )
+        results.append(investigations_results)
+
+    return results

taegis_sdk_python-1.6.0/taegis_sdk_python/commons/events/__init__.py

@@ -0,0 +1 @@
+"""Taegis Common Events Service Implementations."""

taegis_sdk_python-1.6.0/taegis_sdk_python/commons/events/search.py

@@ -0,0 +1,55 @@
+"""Taegis Commons Events search implementations."""
+
+import logging
+from typing import List, Optional
+
+from taegis_sdk_python import GraphQLService
+from taegis_sdk_python.services.events.types import EventQueryOptions, EventQueryResults
+
+log = logging.getLogger(__name__)
+
+
+def get_next_page(events_results: List[EventQueryResults]) -> Optional[str]:
+    """Retrieve events  next page indicator."""
+    try:
+        # the next page could be found in any of the result pages,
+        # but we cannot garuntee which result it will be found in
+        return next(
+            iter({result.next for result in events_results if result.next is not None})
+        )
+    except StopIteration:
+        return None
+
+
+def events_search(
+    service: GraphQLService,
+    query: str,
+    *,
+    options: Optional[EventQueryOptions] = None,
+    caller_name: str = "Taegis SDK Commons",
+) -> List[EventQueryResults]:
+    """Taegis Events search."""
+    if not options:
+        options = EventQueryOptions(
+            timestamp_ascending=True,
+            page_size=1000,
+            max_rows=100000,
+            aggregation_off=False,
+        )
+
+    results = []
+
+    result = service.events.subscription.event_query(
+        query=query,
+        options=options,
+        metadata={"callerName": caller_name},
+    )
+    results.append(result)
+    next_page = get_next_page(result)
+
+    while next_page:
+        result = service.events.subscription.event_page(next_page)
+        results.append(result)
+        next_page = get_next_page(result)
+
+    return results

taegis_sdk_python-1.6.0/taegis_sdk_python/commons/investigations/__init__.py

@@ -0,0 +1 @@
+"""Taegis Common Investigations Service Implementations."""