From ec64caa54575ebc26d1521d62ce5bc4a185d037b Mon Sep 17 00:00:00 2001 From: LegendEvent Date: Tue, 17 Feb 2026 18:12:14 +0100 Subject: [PATCH 1/2] feat: add configurable request timeout support - Add timeout parameter to DarktraceClient (default: None for backwards compatibility) - Add per-request timeout override to all endpoint methods - Support tuple format: timeout=(connect_timeout, read_timeout) - Add TimeoutType export for type hints - Add comprehensive test suite (11 tests) - Update README and docs with timeout documentation This enables users to: - Set client-wide timeout: DarktraceClient(timeout=30) - Override per-request: client.advanced_search.search(query, timeout=600) - Use granular timeouts: timeout=(5, 30) for connect/read --- README.md | 83 +++++++++++-- darktrace/__init__.py | 5 +- darktrace/client.py | 28 ++++- darktrace/dt_advanced_search.py | 20 ++-- darktrace/dt_analyst.py | 57 +++++---- darktrace/dt_antigena.py | 50 +++++--- darktrace/dt_breaches.py | 50 ++++---- darktrace/dt_components.py | 7 +- darktrace/dt_cves.py | 6 +- darktrace/dt_details.py | 6 +- darktrace/dt_deviceinfo.py | 6 +- darktrace/dt_devices.py | 11 +- darktrace/dt_devicesearch.py | 40 ++++--- darktrace/dt_devicesummary.py | 6 +- darktrace/dt_email.py | 92 ++++++++++----- darktrace/dt_endpointdetails.py | 8 +- darktrace/dt_enums.py | 7 +- darktrace/dt_filtertypes.py | 7 +- darktrace/dt_intelfeed.py | 24 ++-- darktrace/dt_mbcomments.py | 21 +++- darktrace/dt_metricdata.py | 7 +- darktrace/dt_metrics.py | 6 +- darktrace/dt_models.py | 7 +- darktrace/dt_network.py | 8 +- darktrace/dt_pcaps.py | 12 +- darktrace/dt_similardevices.py | 7 +- darktrace/dt_status.py | 11 +- darktrace/dt_subnets.py | 16 ++- darktrace/dt_summarystatistics.py | 11 +- darktrace/dt_tags.py | 57 ++++++--- darktrace/dt_utils.py | 11 +- docs/README.md | 21 ++++ tests/test_timeout.py | 186 ++++++++++++++++++++++++++++++ 33 files changed, 680 insertions(+), 214 deletions(-) create mode 100644 tests/test_timeout.py diff --git a/README.md b/README.md index ef34275..6219a31 100644 --- a/README.md +++ b/README.md @@ -12,16 +12,13 @@ --- -## 🆕 Latest Updates (v0.8.55) +## 🆕 Latest Updates (v0.8.56) -- **Feature: Add 13 missing parameters to devicesummary endpoint** - Added support for `device_name`, `ip_address`, `end_timestamp`, `start_timestamp`, `devicesummary_by`, `devicesummary_by_value`, `device_type`, `network_location`, `network_location_id`, `peer_id`, `source`, and `status` parameters to align with Darktrace API specification -- **Documentation: Update devicesummary documentation** - Added examples and parameter descriptions for new filtering options -- **Note: devicesummary HTTP 500 limitation confirmed** - Documentation updated to clarify that all devicesummary parameters return HTTP 500 with API token authentication (Darktrace backend limitation, not SDK bug) +- **Feature: Configurable request timeouts** - Added `timeout` parameter to `DarktraceClient` and all endpoint methods. Supports client-level defaults and per-request overrides. Use for long-running queries like advanced search. +- **Security: Enable SSL certificate verification by default (fixes #47)** - Changed `verify_ssl` default from `False` to `True` for secure-by-default behavior. All endpoint modules now inherit the client's SSL verification setting. +- **Documentation: Add SSL certificate setup guide** - Added instructions for using self-signed certificates with `verify_ssl=True` via system trust store or environment variable. -## 📝 Previous Updates (v0.8.54) - -- **Fix: Multi-parameter devicesearch query format (fixes #45)** - Changed query parameter joining from explicit ' AND ' to space separation per Darktrace API specification -- **Fix: ensure host URL includes protocol (default to https if missing)** +> For previous updates, see [GitHub Releases](https://github.com/LegendEvent/darktrace-sdk/releases). --- @@ -32,6 +29,7 @@ - **Modular & Maintainable**: Each endpoint group is a separate Python module/class. - **Easy Authentication**: Secure HMAC-SHA1 signature generation and token management. - **SSL Verification**: SSL certificate verification is enabled by default for secure connections. +- **Configurable Timeouts**: Client-level and per-request timeout support for long-running queries. - **Async-Ready**: Designed for easy extension to async workflows. - **Type Hints & Docstrings**: Full typing and documentation for all public methods. - **Comprehensive Documentation**: Detailed documentation for every module and endpoint. @@ -55,6 +53,75 @@ client = DarktraceClient( > ⚠️ **Warning**: Disabling SSL verification exposes your connection to man-in-the-middle attacks. Never disable in production environments. +### Using Self-Signed Certificates with verify_ssl=True + +For production environments with self-signed certificates, add the certificate to your system trust store instead of disabling verification: + +```bash +# 1. Get the certificate from your Darktrace instance +openssl s_client -showcerts -connect your-darktrace-instance:443 /dev/null | openssl x509 -outform PEM > ~/darktrace-cert.pem + +# 2. Copy to system CA store (Linux/Ubuntu/Debian) +sudo cp ~/darktrace-cert.pem /usr/local/share/ca-certificates/darktrace-cert.crt +sudo update-ca-certificates + +# 3. Now verify_ssl=True will work +``` + +**Alternative (no sudo required):** +```bash +# Create a custom CA bundle and set environment variable +cat /etc/ssl/certs/ca-certificates.crt ~/darktrace-cert.pem > ~/.custom-ca-bundle.pem +export REQUESTS_CA_BUNDLE=~/.custom-ca-bundle.pem +``` + +--- + +## ⏱️ Request Timeouts + +The SDK supports configurable request timeouts at both client and per-request levels. + +### Client-Level Timeout + +Set a default timeout for all requests: + +```python +from darktrace import DarktraceClient + +# 30 second timeout for all requests +client = DarktraceClient( + host="https://your-darktrace-instance", + public_token="YOUR_PUBLIC_TOKEN", + private_token="YOUR_PRIVATE_TOKEN", + timeout=30 +) +``` + +### Per-Request Timeout + +Override the timeout for specific requests (e.g., long-running advanced searches): + +```python +# Client default: 30 seconds +client = DarktraceClient(host="...", public_token="...", private_token="...", timeout=30) + +# Override for slow query (5 minutes) +results = client.advanced_search.search(query, timeout=300) + +# Override with tuple format (connect_timeout, read_timeout) +results = client.advanced_search.search(query, timeout=(5, 300)) +``` + +### Timeout Format + +| Format | Description | +|--------|-------------| +| `timeout=None` | No timeout (default, waits indefinitely) | +| `timeout=30` | 30 seconds total (both connect and read) | +| `timeout=(5, 30)` | 5 seconds to connect, 30 seconds to read | + +> **Note**: Advanced search queries can take 5-10 minutes for complex queries. Consider using per-request timeouts for these endpoints. + --- ## 📦 Installation diff --git a/darktrace/__init__.py b/darktrace/__init__.py index 969a36e..088d4c5 100644 --- a/darktrace/__init__.py +++ b/darktrace/__init__.py @@ -22,7 +22,7 @@ from .dt_subnets import Subnets from .dt_summarystatistics import SummaryStatistics from .dt_tags import Tags -from .dt_utils import debug_print +from .dt_utils import debug_print, TimeoutType from .dt_components import Components from .dt_cves import CVEs from .dt_details import Details @@ -61,5 +61,6 @@ 'DeviceSearch', 'ModelBreaches', 'AdvancedSearch', - 'debug_print' + 'debug_print', + 'TimeoutType', ] \ No newline at end of file diff --git a/darktrace/client.py b/darktrace/client.py index 73f67fd..63af342 100644 --- a/darktrace/client.py +++ b/darktrace/client.py @@ -1,10 +1,12 @@ +from typing import TYPE_CHECKING, Union, Tuple + from .auth import DarktraceAuth from .dt_antigena import Antigena from .dt_analyst import Analyst from .dt_breaches import ModelBreaches from .dt_devices import Devices from .dt_email import DarktraceEmail -from .dt_utils import debug_print +from .dt_utils import debug_print, TimeoutType from .dt_advanced_search import AdvancedSearch from .dt_components import Components from .dt_cves import CVEs @@ -28,8 +30,6 @@ from .dt_summarystatistics import SummaryStatistics from .dt_tags import Tags -from typing import TYPE_CHECKING - if TYPE_CHECKING: from .dt_antigena import Antigena from .dt_analyst import Analyst @@ -65,6 +65,7 @@ class DarktraceClient: auth: DarktraceAuth debug: bool verify_ssl: bool + timeout: TimeoutType advanced_search: 'AdvancedSearch' antigena: 'Antigena' analyst: 'Analyst' @@ -93,7 +94,15 @@ class DarktraceClient: summarystatistics: 'SummaryStatistics' tags: 'Tags' - def __init__(self, host: str, public_token: str, private_token: str, debug: bool = False, verify_ssl: bool = True) -> None: + def __init__( + self, + host: str, + public_token: str, + private_token: str, + debug: bool = False, + verify_ssl: bool = True, + timeout: TimeoutType = None + ) -> None: """ Initialize the Darktrace API client. @@ -104,6 +113,8 @@ def __init__(self, host: str, public_token: str, private_token: str, debug: bool debug (bool, optional): Enable debug logging. Defaults to False. verify_ssl (bool, optional): Enable SSL certificate verification. Defaults to True. Set to False only for development/testing with self-signed certificates. + timeout (float|tuple, optional): Request timeout in seconds. Can be a single float + or a tuple of (connect_timeout, read_timeout). None means no timeout (default). Example: >>> client = DarktraceClient( @@ -112,6 +123,14 @@ def __init__(self, host: str, public_token: str, private_token: str, debug: bool ... private_token="your_private_token", ... debug=True ... ) + + >>> # With timeout + >>> client = DarktraceClient( + ... host="https://your-instance.darktrace.com", + ... public_token="your_public_token", + ... private_token="your_private_token", + ... timeout=30 # 30 second timeout for all requests + ... ) """ # Ensure host has a protocol @@ -123,6 +142,7 @@ def __init__(self, host: str, public_token: str, private_token: str, debug: bool self.auth = DarktraceAuth(public_token, private_token) self.debug = debug self.verify_ssl = verify_ssl + self.timeout = timeout # Endpoint groups self.advanced_search = AdvancedSearch(self) diff --git a/darktrace/dt_advanced_search.py b/darktrace/dt_advanced_search.py index 6d335a0..17559fe 100644 --- a/darktrace/dt_advanced_search.py +++ b/darktrace/dt_advanced_search.py @@ -1,13 +1,13 @@ import requests import json -from typing import Dict, Any +from typing import Dict, Any, Optional, Union, Tuple from .dt_utils import debug_print, BaseEndpoint, encode_query class AdvancedSearch(BaseEndpoint): def __init__(self, client): super().__init__(client) - def search(self, query: Dict[str, Any], post_request: bool = False): + def search(self, query: Dict[str, Any], post_request: bool = False, timeout: Optional[Union[float, Tuple[float, float]]] = None): """Perform Advanced Search query. Parameters: @@ -59,7 +59,8 @@ def search(self, query: Dict[str, Any], post_request: bool = False): headers, sorted_params = self._get_headers(endpoint, json_body=body) headers['Content-Type'] = 'application/json' self.client._debug(f"POST {url} body={body}") - response = requests.post(url, headers=headers, data=json.dumps(body, separators=(',', ':')), verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.post(url, headers=headers, data=json.dumps(body, separators=(',', ':')), verify=self.client.verify_ssl, timeout=resolved_timeout) self.client._debug(f"Response status: {response.status_code}") self.client._debug(f"Response text: {response.text}") response.raise_for_status() @@ -89,28 +90,31 @@ def search(self, query: Dict[str, Any], post_request: bool = False): url = f"{self.client.host}{endpoint}/{encoded_query}" headers, sorted_params = self._get_headers(f"{endpoint}/{encoded_query}") self.client._debug(f"GET {url}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def analyze(self, field: str, analysis_type: str, query: Dict[str, Any]): + def analyze(self, field: str, analysis_type: str, query: Dict[str, Any], timeout: Optional[Union[float, Tuple[float, float]]] = None): """Analyze field data.""" encoded_query = encode_query(query) endpoint = f'/advancedsearch/api/analyze/{field}/{analysis_type}/{encoded_query}' url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint) self.client._debug(f"GET {url}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def graph(self, graph_type: str, interval: int, query: Dict[str, Any]): + def graph(self, graph_type: str, interval: int, query: Dict[str, Any], timeout: Optional[Union[float, Tuple[float, float]]] = None): """Get graph data.""" encoded_query = encode_query(query) endpoint = f'/advancedsearch/api/graph/{graph_type}/{interval}/{encoded_query}' url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint) self.client._debug(f"GET {url}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() \ No newline at end of file diff --git a/darktrace/dt_analyst.py b/darktrace/dt_analyst.py index e3bae86..0a3cbe5 100644 --- a/darktrace/dt_analyst.py +++ b/darktrace/dt_analyst.py @@ -1,13 +1,13 @@ import requests import json -from typing import Union, List, Dict, Any, Optional +from typing import Union, List, Dict, Any, Optional, Tuple from .dt_utils import debug_print, BaseEndpoint class Analyst(BaseEndpoint): def __init__(self, client): super().__init__(client) - def get_groups(self, **params): + def get_groups(self, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): """Get AI Analyst incident groups. Available parameters: @@ -32,12 +32,13 @@ def get_groups(self, **params): endpoint = '/aianalyst/groups' url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint, params) + resolved_timeout = self._resolve_timeout(timeout) self.client._debug(f"GET {url} params={params}") - response = requests.get(url, headers=headers, params=sorted_params or params, verify=self.client.verify_ssl) + response = requests.get(url, headers=headers, params=sorted_params or params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def get_incident_events(self, **params): + def get_incident_events(self, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): """Get AI Analyst incident events. Available parameters: @@ -67,12 +68,13 @@ def get_incident_events(self, **params): endpoint = '/aianalyst/incidentevents' url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint, params) + resolved_timeout = self._resolve_timeout(timeout) self.client._debug(f"GET {url} params={params}") - response = requests.get(url, headers=headers, params=sorted_params or params, verify=self.client.verify_ssl) + response = requests.get(url, headers=headers, params=sorted_params or params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def acknowledge(self, uuids: Union[str, List[str]]) -> dict: + def acknowledge(self, uuids: Union[str, List[str]], timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: """ Acknowledge AI Analyst incident events. @@ -84,12 +86,13 @@ def acknowledge(self, uuids: Union[str, List[str]]) -> dict: url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint) headers['Content-Type'] = 'application/x-www-form-urlencoded' + resolved_timeout = self._resolve_timeout(timeout) self.client._debug(f"POST {url} data=uuid={uuids}") - response = requests.post(url, headers=headers, params=sorted_params, data={'uuid': uuids}, verify=self.client.verify_ssl) + response = requests.post(url, headers=headers, params=sorted_params, data={'uuid': uuids}, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def unacknowledge(self, uuids: Union[str, List[str]]) -> dict: + def unacknowledge(self, uuids: Union[str, List[str]], timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: """ Unacknowledge AI Analyst incident events. @@ -101,12 +104,13 @@ def unacknowledge(self, uuids: Union[str, List[str]]) -> dict: url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint) headers['Content-Type'] = 'application/x-www-form-urlencoded' + resolved_timeout = self._resolve_timeout(timeout) self.client._debug(f"POST {url} data=uuid={uuids}") - response = requests.post(url, headers=headers, params=sorted_params, data={'uuid': uuids}, verify=self.client.verify_ssl) + response = requests.post(url, headers=headers, params=sorted_params, data={'uuid': uuids}, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def pin(self, uuids: Union[str, List[str]]) -> dict: + def pin(self, uuids: Union[str, List[str]], timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: """ Pin AI Analyst incident events. @@ -118,12 +122,13 @@ def pin(self, uuids: Union[str, List[str]]) -> dict: url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint) headers['Content-Type'] = 'application/x-www-form-urlencoded' + resolved_timeout = self._resolve_timeout(timeout) self.client._debug(f"POST {url} data=uuid={uuids}") - response = requests.post(url, headers=headers, params=sorted_params, data={'uuid': uuids}, verify=self.client.verify_ssl) + response = requests.post(url, headers=headers, params=sorted_params, data={'uuid': uuids}, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def unpin(self, uuids: Union[str, List[str]]) -> dict: + def unpin(self, uuids: Union[str, List[str]], timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: """ Unpin AI Analyst incident events. @@ -135,12 +140,13 @@ def unpin(self, uuids: Union[str, List[str]]) -> dict: url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint) headers['Content-Type'] = 'application/x-www-form-urlencoded' + resolved_timeout = self._resolve_timeout(timeout) self.client._debug(f"POST {url} data=uuid={uuids}") - response = requests.post(url, headers=headers, params=sorted_params, data={'uuid': uuids}, verify=self.client.verify_ssl) + response = requests.post(url, headers=headers, params=sorted_params, data={'uuid': uuids}, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def get_comments(self, incident_id: str, response_data: Optional[str] = ""): + def get_comments(self, incident_id: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, response_data: Optional[str] = ""): """Get comments for an AI Analyst incident event. Parameters: @@ -153,12 +159,13 @@ def get_comments(self, incident_id: str, response_data: Optional[str] = ""): params['responsedata'] = response_data url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint, params) + resolved_timeout = self._resolve_timeout(timeout) self.client._debug(f"GET {url} params={params}") - response = requests.get(url, headers=headers, params=sorted_params or params, verify=self.client.verify_ssl) + response = requests.get(url, headers=headers, params=sorted_params or params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def add_comment(self, incident_id: str, message: str) -> dict: + def add_comment(self, incident_id: str, message: str, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: """Add a comment to an AI Analyst incident event. Parameters: @@ -169,16 +176,17 @@ def add_comment(self, incident_id: str, message: str) -> dict: url = f"{self.client.host}{endpoint}" body: Dict[str, Any] = {"incident_id": incident_id, "message": message} headers, sorted_params = self._get_headers(endpoint, json_body=body) + resolved_timeout = self._resolve_timeout(timeout) self.client._debug(f"POST {url} body={body}") # Send JSON as raw data with consistent formatting (same as signature generation) json_data = json.dumps(body, separators=(',', ':')) - response = requests.post(url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl) + response = requests.post(url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl, timeout=resolved_timeout) self.client._debug(f"Response Status: {response.status_code}") self.client._debug(f"Response Text: {response.text}") response.raise_for_status() return response.json() - def get_stats(self, **params): + def get_stats(self, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): """Get AI Analyst statistics. Available parameters: @@ -198,12 +206,13 @@ def get_stats(self, **params): endpoint = '/aianalyst/stats' url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint, params) + resolved_timeout = self._resolve_timeout(timeout) self.client._debug(f"GET {url} params={params}") - response = requests.get(url, headers=headers, params=sorted_params or params, verify=self.client.verify_ssl) + response = requests.get(url, headers=headers, params=sorted_params or params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def get_investigations(self, **params): + def get_investigations(self, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): """Get AI Analyst investigations (GET request). Available parameters: @@ -225,12 +234,13 @@ def get_investigations(self, **params): endpoint = '/aianalyst/investigations' url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint, params) + resolved_timeout = self._resolve_timeout(timeout) self.client._debug(f"GET {url} params={params}") - response = requests.get(url, headers=headers, params=sorted_params or params, verify=self.client.verify_ssl) + response = requests.get(url, headers=headers, params=sorted_params or params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def create_investigation(self, investigate_time: str, did: int): + def create_investigation(self, investigate_time: str, did: int, timeout: Optional[Union[float, Tuple[float, float]]] = None): """Create a new AI Analyst investigation (POST request). Parameters: @@ -248,11 +258,12 @@ def create_investigation(self, investigate_time: str, did: int): # For POST requests with JSON body, include it in signature headers, sorted_params = self._get_headers(endpoint, json_body=body) + resolved_timeout = self._resolve_timeout(timeout) self.client._debug(f"POST {url} json={body}") # Send JSON as raw data with consistent formatting (same as signature generation) json_data = json.dumps(body, separators=(',', ':')) - response = requests.post(url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl) + response = requests.post(url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl, timeout=resolved_timeout) self.client._debug(f"Response Status: {response.status_code}") self.client._debug(f"Response Text: {response.text}") response.raise_for_status() diff --git a/darktrace/dt_antigena.py b/darktrace/dt_antigena.py index 63346a2..2ef2d18 100644 --- a/darktrace/dt_antigena.py +++ b/darktrace/dt_antigena.py @@ -1,6 +1,6 @@ import requests import json -from typing import Dict, Any, Union, Optional, List +from typing import Dict, Any, Union, Optional, List, Tuple from .dt_utils import debug_print, BaseEndpoint @@ -20,7 +20,7 @@ class Antigena(BaseEndpoint): def __init__(self, client): super().__init__(client) - def get_actions(self, **params): + def get_actions(self, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): """ Get information about current and past Darktrace RESPOND actions. @@ -28,6 +28,8 @@ def get_actions(self, **params): and all historic actions with an expiry date in the last 14 days. Args: + timeout (float or tuple, optional): Timeout for the request in seconds. Can be a single + float for both connect and read timeouts, or a tuple of (connect_timeout, read_timeout). fulldevicedetails (bool): Returns the full device detail objects for all devices referenced by data in an API response. Use of this parameter will alter the JSON structure of the API response for certain calls. @@ -66,14 +68,15 @@ def get_actions(self, **params): url = f"{self.client.host}{endpoint}" self.client._debug(f"GET {url} params={sorted_params}") + resolved_timeout = self._resolve_timeout(timeout) response = requests.get( - url, headers=headers, params=sorted_params, verify=self.client.verify_ssl + url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout ) response.raise_for_status() return response.json() def activate_action( - self, codeid: int, reason: str = "", duration: Optional[int] = None + self, codeid: int, reason: str = "", duration: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None ) -> dict: """ Activate a pending Darktrace RESPOND action. @@ -88,6 +91,8 @@ def activate_action( reason (str, optional): Free text field to specify the action purpose. Required if "Audit Antigena" setting is enabled on the Darktrace System Config page. duration (int, optional): Specify how long the action should be active for in seconds. + timeout (float or tuple, optional): Timeout for the request in seconds. Can be a single + float for both connect and read timeouts, or a tuple of (connect_timeout, read_timeout). Returns: dict: API response containing activation result @@ -114,15 +119,16 @@ def activate_action( # Send JSON as raw data with consistent formatting (same as signature generation) json_data = json.dumps(body, separators=(",", ":")) + resolved_timeout = self._resolve_timeout(timeout) response = requests.post( - url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl + url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl, timeout=resolved_timeout ) self.client._debug(f"Response Status: {response.status_code}") self.client._debug(f"Response Text: {response.text}") response.raise_for_status() return response.json() - def extend_action(self, codeid: int, duration: int, reason: str = "") -> dict: + def extend_action(self, codeid: int, duration: int, reason: str = "", timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: """ Extend an active Darktrace RESPOND action. @@ -137,6 +143,8 @@ def extend_action(self, codeid: int, duration: int, reason: str = "") -> dict: duration (int): New total duration for the action in seconds. This should be the current duration plus the amount the action should be extended for. reason (str, optional): Free text field to specify the extension purpose. + timeout (float or tuple, optional): Timeout for the request in seconds. Can be a single + float for both connect and read timeouts, or a tuple of (connect_timeout, read_timeout). Returns: dict: API response containing extension result @@ -162,15 +170,16 @@ def extend_action(self, codeid: int, duration: int, reason: str = "") -> dict: # Send JSON as raw data with consistent formatting (same as signature generation) json_data = json.dumps(body, separators=(",", ":")) + resolved_timeout = self._resolve_timeout(timeout) response = requests.post( - url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl + url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl, timeout=resolved_timeout ) self.client._debug(f"Response Status: {response.status_code}") self.client._debug(f"Response Text: {response.text}") response.raise_for_status() return response.json() - def clear_action(self, codeid: int, reason: str = "") -> dict: + def clear_action(self, codeid: int, reason: str = "", timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: """ Clear an active, pending or expired Darktrace RESPOND action. @@ -185,6 +194,8 @@ def clear_action(self, codeid: int, reason: str = "") -> dict: Args: codeid (int): Unique numeric identifier of a RESPOND action. reason (str, optional): Free text field to specify the clearing purpose. + timeout (float or tuple, optional): Timeout for the request in seconds. Can be a single + float for both connect and read timeouts, or a tuple of (connect_timeout, read_timeout). Returns: bool: True if clearing was successful, False otherwise. @@ -206,15 +217,16 @@ def clear_action(self, codeid: int, reason: str = "") -> dict: # Send JSON as raw data with consistent formatting (same as signature generation) json_data = json.dumps(body, separators=(",", ":")) + resolved_timeout = self._resolve_timeout(timeout) response = requests.post( - url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl + url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl, timeout=resolved_timeout ) self.client._debug(f"Response Status: {response.status_code}") self.client._debug(f"Response Text: {response.text}") response.raise_for_status() return response.json() - def reactivate_action(self, codeid: int, duration: int, reason: str = "") -> dict: + def reactivate_action(self, codeid: int, duration: int, reason: str = "", timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: """ Reactivate a cleared or expired Darktrace RESPOND action. @@ -224,6 +236,8 @@ def reactivate_action(self, codeid: int, duration: int, reason: str = "") -> dic codeid (int): Unique numeric identifier of a RESPOND action. duration (int): Duration for the reactivated action in seconds. Required. reason (str, optional): Free text field to specify the reactivation purpose. + timeout (float or tuple, optional): Timeout for the request in seconds. Can be a single + float for both connect and read timeouts, or a tuple of (connect_timeout, read_timeout). Returns: dict: API response containing reactivation result @@ -249,8 +263,9 @@ def reactivate_action(self, codeid: int, duration: int, reason: str = "") -> dic # Send JSON as raw data with consistent formatting (same as signature generation) json_data = json.dumps(body, separators=(",", ":")) + resolved_timeout = self._resolve_timeout(timeout) response = requests.post( - url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl + url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl, timeout=resolved_timeout ) self.client._debug(f"Response Status: {response.status_code}") self.client._debug(f"Response Text: {response.text}") @@ -264,6 +279,7 @@ def create_manual_action( duration: int, reason: str = "", connections: Optional[List[Dict[str, Union[str, int]]]] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = None, ) -> dict: """ Create a manual Darktrace RESPOND/Network action. @@ -287,6 +303,8 @@ def create_manual_action( - 'src' (str): IP or hostname of source endpoint - 'dst' (str): IP or hostname of destination endpoint - 'port' (int, optional): Port for destination value + timeout (float or tuple, optional): Timeout for the request in seconds. Can be a single + float for both connect and read timeouts, or a tuple of (connect_timeout, read_timeout). Returns: int: The codeid (unique numeric ID) for the created action, or 0 if creation failed. @@ -336,8 +354,9 @@ def create_manual_action( # Send JSON as raw data with consistent formatting (same as signature generation) json_data = json.dumps(body, separators=(",", ":")) + resolved_timeout = self._resolve_timeout(timeout) response = requests.post( - url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl + url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl, timeout=resolved_timeout ) self.client._debug(f"Response Status: {response.status_code}") self.client._debug(f"Response Text: {response.text}") @@ -345,7 +364,7 @@ def create_manual_action( response.raise_for_status() return response.json() - def get_summary(self, **params): + def get_summary(self, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): """ Get a summary of active and pending Darktrace RESPOND actions. @@ -355,6 +374,8 @@ def get_summary(self, **params): will return information about active actions during that time window. Args: + timeout (float or tuple, optional): Timeout for the request in seconds. Can be a single + float for both connect and read timeouts, or a tuple of (connect_timeout, read_timeout). endtime (int): End time of data to return in millisecond format, relative to midnight January 1st 1970 UTC. starttime (int): Start time of data to return in millisecond format, relative to @@ -403,8 +424,9 @@ def get_summary(self, **params): url = f"{self.client.host}{endpoint}" self.client._debug(f"GET {url} params={sorted_params}") + resolved_timeout = self._resolve_timeout(timeout) response = requests.get( - url, headers=headers, params=sorted_params, verify=self.client.verify_ssl + url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout ) response.raise_for_status() return response.json() diff --git a/darktrace/dt_breaches.py b/darktrace/dt_breaches.py index add4e46..72463ae 100644 --- a/darktrace/dt_breaches.py +++ b/darktrace/dt_breaches.py @@ -1,6 +1,6 @@ import requests import json -from typing import Dict, Any, Optional, Union +from typing import Dict, Any, Optional, Union, Tuple from datetime import datetime from .dt_utils import debug_print, BaseEndpoint @@ -8,7 +8,7 @@ class ModelBreaches(BaseEndpoint): def __init__(self, client): super().__init__(client) - def get(self, **params): + def get(self, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): """ Get model breach alerts from the /modelbreaches endpoint. @@ -64,16 +64,18 @@ def get(self, **params): headers, sorted_params = self._get_headers(endpoint, dict(params_list)) self.client._debug(f"GET {url} params={sorted_params}") + resolved_timeout = self._resolve_timeout(timeout) response = requests.get( url, headers=headers, params=sorted_params, - verify=self.client.verify_ssl + verify=self.client.verify_ssl, + timeout=resolved_timeout ) response.raise_for_status() return response.json() - def get_comments(self, pbid: Union[int, list], **params): + def get_comments(self, pbid: Union[int, list], timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): """ Get comments for a specific model breach alert. @@ -85,16 +87,17 @@ def get_comments(self, pbid: Union[int, list], **params): """ if isinstance(pbid, (list, tuple)): # Build dict with string keys for valid JSON - return {str(single_pbid): self.get_comments(single_pbid, **params) for single_pbid in pbid} + return {str(single_pbid): self.get_comments(single_pbid, timeout=timeout, **params) for single_pbid in pbid} endpoint = f'/modelbreaches/{pbid}/comments' url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def add_comment(self, pbid: int, message: str, **params) -> dict: + def add_comment(self, pbid: int, message: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params) -> dict: """ Add a comment to a model breach alert. @@ -126,7 +129,7 @@ def add_comment(self, pbid: int, message: str, **params) -> dict: debug_print(f" - sorted_params: {sorted_params}", self.client.debug) self.client._debug(f"POST {url} params={sorted_params} body={body}") - + try: # Send JSON as raw data, not as json parameter (as per Darktrace docs) # IMPORTANT: Must use same JSON formatting as in signature generation! @@ -136,8 +139,9 @@ def add_comment(self, pbid: int, message: str, **params) -> dict: debug_print(f"BREACHES: With headers: {headers}", self.client.debug) debug_print(f"BREACHES: With params: {sorted_params}", self.client.debug) debug_print(f"BREACHES: With data: '{json_data}'", self.client.debug) - - response = requests.post(url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl) + + resolved_timeout = self._resolve_timeout(timeout) + response = requests.post(url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl, timeout=resolved_timeout) self.client._debug(f"Response Status: {response.status_code}") self.client._debug(f"Response Text: {response.text}") debug_print(f"BREACHES: Response status: {response.status_code}", self.client.debug) @@ -150,7 +154,7 @@ def add_comment(self, pbid: int, message: str, **params) -> dict: debug_print(f"BREACHES: Exception: {str(e)}", self.client.debug) return {"error": str(e)} - def acknowledge(self, pbid: Union[int, list], **params) -> dict: + def acknowledge(self, pbid: Union[int, list], timeout: Optional[Union[float, Tuple[float, float]]] = None, **params) -> dict: """ Acknowledge a model breach alert. @@ -161,7 +165,7 @@ def acknowledge(self, pbid: Union[int, list], **params) -> dict: dict: The full JSON response from Darktrace (or error info as dict), or a dict mapping pbid to response if pbid is a list. """ if isinstance(pbid, (list, tuple)): - return {single_pbid: self.acknowledge(single_pbid, **params) for single_pbid in pbid} + return {single_pbid: self.acknowledge(single_pbid, timeout=timeout, **params) for single_pbid in pbid} endpoint = f'/modelbreaches/{pbid}/acknowledge' url = f"{self.client.host}{endpoint}" body: Dict[str, bool] = {'acknowledge': True} @@ -171,7 +175,8 @@ def acknowledge(self, pbid: Union[int, list], **params) -> dict: # Send JSON as raw data, not as json parameter (as per Darktrace docs) # IMPORTANT: Must use same JSON formatting as in signature generation! json_data = json.dumps(body, separators=(',', ':')) - response = requests.post(url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.post(url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl, timeout=resolved_timeout) self.client._debug(f"Response Status: {response.status_code}") self.client._debug(f"Response Text: {response.text}") response.raise_for_status() @@ -180,7 +185,7 @@ def acknowledge(self, pbid: Union[int, list], **params) -> dict: self.client._debug(f"Exception occurred while acknowledging breach: {str(e)}") return {"error": str(e)} - def unacknowledge(self, pbid: Union[int, list], **params) -> dict: + def unacknowledge(self, pbid: Union[int, list], timeout: Optional[Union[float, Tuple[float, float]]] = None, **params) -> dict: """ Unacknowledge a model breach alert. @@ -191,7 +196,7 @@ def unacknowledge(self, pbid: Union[int, list], **params) -> dict: dict: The full JSON response from Darktrace (or error info as dict), or a dict mapping pbid to response if pbid is a list. """ if isinstance(pbid, (list, tuple)): - return {single_pbid: self.unacknowledge(single_pbid, **params) for single_pbid in pbid} + return {single_pbid: self.unacknowledge(single_pbid, timeout=timeout, **params) for single_pbid in pbid} endpoint = f'/modelbreaches/{pbid}/unacknowledge' url = f"{self.client.host}{endpoint}" body: Dict[str, bool] = {'unacknowledge': True} @@ -201,7 +206,8 @@ def unacknowledge(self, pbid: Union[int, list], **params) -> dict: # Send JSON as raw data, not as json parameter (as per Darktrace docs) # IMPORTANT: Must use same JSON formatting as in signature generation! json_data = json.dumps(body, separators=(',', ':')) - response = requests.post(url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.post(url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl, timeout=resolved_timeout) self.client._debug(f"Response Status: {response.status_code}") self.client._debug(f"Response Text: {response.text}") response.raise_for_status() @@ -210,7 +216,7 @@ def unacknowledge(self, pbid: Union[int, list], **params) -> dict: self.client._debug(f"Exception occurred while unacknowledging breach: {str(e)}") return {"error": str(e)} - def acknowledge_with_comment(self, pbid: int, message: str, **params) -> dict: + def acknowledge_with_comment(self, pbid: int, message: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params) -> dict: """ Acknowledge a model breach and add a comment in one call. @@ -222,14 +228,14 @@ def acknowledge_with_comment(self, pbid: int, message: str, **params) -> dict: Returns: dict: Contains the responses from both acknowledge and add_comment. """ - ack_response = self.acknowledge(pbid, **params) - comment_response = self.add_comment(pbid, message, **params) + ack_response = self.acknowledge(pbid, timeout=timeout, **params) + comment_response = self.add_comment(pbid, message, timeout=timeout, **params) return { "acknowledge": ack_response, "add_comment": comment_response } - def unacknowledge_with_comment(self, pbid: int, message: str, **params) -> dict: + def unacknowledge_with_comment(self, pbid: int, message: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params) -> dict: """ Unacknowledge a model breach and add a comment in one call. @@ -241,8 +247,8 @@ def unacknowledge_with_comment(self, pbid: int, message: str, **params) -> dict: Returns: dict: Contains the responses from both unacknowledge and add_comment. """ - unack_response = self.unacknowledge(pbid, **params) - comment_response = self.add_comment(pbid, message, **params) + unack_response = self.unacknowledge(pbid, timeout=timeout, **params) + comment_response = self.add_comment(pbid, message, timeout=timeout, **params) return { "unacknowledge": unack_response, "add_comment": comment_response diff --git a/darktrace/dt_components.py b/darktrace/dt_components.py index 1d8cfb6..9d98bcd 100644 --- a/darktrace/dt_components.py +++ b/darktrace/dt_components.py @@ -1,12 +1,12 @@ import requests -from typing import Optional +from typing import Optional, Union, Tuple from .dt_utils import debug_print, BaseEndpoint class Components(BaseEndpoint): def __init__(self, client): super().__init__(client) - def get(self, cid: Optional[int] = None, responsedata: Optional[str] = None, **params): + def get(self, cid: Optional[int] = None, responsedata: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): """ Get information about model components. @@ -30,6 +30,7 @@ def get(self, cid: Optional[int] = None, responsedata: Optional[str] = None, **p url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() \ No newline at end of file diff --git a/darktrace/dt_cves.py b/darktrace/dt_cves.py index 6ad9692..368a690 100644 --- a/darktrace/dt_cves.py +++ b/darktrace/dt_cves.py @@ -1,5 +1,5 @@ import requests -from typing import Optional +from typing import Optional, Union, Tuple from .dt_utils import debug_print, BaseEndpoint class CVEs(BaseEndpoint): @@ -10,6 +10,7 @@ def get( self, did: Optional[int] = None, fulldevicedetails: Optional[bool] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = None, **params ): """ @@ -37,6 +38,7 @@ def get( # Use consistent parameter/header handling headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() \ No newline at end of file diff --git a/darktrace/dt_details.py b/darktrace/dt_details.py index 9a58998..57d1f9f 100644 --- a/darktrace/dt_details.py +++ b/darktrace/dt_details.py @@ -1,5 +1,5 @@ import requests -from typing import Optional +from typing import Optional, Union, Tuple from .dt_utils import debug_print, BaseEndpoint class Details(BaseEndpoint): @@ -31,6 +31,7 @@ def get( deduplicate: bool = False, fulldevicedetails: bool = False, responsedata: Optional[str] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = None, **params ): """ @@ -134,7 +135,8 @@ def get( params['responsedata'] = responsedata headers, sorted_params = self._get_headers(endpoint, params) + resolved_timeout = self._resolve_timeout(timeout) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() \ No newline at end of file diff --git a/darktrace/dt_deviceinfo.py b/darktrace/dt_deviceinfo.py index 91e8953..68f73b5 100644 --- a/darktrace/dt_deviceinfo.py +++ b/darktrace/dt_deviceinfo.py @@ -1,5 +1,5 @@ import requests -from typing import Optional +from typing import Optional, Union, Tuple from .dt_utils import debug_print, BaseEndpoint class DeviceInfo(BaseEndpoint): @@ -17,6 +17,7 @@ def get( showallgraphdata: bool = True, similardevices: Optional[int] = None, intervalhours: int = 1, + timeout: Optional[Union[float, Tuple[float, float]]] = None, **params ): """ @@ -69,7 +70,8 @@ def get( url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint, params) + resolved_timeout = self._resolve_timeout(timeout) self.client._debug(f"GET {url} params={params}") - response = requests.get(url, headers=headers, params=sorted_params or params, verify=self.client.verify_ssl) + response = requests.get(url, headers=headers, params=sorted_params or params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() \ No newline at end of file diff --git a/darktrace/dt_devices.py b/darktrace/dt_devices.py index 780bb71..cb9e5e3 100644 --- a/darktrace/dt_devices.py +++ b/darktrace/dt_devices.py @@ -1,6 +1,6 @@ import requests import json -from typing import List, Dict, Any +from typing import List, Dict, Any, Optional, Union, Tuple from .dt_utils import debug_print, BaseEndpoint @@ -21,6 +21,7 @@ def get( responsedata: str = None, cloudsecurity: bool = None, saasfilter: Any = None, + timeout: Optional[Union[float, Tuple[float, float]]] = None, ): """ Update a single device. @@ -79,13 +80,14 @@ def get( headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={sorted_params}") + resolved_timeout = self._resolve_timeout(timeout) response = requests.get( - url, headers=headers, params=sorted_params, verify=self.client.verify_ssl + url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout ) response.raise_for_status() return response.json() - def update(self, did: int, **kwargs) -> dict: + def update(self, did: int, timeout: Optional[Union[float, Tuple[float, float]]] = None, **kwargs) -> dict: """Update device properties in Darktrace. Args: @@ -108,8 +110,9 @@ def update(self, did: int, **kwargs) -> dict: # Send JSON as raw data with consistent formatting (same as signature generation) json_data = json.dumps(body, separators=(",", ":")) + resolved_timeout = self._resolve_timeout(timeout) response = requests.post( - url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl + url, headers=headers, params=sorted_params, data=json_data, verify=self.client.verify_ssl, timeout=resolved_timeout ) self.client._debug(f"Response Status: {response.status_code}") self.client._debug(f"Response Text: {response.text}") diff --git a/darktrace/dt_devicesearch.py b/darktrace/dt_devicesearch.py index 6ac3032..09d75cc 100644 --- a/darktrace/dt_devicesearch.py +++ b/darktrace/dt_devicesearch.py @@ -1,4 +1,5 @@ import requests +from typing import Optional, Union, Tuple from .dt_utils import debug_print, BaseEndpoint @@ -36,6 +37,7 @@ def get( offset=None, responsedata=None, seensince=None, + timeout: Optional[Union[float, Tuple[float, float]]] = None, **kwargs, ): """ @@ -114,106 +116,114 @@ def get( headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={params}") + resolved_timeout = self._resolve_timeout(timeout) response = requests.get( - url, headers=headers, params=sorted_params, verify=self.client.verify_ssl + url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout ) response.raise_for_status() return response.json() - def get_tag(self, tag: str, **kwargs): + def get_tag(self, tag: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **kwargs): """ Search for devices with a specific tag. Args: tag (str): The tag to search for. + timeout (Optional[Union[float, Tuple[float, float]]]): Request timeout in seconds. **kwargs: Additional parameters for the search. Returns: dict: API response """ query = f'tag:"{tag}"' - return self.get(query=query, **kwargs) + return self.get(query=query, timeout=timeout, **kwargs) - def get_type(self, type: str, **kwargs): + def get_type(self, type: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **kwargs): """ Search for devices with a specific type. Args: type (str): The type to search for. + timeout (Optional[Union[float, Tuple[float, float]]]): Request timeout in seconds. **kwargs: Additional parameters for the search. Returns: dict: API response """ query = f'type:"{type}"' - return self.get(query=query, **kwargs) + return self.get(query=query, timeout=timeout, **kwargs) - def get_label(self, label: str, **kwargs): + def get_label(self, label: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **kwargs): """ Search for devices with a specific label. Args: label (str): The label to search for. + timeout (Optional[Union[float, Tuple[float, float]]]): Request timeout in seconds. **kwargs: Additional parameters for the search. Returns: dict: API response """ query = f'label:"{label}"' - return self.get(query=query, **kwargs) + return self.get(query=query, timeout=timeout, **kwargs) - def get_vendor(self, vendor: str, **kwargs): + def get_vendor(self, vendor: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **kwargs): """ Search for devices with a specific vendor. Args: vendor (str): The vendor to search for. + timeout (Optional[Union[float, Tuple[float, float]]]): Request timeout in seconds. **kwargs: Additional parameters for the search. Returns: dict: API response """ query = f'vendor:"{vendor}"' - return self.get(query=query, **kwargs) + return self.get(query=query, timeout=timeout, **kwargs) - def get_hostname(self, hostname: str, **kwargs): + def get_hostname(self, hostname: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **kwargs): """ Search for devices with a specific hostname. Args: hostname (str): The hostname to search for. + timeout (Optional[Union[float, Tuple[float, float]]]): Request timeout in seconds. **kwargs: Additional parameters for the search. Returns: dict: API response """ query = f'hostname:"{hostname}"' - return self.get(query=query, **kwargs) + return self.get(query=query, timeout=timeout, **kwargs) - def get_ip(self, ip: str, **kwargs): + def get_ip(self, ip: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **kwargs): """ Search for devices with a specific IP address. Args: ip (str): The IP address to search for. + timeout (Optional[Union[float, Tuple[float, float]]]): Request timeout in seconds. **kwargs: Additional parameters for the search. Returns: dict: API response """ query = f'ip:"{ip}"' - return self.get(query=query, **kwargs) + return self.get(query=query, timeout=timeout, **kwargs) - def get_mac(self, mac: str, **kwargs): + def get_mac(self, mac: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **kwargs): """ Search for devices with a specific MAC address. Args: mac (str): The MAC address to search for. + timeout (Optional[Union[float, Tuple[float, float]]]): Request timeout in seconds. **kwargs: Additional parameters for the search. Returns: dict: API response """ query = f'mac:"{mac}"' - return self.get(query=query, **kwargs) + return self.get(query=query, timeout=timeout, **kwargs) diff --git a/darktrace/dt_devicesummary.py b/darktrace/dt_devicesummary.py index c729fa0..869bccf 100644 --- a/darktrace/dt_devicesummary.py +++ b/darktrace/dt_devicesummary.py @@ -1,5 +1,5 @@ import requests -from typing import Optional, Dict, Any +from typing import Optional, Union, Tuple, Dict, Any from .dt_utils import debug_print, BaseEndpoint class DeviceSummary(BaseEndpoint): @@ -44,6 +44,7 @@ def get( source: Optional[str] = None, status: Optional[str] = None, responsedata: Optional[str] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = None, **kwargs ) -> Dict[str, Any]: """ @@ -100,7 +101,8 @@ def get( params['responsedata'] = responsedata params.update(kwargs) headers, sorted_params = self._get_headers(endpoint, params) + resolved_timeout = self._resolve_timeout(timeout) self.client._debug(f"GET {url} params={params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() \ No newline at end of file diff --git a/darktrace/dt_email.py b/darktrace/dt_email.py index e1db377..cb766c8 100644 --- a/darktrace/dt_email.py +++ b/darktrace/dt_email.py @@ -1,18 +1,19 @@ import requests import json -from typing import Dict, Any, Optional, Union, List +from typing import Dict, Any, Optional, Union, List, Tuple from .dt_utils import debug_print, BaseEndpoint class DarktraceEmail(BaseEndpoint): def __init__(self, client): super().__init__(client) - def decode_link(self, link: str) -> Dict[str, Any]: + def decode_link(self, link: str, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: """ Decode a link using the Darktrace/Email API. Args: link (str): The encoded link to decode. + timeout (float, tuple[float, float], optional): Request timeout in seconds. Returns: dict: Decoded link information. @@ -24,17 +25,19 @@ def decode_link(self, link: str) -> Dict[str, Any]: params = {"link": link} headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def get_action_summary(self, days: Optional[int] = None, limit: Optional[int] = None) -> Dict[str, Any]: + def get_action_summary(self, days: Optional[int] = None, limit: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: """ Get action summary from Darktrace/Email API. Args: days (int, optional): Number of days to include in the summary. limit (int, optional): Limit the number of results. + timeout (float, tuple[float, float], optional): Request timeout in seconds. Returns: dict: Action summary data. @@ -50,17 +53,19 @@ def get_action_summary(self, days: Optional[int] = None, limit: Optional[int] = params["limit"] = limit headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def get_dash_stats(self, days: Optional[int] = None, limit: Optional[int] = None) -> Dict[str, Any]: + def get_dash_stats(self, days: Optional[int] = None, limit: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: """ Get dashboard stats from Darktrace/Email API. Args: days (int, optional): Number of days to include in the stats. limit (int, optional): Limit the number of results. + timeout (float, tuple[float, float], optional): Request timeout in seconds. Returns: dict: Dashboard statistics. @@ -76,17 +81,19 @@ def get_dash_stats(self, days: Optional[int] = None, limit: Optional[int] = None params["limit"] = limit headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def get_data_loss(self, days: Optional[int] = None, limit: Optional[int] = None) -> Dict[str, Any]: + def get_data_loss(self, days: Optional[int] = None, limit: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: """ Get data loss information from Darktrace/Email API. Args: days (int, optional): Number of days to include in the data loss stats. limit (int, optional): Limit the number of results. + timeout (float, tuple[float, float], optional): Request timeout in seconds. Returns: dict: Data loss information. @@ -102,17 +109,19 @@ def get_data_loss(self, days: Optional[int] = None, limit: Optional[int] = None) params["limit"] = limit headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def get_user_anomaly(self, days: Optional[int] = None, limit: Optional[int] = None) -> Dict[str, Any]: + def get_user_anomaly(self, days: Optional[int] = None, limit: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: """ Get user anomaly data from Darktrace/Email API. Args: days (int, optional): Number of days to include in the anomaly stats. limit (int, optional): Limit the number of results. + timeout (float, tuple[float, float], optional): Request timeout in seconds. Returns: dict: User anomaly data. @@ -128,30 +137,33 @@ def get_user_anomaly(self, days: Optional[int] = None, limit: Optional[int] = No params["limit"] = limit headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def email_action(self, uuid: str, data: Dict[str, Any]): + def email_action(self, uuid: str, data: Dict[str, Any], timeout: Optional[Union[float, Tuple[float, float]]] = None): """Perform an action on an email by UUID in Darktrace/Email API.""" endpoint = f'/agemail/api/ep/api/v1.0/emails/{uuid}/action' url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint, json_body=data) headers['Content-Type'] = 'application/json' self.client._debug(f"POST {url} data={data}") - response = requests.post(url, headers=headers, data=json.dumps(data, separators=(',', ':')), verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.post(url, headers=headers, data=json.dumps(data, separators=(',', ':')), verify=self.client.verify_ssl, timeout=resolved_timeout) self.client._debug(f"Response status: {response.status_code}") self.client._debug(f"Response text: {response.text}") response.raise_for_status() return response.json() - def get_email(self, uuid: str, include_headers: Optional[bool] = None) -> Dict[str, Any]: + def get_email(self, uuid: str, include_headers: Optional[bool] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: """ Get email details by UUID from Darktrace/Email API. Args: uuid (str): Email UUID. include_headers (bool, optional): Whether to include email headers in the response. + timeout (float, tuple[float, float], optional): Request timeout in seconds. Returns: dict: Email details. @@ -165,16 +177,18 @@ def get_email(self, uuid: str, include_headers: Optional[bool] = None) -> Dict[s params["include_headers"] = include_headers headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def download_email(self, uuid: str) -> bytes: + def download_email(self, uuid: str, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> bytes: """ Download an email by UUID from Darktrace/Email API. Args: uuid (str): Email UUID. + timeout (float, tuple[float, float], optional): Request timeout in seconds. Returns: bytes: Raw email content (MIME). @@ -185,27 +199,32 @@ def download_email(self, uuid: str) -> bytes: url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint) self.client._debug(f"GET {url} params={{}}") - response = requests.get(url, headers=headers, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.content - def search_emails(self, data: Dict[str, Any]): + def search_emails(self, data: Dict[str, Any], timeout: Optional[Union[float, Tuple[float, float]]] = None): """Search emails in Darktrace/Email API.""" endpoint = '/agemail/api/ep/api/v1.0/emails/search' url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint, json_body=data) headers['Content-Type'] = 'application/json' self.client._debug(f"POST {url} data={data}") - response = requests.post(url, headers=headers, data=json.dumps(data, separators=(',', ':')), verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.post(url, headers=headers, data=json.dumps(data, separators=(',', ':')), verify=self.client.verify_ssl, timeout=resolved_timeout) self.client._debug(f"Response status: {response.status_code}") self.client._debug(f"Response text: {response.text}") response.raise_for_status() return response.json() - def get_tags(self) -> Dict[str, Any]: + def get_tags(self, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: """ Get tags from Darktrace/Email API. + Args: + timeout (float, tuple[float, float], optional): Request timeout in seconds. + Returns: dict: Tags data. Example: @@ -215,14 +234,18 @@ def get_tags(self) -> Dict[str, Any]: url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint) self.client._debug(f"GET {url} params={{}}") - response = requests.get(url, headers=headers, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def get_actions(self) -> Dict[str, Any]: + def get_actions(self, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: """ Get actions from Darktrace/Email API. + Args: + timeout (float, tuple[float, float], optional): Request timeout in seconds. + Returns: dict: Actions data. Example: @@ -232,14 +255,18 @@ def get_actions(self) -> Dict[str, Any]: url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint) self.client._debug(f"GET {url} params={{}}") - response = requests.get(url, headers=headers, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def get_filters(self) -> Dict[str, Any]: + def get_filters(self, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: """ Get filters from Darktrace/Email API. + Args: + timeout (float, tuple[float, float], optional): Request timeout in seconds. + Returns: dict: Filters data. Example: @@ -249,14 +276,18 @@ def get_filters(self) -> Dict[str, Any]: url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint) self.client._debug(f"GET {url} params={{}}") - response = requests.get(url, headers=headers, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def get_event_types(self) -> Dict[str, Any]: + def get_event_types(self, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: """ Get audit event types from Darktrace/Email API. + Args: + timeout (float, tuple[float, float], optional): Request timeout in seconds. + Returns: dict: Audit event types. Example: @@ -266,11 +297,12 @@ def get_event_types(self) -> Dict[str, Any]: url = f"{self.client.host}{endpoint}" headers, sorted_params = self._get_headers(endpoint) self.client._debug(f"GET {url} params={{}}") - response = requests.get(url, headers=headers, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def get_audit_events(self, event_type: Optional[str] = None, limit: Optional[int] = None, offset: Optional[int] = None) -> Dict[str, Any]: + def get_audit_events(self, event_type: Optional[str] = None, limit: Optional[int] = None, offset: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: """ Get audit events from Darktrace/Email API. @@ -278,6 +310,7 @@ def get_audit_events(self, event_type: Optional[str] = None, limit: Optional[int event_type (str, optional): Filter by event type. limit (int, optional): Limit the number of results. offset (int, optional): Offset for pagination. + timeout (float, tuple[float, float], optional): Request timeout in seconds. Returns: dict: Audit events data. @@ -295,6 +328,7 @@ def get_audit_events(self, event_type: Optional[str] = None, limit: Optional[int params["offset"] = offset headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() \ No newline at end of file diff --git a/darktrace/dt_endpointdetails.py b/darktrace/dt_endpointdetails.py index ec8ef0f..e7cedcd 100644 --- a/darktrace/dt_endpointdetails.py +++ b/darktrace/dt_endpointdetails.py @@ -1,5 +1,5 @@ import requests -from typing import Optional, Any, Dict +from typing import Optional, Any, Dict, Union, Tuple from .dt_utils import debug_print, BaseEndpoint class EndpointDetails(BaseEndpoint): @@ -12,7 +12,8 @@ def get(self, additionalinfo: Optional[bool] = None, devices: Optional[bool] = None, score: Optional[bool] = None, - responsedata: Optional[str] = None + responsedata: Optional[str] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = None ) -> Dict[str, Any]: """ Get endpoint details from Darktrace. @@ -46,6 +47,7 @@ def get(self, headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() \ No newline at end of file diff --git a/darktrace/dt_enums.py b/darktrace/dt_enums.py index fbc1711..21bbc41 100644 --- a/darktrace/dt_enums.py +++ b/darktrace/dt_enums.py @@ -1,5 +1,5 @@ import requests -from typing import Optional +from typing import Optional, Union, Tuple from .dt_utils import debug_print, BaseEndpoint class Enums(BaseEndpoint): @@ -11,7 +11,7 @@ class Enums(BaseEndpoint): def __init__(self, client): super().__init__(client) - def get(self, responsedata: Optional[str] = None, **params): + def get(self, responsedata: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): """ Get enum values for all types or restrict to a specific field/object. @@ -31,6 +31,7 @@ def get(self, responsedata: Optional[str] = None, **params): query_params.update(params) headers, sorted_params = self._get_headers(endpoint, query_params) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() \ No newline at end of file diff --git a/darktrace/dt_filtertypes.py b/darktrace/dt_filtertypes.py index ee3352d..94dac50 100644 --- a/darktrace/dt_filtertypes.py +++ b/darktrace/dt_filtertypes.py @@ -1,5 +1,5 @@ import requests -from typing import Optional +from typing import Optional, Union, Tuple from .dt_utils import debug_print, BaseEndpoint class FilterTypes(BaseEndpoint): @@ -22,7 +22,7 @@ class FilterTypes(BaseEndpoint): def __init__(self, client): super().__init__(client) - def get(self, responsedata: Optional[str] = None, **params): + def get(self, responsedata: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): """ Get all filter types or restrict to a specific field/object. @@ -41,6 +41,7 @@ def get(self, responsedata: Optional[str] = None, **params): query_params.update(params) headers, sorted_params = self._get_headers(endpoint, query_params) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() \ No newline at end of file diff --git a/darktrace/dt_intelfeed.py b/darktrace/dt_intelfeed.py index 24450af..ef6cfbe 100644 --- a/darktrace/dt_intelfeed.py +++ b/darktrace/dt_intelfeed.py @@ -1,6 +1,6 @@ import requests import json -from typing import Optional, List, Dict, Any, Union +from typing import Optional, List, Dict, Any, Union, Tuple from .dt_utils import debug_print, BaseEndpoint @@ -26,7 +26,8 @@ def __init__(self, client): super().__init__(client) def get(self, sources: Optional[bool] = None, source: Optional[str] = None, - fulldetails: Optional[bool] = None, responsedata: Optional[str] = None, **params): + fulldetails: Optional[bool] = None, responsedata: Optional[str] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): """ Get the intelfeed list, sources, or detailed entries. @@ -53,8 +54,9 @@ def get(self, sources: Optional[bool] = None, source: Optional[str] = None, query_params['responsedata'] = responsedata query_params.update(params) headers, sorted_params = self._get_headers(endpoint, query_params) + resolved_timeout = self._resolve_timeout(timeout) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() @@ -70,11 +72,12 @@ def get_with_details(self): """Get intel feed with full details about expiry time and description for each entry.""" return self.get(full_details=True) - def update(self, add_entry: Optional[str] = None, add_list: Optional[List[str]] = None, - description: Optional[str] = None, source: Optional[str] = None, - expiry: Optional[str] = None, is_hostname: bool = False, - remove_entry: Optional[str] = None, remove_all: bool = False, - enable_antigena: bool = False): + def update(self, add_entry: Optional[str] = None, add_list: Optional[List[str]] = None, + description: Optional[str] = None, source: Optional[str] = None, + expiry: Optional[str] = None, is_hostname: bool = False, + remove_entry: Optional[str] = None, remove_all: bool = False, + enable_antigena: bool = False, + timeout: Optional[Union[float, Tuple[float, float]]] = None): """Update the intel feed (watched domains) in Darktrace. Args: @@ -116,9 +119,10 @@ def update(self, add_entry: Optional[str] = None, add_list: Optional[List[str]] # For POST requests with JSON body, we need to include the body in the signature headers, _ = self._get_headers(endpoint, json_body=body) headers['Content-Type'] = 'application/json' - + + resolved_timeout = self._resolve_timeout(timeout) self.client._debug(f"POST {url} body={body}") - response = requests.post(url, headers=headers, data=json.dumps(body, separators=(',', ':')), verify=self.client.verify_ssl) + response = requests.post(url, headers=headers, data=json.dumps(body, separators=(',', ':')), verify=self.client.verify_ssl, timeout=resolved_timeout) self.client._debug(f"Response status: {response.status_code}") self.client._debug(f"Response text: {response.text}") response.raise_for_status() diff --git a/darktrace/dt_mbcomments.py b/darktrace/dt_mbcomments.py index 9392121..55b3926 100644 --- a/darktrace/dt_mbcomments.py +++ b/darktrace/dt_mbcomments.py @@ -1,6 +1,6 @@ import requests import json -from typing import Optional, Dict, Any +from typing import Optional, Dict, Any, Union, Tuple from .dt_utils import debug_print, BaseEndpoint class MBComments(BaseEndpoint): @@ -14,6 +14,7 @@ def get(self, responsedata: Optional[str] = None, count: Optional[int] = None, pbid: Optional[int] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = None, **params ): """ @@ -26,6 +27,7 @@ def get(self, responsedata (str, optional): Restrict the returned JSON to only the specified field/object. count (int, optional): Number of comments to return (default 100). pbid (int, optional): Only return comments for the model breach with this ID. + timeout (float or tuple, optional): Timeout for the request in seconds. **params: Additional query parameters. Returns: @@ -46,21 +48,30 @@ def get(self, query_params['pbid'] = pbid query_params.update(params) headers, sorted_params = self._get_headers(endpoint, query_params) + resolved_timeout = self._resolve_timeout(timeout) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def post(self, breach_id: str, comment: str, **params): - """Add a comment to a model breach.""" + def post(self, breach_id: str, comment: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): + """Add a comment to a model breach. + + Args: + breach_id (str): Model breach ID. + comment (str): Comment text to add. + timeout (float or tuple, optional): Timeout for the request in seconds. + **params: Additional parameters. + """ endpoint = '/mbcomments' url = f"{self.client.host}{endpoint}" data: Dict[str, Any] = {'breachid': breach_id, 'comment': comment} data.update(params) headers, sorted_params = self._get_headers(endpoint, json_body=data) headers['Content-Type'] = 'application/json' + resolved_timeout = self._resolve_timeout(timeout) self.client._debug(f"POST {url} data={data}") - response = requests.post(url, headers=headers, data=json.dumps(data, separators=(',', ':')), verify=self.client.verify_ssl) + response = requests.post(url, headers=headers, data=json.dumps(data, separators=(',', ':')), verify=self.client.verify_ssl, timeout=resolved_timeout) self.client._debug(f"Response status: {response.status_code}") self.client._debug(f"Response text: {response.text}") response.raise_for_status() diff --git a/darktrace/dt_metricdata.py b/darktrace/dt_metricdata.py index 6966ea8..3ea7294 100644 --- a/darktrace/dt_metricdata.py +++ b/darktrace/dt_metricdata.py @@ -1,5 +1,5 @@ import requests -from typing import Optional, List +from typing import Optional, List, Union, Tuple from .dt_utils import debug_print, BaseEndpoint class MetricData(BaseEndpoint): @@ -26,6 +26,7 @@ def get( breachtimes: Optional[bool] = None, fulldevicedetails: Optional[bool] = None, devices: Optional[List[str]] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = None, **params ): """ @@ -50,6 +51,7 @@ def get( breachtimes (bool, optional): Whether to include breach times. fulldevicedetails (bool, optional): Whether to include full device details. devices (list of str, optional): List of device IDs or names. + timeout (float or tuple, optional): Request timeout in seconds. Can be a single value or (connect_timeout, read_timeout). **params: Additional parameters for future compatibility. Returns: @@ -103,6 +105,7 @@ def get( headers, sorted_params = self._get_headers(endpoint, query_params) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() \ No newline at end of file diff --git a/darktrace/dt_metrics.py b/darktrace/dt_metrics.py index e7ce896..971488e 100644 --- a/darktrace/dt_metrics.py +++ b/darktrace/dt_metrics.py @@ -1,5 +1,5 @@ import requests -from typing import Optional +from typing import Optional, Union, Tuple from .dt_utils import debug_print, BaseEndpoint class Metrics(BaseEndpoint): @@ -10,6 +10,7 @@ def get( self, metric_id: Optional[int] = None, responsedata: Optional[str] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = None, **params ): """ @@ -37,6 +38,7 @@ def get( query_params.update(params) headers, sorted_params = self._get_headers(endpoint, query_params) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() \ No newline at end of file diff --git a/darktrace/dt_models.py b/darktrace/dt_models.py index 64569de..d2492b4 100644 --- a/darktrace/dt_models.py +++ b/darktrace/dt_models.py @@ -1,12 +1,12 @@ import requests -from typing import Optional +from typing import Optional, Union, Tuple from .dt_utils import debug_print, BaseEndpoint class Models(BaseEndpoint): def __init__(self, client): super().__init__(client) - def get(self, uuid: Optional[str] = None, responsedata: Optional[str] = None): + def get(self, uuid: Optional[str] = None, responsedata: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None): """ Get model information from Darktrace. @@ -26,6 +26,7 @@ def get(self, uuid: Optional[str] = None, responsedata: Optional[str] = None): params['responsedata'] = responsedata headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() \ No newline at end of file diff --git a/darktrace/dt_network.py b/darktrace/dt_network.py index c575ea4..7346960 100644 --- a/darktrace/dt_network.py +++ b/darktrace/dt_network.py @@ -1,5 +1,5 @@ import requests -from typing import Optional +from typing import Optional, Union, Tuple from .dt_utils import debug_print, BaseEndpoint class Network(BaseEndpoint): @@ -22,7 +22,8 @@ def get(self, starttime: Optional[int] = None, to: Optional[str] = None, viewsubnet: Optional[int] = None, - responsedata: Optional[str] = None + responsedata: Optional[str] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = None ): """ Get network connectivity and statistics information from Darktrace. @@ -85,7 +86,8 @@ def get(self, params['responsedata'] = responsedata headers, sorted_params = self._get_headers(endpoint, params) + resolved_timeout = self._resolve_timeout(timeout) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() \ No newline at end of file diff --git a/darktrace/dt_pcaps.py b/darktrace/dt_pcaps.py index f52600d..832bd01 100644 --- a/darktrace/dt_pcaps.py +++ b/darktrace/dt_pcaps.py @@ -1,12 +1,12 @@ import requests -from typing import Optional +from typing import Optional, Union, Tuple from .dt_utils import debug_print, BaseEndpoint class PCAPs(BaseEndpoint): def __init__(self, client): super().__init__(client) - def get(self, pcap_id: Optional[str] = None, responsedata: Optional[str] = None): + def get(self, pcap_id: Optional[str] = None, responsedata: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None): """ Retrieve PCAP information or download a specific PCAP file from Darktrace. @@ -24,12 +24,13 @@ def get(self, pcap_id: Optional[str] = None, responsedata: Optional[str] = None) params['responsedata'] = responsedata headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() # Return JSON if possible, else return raw content (for PCAP file download) return response.json() if 'application/json' in response.headers.get('Content-Type', '') else response.content - def create(self, ip1: str, start: int, end: int, ip2: Optional[str] = None, port1: Optional[int] = None, port2: Optional[int] = None, protocol: Optional[str] = None): + def create(self, ip1: str, start: int, end: int, ip2: Optional[str] = None, port1: Optional[int] = None, port2: Optional[int] = None, protocol: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None): """ Create a new PCAP capture request in Darktrace. @@ -58,6 +59,7 @@ def create(self, ip1: str, start: int, end: int, ip2: Optional[str] = None, port body["protocol"] = protocol headers, _ = self._get_headers(endpoint) self.client._debug(f"POST {url} body={body}") - response = requests.post(url, headers=headers, json=body, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.post(url, headers=headers, json=body, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() \ No newline at end of file diff --git a/darktrace/dt_similardevices.py b/darktrace/dt_similardevices.py index 4927e9b..6a7920a 100644 --- a/darktrace/dt_similardevices.py +++ b/darktrace/dt_similardevices.py @@ -1,5 +1,5 @@ import requests -from typing import Optional +from typing import Optional, Union, Tuple from .dt_utils import debug_print, BaseEndpoint class SimilarDevices(BaseEndpoint): @@ -13,6 +13,7 @@ def get( fulldevicedetails: Optional[bool] = None, token: Optional[str] = None, responsedata: Optional[str] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = None, **kwargs ): """ @@ -24,6 +25,7 @@ def get( fulldevicedetails (bool, optional): Whether to include full device details in the response. token (str, optional): Pagination token for large result sets. responsedata (str, optional): Restrict the returned JSON to only the specified field(s). + timeout (float or tuple, optional): Request timeout in seconds. Can be a single value or (connect_timeout, read_timeout). **kwargs: Additional parameters for future compatibility. Returns: @@ -46,7 +48,8 @@ def get( headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() try: return response.json() diff --git a/darktrace/dt_status.py b/darktrace/dt_status.py index 7cf11ae..68476d4 100644 --- a/darktrace/dt_status.py +++ b/darktrace/dt_status.py @@ -1,5 +1,5 @@ import requests -from typing import Optional +from typing import Optional, Union, Tuple from .dt_utils import debug_print, BaseEndpoint class Status(BaseEndpoint): @@ -9,8 +9,9 @@ def __init__(self, client): def get(self, includechildren: Optional[bool] = None, fast: Optional[bool] = None, - responsedata: Optional[str] = None - ): + responsedata: Optional[str] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = None + ): """ Get detailed system health and status information from Darktrace. @@ -18,6 +19,7 @@ def get(self, includechildren (bool, optional): Whether to include information about probes (children). True by default. fast (bool, optional): When true, returns data faster but may omit subnet connectivity information if not cached. responsedata (str, optional): Restrict the returned JSON to only the specified top-level field(s) or object(s). + timeout (float or tuple, optional): Request timeout in seconds. Can be a single value or (connect_timeout, read_timeout). Returns: dict: System health and status information from Darktrace. @@ -35,6 +37,7 @@ def get(self, headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() \ No newline at end of file diff --git a/darktrace/dt_subnets.py b/darktrace/dt_subnets.py index f050d21..95629f6 100644 --- a/darktrace/dt_subnets.py +++ b/darktrace/dt_subnets.py @@ -1,5 +1,5 @@ import requests -from typing import Optional +from typing import Optional, Union, Tuple from .dt_utils import debug_print, BaseEndpoint class Subnets(BaseEndpoint): @@ -10,7 +10,8 @@ def get(self, subnet_id: Optional[int] = None, seensince: Optional[str] = None, sid: Optional[int] = None, - responsedata: Optional[str] = None + responsedata: Optional[str] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = None ): """ Get subnet information from Darktrace. @@ -21,6 +22,7 @@ def get(self, Minimum=1 second, Maximum=6 months. Subnets with activity in the specified time period are returned. sid (int, optional): Identification number of a subnet modeled in the Darktrace system. responsedata (str, optional): Restrict the returned JSON to only the specified top-level field(s) or object(s). + timeout (float or tuple, optional): Timeout for the request in seconds. Can be a single float or a tuple of (connect_timeout, read_timeout). Returns: list or dict: Subnet information from Darktrace. @@ -38,7 +40,8 @@ def get(self, headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() @@ -53,7 +56,8 @@ def post(self, uniqueHostnames: Optional[bool] = None, excluded: Optional[bool] = None, modelExcluded: Optional[bool] = None, - responsedata: Optional[str] = None + responsedata: Optional[str] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = None ): """ Create or update a subnet in Darktrace. @@ -70,6 +74,7 @@ def post(self, excluded (bool, optional): Whether traffic in this subnet should not be processed at all. modelExcluded (bool, optional): Whether devices within this subnet should be fully modeled. If true, the devices will be added to the Internal Traffic subnet. responsedata (str, optional): Restrict the returned JSON to only the specified top-level field(s) or object(s). + timeout (float or tuple, optional): Timeout for the request in seconds. Can be a single float or a tuple of (connect_timeout, read_timeout). Returns: dict: Result of the subnet creation or update operation. @@ -101,6 +106,7 @@ def post(self, headers, _ = self._get_headers(endpoint) self.client._debug(f"POST {url} body={body}") - response = requests.post(url, headers=headers, json=body, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.post(url, headers=headers, json=body, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() \ No newline at end of file diff --git a/darktrace/dt_summarystatistics.py b/darktrace/dt_summarystatistics.py index a9264e5..4f31e41 100644 --- a/darktrace/dt_summarystatistics.py +++ b/darktrace/dt_summarystatistics.py @@ -1,5 +1,5 @@ import requests -from typing import Optional +from typing import Optional, Union, Tuple from .dt_utils import debug_print, BaseEndpoint class SummaryStatistics(BaseEndpoint): @@ -13,8 +13,9 @@ def get(self, to: Optional[str] = None, hours: Optional[int] = None, csensor: Optional[bool] = None, - mitreTactics: Optional[bool] = None - ): + mitreTactics: Optional[bool] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = None + ): """ Get summary statistics information from Darktrace. @@ -26,6 +27,7 @@ def get(self, hours (int, optional): Number of hour intervals from the end time (or current time) to return. Requires eventtype. csensor (bool, optional): When true, only bandwidth statistics for cSensor agents are returned. When false, statistics for Darktrace/Network bandwidth. mitreTactics (bool, optional): When true, alters the returned data to display MITRE ATT&CK Framework breakdown. + timeout (float or tuple, optional): Request timeout in seconds. Can be a single value or (connect_timeout, read_timeout). Returns: dict: Summary statistics information from Darktrace. @@ -51,6 +53,7 @@ def get(self, headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() \ No newline at end of file diff --git a/darktrace/dt_tags.py b/darktrace/dt_tags.py index de24bdd..635990c 100644 --- a/darktrace/dt_tags.py +++ b/darktrace/dt_tags.py @@ -1,5 +1,5 @@ import requests -from typing import Optional +from typing import Optional, Union, Tuple from .dt_utils import debug_print, BaseEndpoint class Tags(BaseEndpoint): @@ -11,7 +11,8 @@ def __init__(self, client): def get(self, tag_id: Optional[str] = None, tag: Optional[str] = None, - responsedata: Optional[str] = None + responsedata: Optional[str] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = None ): """ Get tag information from Darktrace. @@ -20,6 +21,7 @@ def get(self, tag_id (str, optional): Tag ID (tid) to retrieve a specific tag by ID (e.g., /tags/5). tag (str, optional): Name of an existing tag (e.g., /tags?tag=active threat). responsedata (str, optional): Restrict the returned JSON to only the specified field or object. + timeout (float or tuple, optional): Request timeout in seconds. Returns: dict or list: Tag information from Darktrace. @@ -35,11 +37,12 @@ def get(self, headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def create(self, name: str, color: Optional[int] = None, description: Optional[str] = None): + def create(self, name: str, color: Optional[int] = None, description: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None): """ Create a new tag in Darktrace. @@ -47,6 +50,7 @@ def create(self, name: str, color: Optional[int] = None, description: Optional[s name (str): Name for the created tag (required). color (int, optional): The hue value (in HSL) for the tag in the UI. description (str, optional): Optional description for the tag. + timeout (float or tuple, optional): Request timeout in seconds. Returns: dict: The created tag information from Darktrace. @@ -61,16 +65,18 @@ def create(self, name: str, color: Optional[int] = None, description: Optional[s headers, _ = self._get_headers(endpoint) self.client._debug(f"POST {url} body={body}") - response = requests.post(url, headers=headers, json=body, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.post(url, headers=headers, json=body, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def delete(self, tag_id: str) -> dict: + def delete(self, tag_id: str, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: """ Delete a tag by tag ID (tid). Args: tag_id (str): Tag ID (tid) to delete (e.g., /tags/5). + timeout (float or tuple, optional): Request timeout in seconds. Returns: dict: The response from the Darktrace API. @@ -79,14 +85,15 @@ def delete(self, tag_id: str) -> dict: url = f"{self.client.host}{endpoint}" headers, _ = self._get_headers(endpoint) self.client._debug(f"DELETE {url}") - response = requests.delete(url, headers=headers, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.delete(url, headers=headers, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() #TAGS/ENTITIES ENDPOINT - def get_entities(self, did: Optional[int] = None, tag: Optional[str] = None, responsedata: Optional[str] = None, fulldevicedetails: Optional[bool] = None): + def get_entities(self, did: Optional[int] = None, tag: Optional[str] = None, responsedata: Optional[str] = None, fulldevicedetails: Optional[bool] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None): """ Get tags for a device or devices for a tag via /tags/entities. @@ -95,6 +102,7 @@ def get_entities(self, did: Optional[int] = None, tag: Optional[str] = None, res tag (str, optional): Name of an existing tag to list devices for a tag. responsedata (str, optional): Restrict the returned JSON to only the specified field or object. fulldevicedetails (bool, optional): If true and a tag is queried, adds a devices object to the response with more detailed device data. + timeout (float or tuple, optional): Request timeout in seconds. Returns: list or dict: Tag or device information from Darktrace. @@ -112,11 +120,12 @@ def get_entities(self, did: Optional[int] = None, tag: Optional[str] = None, res params['fulldevicedetails'] = fulldevicedetails headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def post_entities(self, did: int, tag: str, duration: Optional[int] = None): + def post_entities(self, did: int, tag: str, duration: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None): """ Add a tag to a device via /tags/entities (POST, form-encoded). @@ -124,6 +133,7 @@ def post_entities(self, did: int, tag: str, duration: Optional[int] = None): did (int): Device ID to tag. tag (str): Name of the tag to add. duration (int, optional): How long the tag should be set for the device (seconds). + timeout (float or tuple, optional): Request timeout in seconds. Returns: dict: API response from Darktrace. @@ -135,17 +145,19 @@ def post_entities(self, did: int, tag: str, duration: Optional[int] = None): data['duration'] = duration headers, _ = self._get_headers(endpoint) self.client._debug(f"POST {url} data={data}") - response = requests.post(url, headers=headers, data=data, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.post(url, headers=headers, data=data, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def delete_entities(self, did: int, tag: str) -> dict: + def delete_entities(self, did: int, tag: str, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: """ Remove a tag from a device via /tags/entities (DELETE). Args: did (int): Device ID to untag. tag (str): Name of the tag to remove. + timeout (float or tuple, optional): Request timeout in seconds. Returns: dict: The response from the Darktrace API. @@ -155,12 +167,13 @@ def delete_entities(self, did: int, tag: str) -> dict: params = {'did': did, 'tag': tag} headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"DELETE {url} params={sorted_params}") - response = requests.delete(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.delete(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() # /tags/[tid]/entities ENDPOINT - def get_tag_entities(self, tid: int, responsedata: Optional[str] = None, fulldevicedetails: Optional[bool] = None): + def get_tag_entities(self, tid: int, responsedata: Optional[str] = None, fulldevicedetails: Optional[bool] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None): """ Get entities (devices or credentials) associated with a specific tag via /tags/[tid]/entities (GET). @@ -168,6 +181,7 @@ def get_tag_entities(self, tid: int, responsedata: Optional[str] = None, fulldev tid (int): Tag ID (tid) to query. responsedata (str, optional): Restrict the returned JSON to only the specified field or object. fulldevicedetails (bool, optional): If true, adds a devices object to the response with more detailed device data. + timeout (float or tuple, optional): Request timeout in seconds. Returns: list or dict: Entities associated with the tag from Darktrace. @@ -181,11 +195,12 @@ def get_tag_entities(self, tid: int, responsedata: Optional[str] = None, fulldev params['fulldevicedetails'] = fulldevicedetails headers, sorted_params = self._get_headers(endpoint, params) self.client._debug(f"GET {url} params={sorted_params}") - response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.get(url, headers=headers, params=sorted_params, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def post_tag_entities(self, tid: int, entityType: str, entityValue, expiryDuration: Optional[int] = None): + def post_tag_entities(self, tid: int, entityType: str, entityValue, expiryDuration: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None): """ Add a tag to one or more entities (device or credential) via /tags/[tid]/entities (POST, JSON body). @@ -194,6 +209,7 @@ def post_tag_entities(self, tid: int, entityType: str, entityValue, expiryDurati entityType (str): The type of entity to be tagged. Valid values: 'Device', 'Credential'. entityValue (str or list): For devices, the did (as string or list of strings). For credentials, the credential value(s). expiryDuration (int, optional): Duration in seconds the tag should be applied for. + timeout (float or tuple, optional): Request timeout in seconds. Returns: dict: API response from Darktrace. @@ -205,17 +221,19 @@ def post_tag_entities(self, tid: int, entityType: str, entityValue, expiryDurati body["expiryDuration"] = expiryDuration headers, _ = self._get_headers(endpoint) self.client._debug(f"POST {url} body={body}") - response = requests.post(url, headers=headers, json=body, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.post(url, headers=headers, json=body, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() - def delete_tag_entity(self, tid: int, teid: int) -> dict: + def delete_tag_entity(self, tid: int, teid: int, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: """ Remove a tag from an entity via /tags/[tid]/entities/[teid] (DELETE). Args: tid (int): Tag ID (tid). teid (int): Tag entity ID (teid) representing the tag-to-entity relationship. + timeout (float or tuple, optional): Request timeout in seconds. Returns: dict: The response from the Darktrace API. @@ -224,6 +242,7 @@ def delete_tag_entity(self, tid: int, teid: int) -> dict: url = f"{self.client.host}{endpoint}" headers, _ = self._get_headers(endpoint) self.client._debug(f"DELETE {url}") - response = requests.delete(url, headers=headers, verify=self.client.verify_ssl) + resolved_timeout = self._resolve_timeout(timeout) + response = requests.delete(url, headers=headers, verify=self.client.verify_ssl, timeout=resolved_timeout) response.raise_for_status() return response.json() \ No newline at end of file diff --git a/darktrace/dt_utils.py b/darktrace/dt_utils.py index 705949a..bf1dac1 100644 --- a/darktrace/dt_utils.py +++ b/darktrace/dt_utils.py @@ -1,6 +1,9 @@ import base64 import json -from typing import Dict, Any, Optional, Tuple +from typing import Dict, Any, Optional, Tuple, Union + +# Type alias for timeout parameter - can be None, float, or tuple of (connect, read) +TimeoutType = Optional[Union[float, Tuple[float, float]]] def debug_print(message: str, debug: bool = False): if debug: @@ -12,6 +15,12 @@ class BaseEndpoint: def __init__(self, client): self.client = client + def _resolve_timeout(self, timeout: TimeoutType) -> TimeoutType: + """Resolve timeout value, using client default if not specified.""" + if timeout is not None: + return timeout + return getattr(self.client, 'timeout', None) + def _get_headers(self, endpoint: str, params: Optional[Dict[str, Any]] = None, json_body: Optional[Dict[str, Any]] = None) -> Tuple[Dict[str, str], Optional[Dict[str, Any]]]: """ Get authentication headers and sorted parameters for API requests. diff --git a/docs/README.md b/docs/README.md index e0a8566..6327dc7 100644 --- a/docs/README.md +++ b/docs/README.md @@ -26,6 +26,7 @@ client = DarktraceClient( | `private_token` | str | required | Your Darktrace API private token | | `debug` | bool | False | Enable debug logging | | `verify_ssl` | bool | True | Enable SSL certificate verification | +| `timeout` | float \| tuple | None | Request timeout in seconds, or (connect, read) tuple. None = no timeout | ### SSL Verification @@ -42,6 +43,26 @@ client = DarktraceClient( > ⚠️ **Warning**: Disabling SSL verification is not recommended for production environments. +### Request Timeouts + +Configure timeouts for API requests. This is especially important for long-running queries like advanced search: + +```python +# Client-level timeout (applies to all requests) +client = DarktraceClient( + host="https://your-darktrace-instance", + public_token="YOUR_PUBLIC_TOKEN", + private_token="YOUR_PRIVATE_TOKEN", + timeout=30 # 30 seconds +) + +# Per-request override for slow queries +results = client.advanced_search.search(query, timeout=600) # 10 minutes + +# Tuple format: (connect_timeout, read_timeout) +results = client.advanced_search.search(query, timeout=(5, 300)) # 5s connect, 5min read +``` + ## Available Modules The Darktrace SDK provides access to all Darktrace API endpoints through the following modules: diff --git a/tests/test_timeout.py b/tests/test_timeout.py new file mode 100644 index 0000000..310013c --- /dev/null +++ b/tests/test_timeout.py @@ -0,0 +1,186 @@ +""" +Test suite for timeout configuration in Darktrace SDK. + +Tests verify that timeout is properly passed to requests at: +1. Client level (default timeout) +2. Per-request level (override) +3. Tuple format (connect, read timeouts) +""" +import pytest +from unittest.mock import Mock, patch, call +from darktrace import DarktraceClient, TimeoutType + + +@pytest.fixture +def mock_response(): + """Create a mock response object.""" + response = Mock() + response.raise_for_status = Mock() + response.json.return_value = {"devices": []} + return response + + +class TestClientLevelTimeout: + """Test client-level timeout configuration.""" + + def test_client_with_no_timeout_default(self, mock_response): + """Client without timeout should pass None to requests (no timeout).""" + client = DarktraceClient( + host="https://test.darktrace.com", + public_token="test_public", + private_token="test_private", + ) + assert client.timeout is None + + with patch("darktrace.dt_devices.requests.get", return_value=mock_response) as mock_get: + client.devices.get() + + # Verify timeout=None was passed (no timeout) + call_kwargs = mock_get.call_args[1] + assert call_kwargs.get("timeout") is None + + def test_client_with_float_timeout(self, mock_response): + """Client with float timeout should pass it to all requests.""" + client = DarktraceClient( + host="https://test.darktrace.com", + public_token="test_public", + private_token="test_private", + timeout=30.0, + ) + assert client.timeout == 30.0 + + with patch("darktrace.dt_devices.requests.get", return_value=mock_response) as mock_get: + client.devices.get() + + call_kwargs = mock_get.call_args[1] + assert call_kwargs.get("timeout") == 30.0 + + def test_client_with_tuple_timeout(self, mock_response): + """Client with tuple timeout should pass (connect, read) tuple.""" + client = DarktraceClient( + host="https://test.darktrace.com", + public_token="test_public", + private_token="test_private", + timeout=(5.0, 30.0), + ) + assert client.timeout == (5.0, 30.0) + + with patch("darktrace.dt_devices.requests.get", return_value=mock_response) as mock_get: + client.devices.get() + + call_kwargs = mock_get.call_args[1] + assert call_kwargs.get("timeout") == (5.0, 30.0) + + +class TestPerRequestTimeout: + """Test per-request timeout override.""" + + def test_per_request_override_float(self, mock_response): + """Per-request timeout should override client default.""" + client = DarktraceClient( + host="https://test.darktrace.com", + public_token="test_public", + private_token="test_private", + timeout=30.0, # Client default + ) + + with patch("darktrace.dt_devices.requests.get", return_value=mock_response) as mock_get: + client.devices.get(timeout=60.0) # Per-request override + + call_kwargs = mock_get.call_args[1] + assert call_kwargs.get("timeout") == 60.0 + + def test_per_request_override_tuple(self, mock_response): + """Per-request tuple timeout should work.""" + client = DarktraceClient( + host="https://test.darktrace.com", + public_token="test_public", + private_token="test_private", + ) + + with patch("darktrace.dt_devices.requests.get", return_value=mock_response) as mock_get: + client.devices.get(timeout=(10.0, 60.0)) + + call_kwargs = mock_get.call_args[1] + assert call_kwargs.get("timeout") == (10.0, 60.0) + + def test_per_request_none_uses_client_default(self, mock_response): + """Per-request None should use client default.""" + client = DarktraceClient( + host="https://test.darktrace.com", + public_token="test_public", + private_token="test_private", + timeout=30.0, + ) + + with patch("darktrace.dt_devices.requests.get", return_value=mock_response) as mock_get: + # Explicitly passing None should use client default + client.devices.get(timeout=None) + + call_kwargs = mock_get.call_args[1] + assert call_kwargs.get("timeout") == 30.0 + + +class TestTimeoutType: + """Test TimeoutType export for type hints.""" + + def test_timeout_type_import(self): + """TimeoutType should be importable from darktrace.""" + from darktrace import TimeoutType + assert TimeoutType is not None + + def test_timeout_type_annotation(self): + """TimeoutType should work as type annotation.""" + timeout: TimeoutType = 30.0 + assert timeout == 30.0 + + timeout_tuple: TimeoutType = (5.0, 30.0) + assert timeout_tuple == (5.0, 30.0) + + timeout_none: TimeoutType = None + assert timeout_none is None + + +class TestMultipleEndpoints: + """Test timeout works across different endpoint types.""" + + def test_timeout_on_get_request(self, mock_response): + """Timeout should work on GET requests.""" + client = DarktraceClient( + host="https://test.darktrace.com", + public_token="test_public", + private_token="test_private", + timeout=45.0, + ) + + with patch("darktrace.dt_devices.requests.get", return_value=mock_response) as mock_get: + client.devices.get() + assert mock_get.call_args[1]["timeout"] == 45.0 + + def test_timeout_on_post_request(self, mock_response): + """Timeout should work on POST requests.""" + client = DarktraceClient( + host="https://test.darktrace.com", + public_token="test_public", + private_token="test_private", + timeout=60.0, + ) + mock_response.json.return_value = {"tid": 1} + + with patch("darktrace.dt_tags.requests.post", return_value=mock_response) as mock_post: + client.tags.create(name="test-tag") + assert mock_post.call_args[1]["timeout"] == 60.0 + + def test_timeout_on_delete_request(self, mock_response): + """Timeout should work on DELETE requests.""" + client = DarktraceClient( + host="https://test.darktrace.com", + public_token="test_public", + private_token="test_private", + timeout=15.0, + ) + mock_response.json.return_value = {"deleted": True} + + with patch("darktrace.dt_tags.requests.delete", return_value=mock_response) as mock_delete: + client.tags.delete(tag_id="123") + assert mock_delete.call_args[1]["timeout"] == 15.0 From 968445cfa46a6d20315db1ad4720f640535c13e6 Mon Sep 17 00:00:00 2001 From: LegendEvent Date: Tue, 17 Feb 2026 18:53:39 +0100 Subject: [PATCH 2/2] fix: use sentinel pattern for timeout to allow None override - Add _UNSET sentinel to distinguish 'not provided' from 'None' - timeout=None now disables timeout (no timeout) - timeout not provided uses client default - Remove unused Union/Tuple imports from client.py - Update tests and documentation Addresses Copilot PR review comments #1 and #2 --- README.md | 7 ++++--- darktrace/client.py | 2 +- darktrace/dt_advanced_search.py | 8 ++++---- darktrace/dt_analyst.py | 24 ++++++++++++------------ darktrace/dt_antigena.py | 14 +++++++------- darktrace/dt_breaches.py | 16 ++++++++-------- darktrace/dt_components.py | 4 ++-- darktrace/dt_cves.py | 4 ++-- darktrace/dt_details.py | 4 ++-- darktrace/dt_deviceinfo.py | 4 ++-- darktrace/dt_devices.py | 6 +++--- darktrace/dt_devicesearch.py | 18 +++++++++--------- darktrace/dt_devicesummary.py | 4 ++-- darktrace/dt_email.py | 30 +++++++++++++++--------------- darktrace/dt_endpointdetails.py | 2 +- darktrace/dt_enums.py | 4 ++-- darktrace/dt_filtertypes.py | 4 ++-- darktrace/dt_intelfeed.py | 6 +++--- darktrace/dt_mbcomments.py | 6 +++--- darktrace/dt_metricdata.py | 4 ++-- darktrace/dt_metrics.py | 4 ++-- darktrace/dt_models.py | 4 ++-- darktrace/dt_network.py | 2 +- darktrace/dt_pcaps.py | 6 +++--- darktrace/dt_similardevices.py | 4 ++-- darktrace/dt_status.py | 2 +- darktrace/dt_subnets.py | 2 +- darktrace/dt_summarystatistics.py | 2 +- darktrace/dt_tags.py | 18 +++++++++--------- darktrace/dt_utils.py | 19 ++++++++++++++----- tests/test_timeout.py | 27 +++++++++++++++++++++------ 31 files changed, 143 insertions(+), 118 deletions(-) diff --git a/README.md b/README.md index 6219a31..e612f33 100644 --- a/README.md +++ b/README.md @@ -116,9 +116,10 @@ results = client.advanced_search.search(query, timeout=(5, 300)) | Format | Description | |--------|-------------| -| `timeout=None` | No timeout (default, waits indefinitely) | -| `timeout=30` | 30 seconds total (both connect and read) | -| `timeout=(5, 30)` | 5 seconds to connect, 30 seconds to read | +| (not passed) | Uses client default timeout | +| `timeout=None` | No timeout (disables client default) | +| `timeout=30` | 30 seconds total | +| `timeout=(5, 30)` | 5 seconds connect, 30 seconds read | > **Note**: Advanced search queries can take 5-10 minutes for complex queries. Consider using per-request timeouts for these endpoints. diff --git a/darktrace/client.py b/darktrace/client.py index 63af342..1cdcbc3 100644 --- a/darktrace/client.py +++ b/darktrace/client.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Union, Tuple +from typing import TYPE_CHECKING from .auth import DarktraceAuth from .dt_antigena import Antigena diff --git a/darktrace/dt_advanced_search.py b/darktrace/dt_advanced_search.py index 17559fe..96d8a4f 100644 --- a/darktrace/dt_advanced_search.py +++ b/darktrace/dt_advanced_search.py @@ -1,13 +1,13 @@ import requests import json from typing import Dict, Any, Optional, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint, encode_query +from .dt_utils import debug_print, BaseEndpoint, encode_query, _UNSET class AdvancedSearch(BaseEndpoint): def __init__(self, client): super().__init__(client) - def search(self, query: Dict[str, Any], post_request: bool = False, timeout: Optional[Union[float, Tuple[float, float]]] = None): + def search(self, query: Dict[str, Any], post_request: bool = False, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET): # type: ignore[assignment] """Perform Advanced Search query. Parameters: @@ -95,7 +95,7 @@ def search(self, query: Dict[str, Any], post_request: bool = False, timeout: Opt response.raise_for_status() return response.json() - def analyze(self, field: str, analysis_type: str, query: Dict[str, Any], timeout: Optional[Union[float, Tuple[float, float]]] = None): + def analyze(self, field: str, analysis_type: str, query: Dict[str, Any], timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET): # type: ignore[assignment] """Analyze field data.""" encoded_query = encode_query(query) endpoint = f'/advancedsearch/api/analyze/{field}/{analysis_type}/{encoded_query}' @@ -107,7 +107,7 @@ def analyze(self, field: str, analysis_type: str, query: Dict[str, Any], timeout response.raise_for_status() return response.json() - def graph(self, graph_type: str, interval: int, query: Dict[str, Any], timeout: Optional[Union[float, Tuple[float, float]]] = None): + def graph(self, graph_type: str, interval: int, query: Dict[str, Any], timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET): # type: ignore[assignment] """Get graph data.""" encoded_query = encode_query(query) endpoint = f'/advancedsearch/api/graph/{graph_type}/{interval}/{encoded_query}' diff --git a/darktrace/dt_analyst.py b/darktrace/dt_analyst.py index 0a3cbe5..79ca9d0 100644 --- a/darktrace/dt_analyst.py +++ b/darktrace/dt_analyst.py @@ -1,13 +1,13 @@ import requests import json from typing import Union, List, Dict, Any, Optional, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class Analyst(BaseEndpoint): def __init__(self, client): super().__init__(client) - def get_groups(self, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): + def get_groups(self, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **params): # type: ignore[assignment] """Get AI Analyst incident groups. Available parameters: @@ -38,7 +38,7 @@ def get_groups(self, timeout: Optional[Union[float, Tuple[float, float]]] = None response.raise_for_status() return response.json() - def get_incident_events(self, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): + def get_incident_events(self, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **params): # type: ignore[assignment] """Get AI Analyst incident events. Available parameters: @@ -74,7 +74,7 @@ def get_incident_events(self, timeout: Optional[Union[float, Tuple[float, float] response.raise_for_status() return response.json() - def acknowledge(self, uuids: Union[str, List[str]], timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: + def acknowledge(self, uuids: Union[str, List[str]], timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> dict: # type: ignore[assignment] """ Acknowledge AI Analyst incident events. @@ -92,7 +92,7 @@ def acknowledge(self, uuids: Union[str, List[str]], timeout: Optional[Union[floa response.raise_for_status() return response.json() - def unacknowledge(self, uuids: Union[str, List[str]], timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: + def unacknowledge(self, uuids: Union[str, List[str]], timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> dict: # type: ignore[assignment] """ Unacknowledge AI Analyst incident events. @@ -110,7 +110,7 @@ def unacknowledge(self, uuids: Union[str, List[str]], timeout: Optional[Union[fl response.raise_for_status() return response.json() - def pin(self, uuids: Union[str, List[str]], timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: + def pin(self, uuids: Union[str, List[str]], timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> dict: # type: ignore[assignment] """ Pin AI Analyst incident events. @@ -128,7 +128,7 @@ def pin(self, uuids: Union[str, List[str]], timeout: Optional[Union[float, Tuple response.raise_for_status() return response.json() - def unpin(self, uuids: Union[str, List[str]], timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: + def unpin(self, uuids: Union[str, List[str]], timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> dict: # type: ignore[assignment] """ Unpin AI Analyst incident events. @@ -146,7 +146,7 @@ def unpin(self, uuids: Union[str, List[str]], timeout: Optional[Union[float, Tup response.raise_for_status() return response.json() - def get_comments(self, incident_id: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, response_data: Optional[str] = ""): + def get_comments(self, incident_id: str, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, response_data: Optional[str] = ""): # type: ignore[assignment] """Get comments for an AI Analyst incident event. Parameters: @@ -165,7 +165,7 @@ def get_comments(self, incident_id: str, timeout: Optional[Union[float, Tuple[fl response.raise_for_status() return response.json() - def add_comment(self, incident_id: str, message: str, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: + def add_comment(self, incident_id: str, message: str, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> dict: # type: ignore[assignment] """Add a comment to an AI Analyst incident event. Parameters: @@ -186,7 +186,7 @@ def add_comment(self, incident_id: str, message: str, timeout: Optional[Union[fl response.raise_for_status() return response.json() - def get_stats(self, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): + def get_stats(self, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **params): # type: ignore[assignment] """Get AI Analyst statistics. Available parameters: @@ -212,7 +212,7 @@ def get_stats(self, timeout: Optional[Union[float, Tuple[float, float]]] = None, response.raise_for_status() return response.json() - def get_investigations(self, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): + def get_investigations(self, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **params): # type: ignore[assignment] """Get AI Analyst investigations (GET request). Available parameters: @@ -240,7 +240,7 @@ def get_investigations(self, timeout: Optional[Union[float, Tuple[float, float]] response.raise_for_status() return response.json() - def create_investigation(self, investigate_time: str, did: int, timeout: Optional[Union[float, Tuple[float, float]]] = None): + def create_investigation(self, investigate_time: str, did: int, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET): # type: ignore[assignment] """Create a new AI Analyst investigation (POST request). Parameters: diff --git a/darktrace/dt_antigena.py b/darktrace/dt_antigena.py index 2ef2d18..118b576 100644 --- a/darktrace/dt_antigena.py +++ b/darktrace/dt_antigena.py @@ -1,7 +1,7 @@ import requests import json from typing import Dict, Any, Union, Optional, List, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class Antigena(BaseEndpoint): @@ -20,7 +20,7 @@ class Antigena(BaseEndpoint): def __init__(self, client): super().__init__(client) - def get_actions(self, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): + def get_actions(self, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **params): # type: ignore[assignment] """ Get information about current and past Darktrace RESPOND actions. @@ -128,7 +128,7 @@ def activate_action( response.raise_for_status() return response.json() - def extend_action(self, codeid: int, duration: int, reason: str = "", timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: + def extend_action(self, codeid: int, duration: int, reason: str = "", timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> dict: # type: ignore[assignment] """ Extend an active Darktrace RESPOND action. @@ -179,7 +179,7 @@ def extend_action(self, codeid: int, duration: int, reason: str = "", timeout: O response.raise_for_status() return response.json() - def clear_action(self, codeid: int, reason: str = "", timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: + def clear_action(self, codeid: int, reason: str = "", timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> dict: # type: ignore[assignment] """ Clear an active, pending or expired Darktrace RESPOND action. @@ -226,7 +226,7 @@ def clear_action(self, codeid: int, reason: str = "", timeout: Optional[Union[fl response.raise_for_status() return response.json() - def reactivate_action(self, codeid: int, duration: int, reason: str = "", timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: + def reactivate_action(self, codeid: int, duration: int, reason: str = "", timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> dict: # type: ignore[assignment] """ Reactivate a cleared or expired Darktrace RESPOND action. @@ -279,7 +279,7 @@ def create_manual_action( duration: int, reason: str = "", connections: Optional[List[Dict[str, Union[str, int]]]] = None, - timeout: Optional[Union[float, Tuple[float, float]]] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, # type: ignore[assignment] ) -> dict: """ Create a manual Darktrace RESPOND/Network action. @@ -364,7 +364,7 @@ def create_manual_action( response.raise_for_status() return response.json() - def get_summary(self, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): + def get_summary(self, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **params): # type: ignore[assignment] """ Get a summary of active and pending Darktrace RESPOND actions. diff --git a/darktrace/dt_breaches.py b/darktrace/dt_breaches.py index 72463ae..d7a850a 100644 --- a/darktrace/dt_breaches.py +++ b/darktrace/dt_breaches.py @@ -2,13 +2,13 @@ import json from typing import Dict, Any, Optional, Union, Tuple from datetime import datetime -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class ModelBreaches(BaseEndpoint): def __init__(self, client): super().__init__(client) - def get(self, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): + def get(self, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **params): # type: ignore[assignment] """ Get model breach alerts from the /modelbreaches endpoint. @@ -75,7 +75,7 @@ def get(self, timeout: Optional[Union[float, Tuple[float, float]]] = None, **par response.raise_for_status() return response.json() - def get_comments(self, pbid: Union[int, list], timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): + def get_comments(self, pbid: Union[int, list], timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **params): # type: ignore[assignment] """ Get comments for a specific model breach alert. @@ -97,7 +97,7 @@ def get_comments(self, pbid: Union[int, list], timeout: Optional[Union[float, Tu response.raise_for_status() return response.json() - def add_comment(self, pbid: int, message: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params) -> dict: + def add_comment(self, pbid: int, message: str, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **params) -> dict: # type: ignore[assignment] """ Add a comment to a model breach alert. @@ -154,7 +154,7 @@ def add_comment(self, pbid: int, message: str, timeout: Optional[Union[float, Tu debug_print(f"BREACHES: Exception: {str(e)}", self.client.debug) return {"error": str(e)} - def acknowledge(self, pbid: Union[int, list], timeout: Optional[Union[float, Tuple[float, float]]] = None, **params) -> dict: + def acknowledge(self, pbid: Union[int, list], timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **params) -> dict: # type: ignore[assignment] """ Acknowledge a model breach alert. @@ -185,7 +185,7 @@ def acknowledge(self, pbid: Union[int, list], timeout: Optional[Union[float, Tup self.client._debug(f"Exception occurred while acknowledging breach: {str(e)}") return {"error": str(e)} - def unacknowledge(self, pbid: Union[int, list], timeout: Optional[Union[float, Tuple[float, float]]] = None, **params) -> dict: + def unacknowledge(self, pbid: Union[int, list], timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **params) -> dict: # type: ignore[assignment] """ Unacknowledge a model breach alert. @@ -216,7 +216,7 @@ def unacknowledge(self, pbid: Union[int, list], timeout: Optional[Union[float, T self.client._debug(f"Exception occurred while unacknowledging breach: {str(e)}") return {"error": str(e)} - def acknowledge_with_comment(self, pbid: int, message: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params) -> dict: + def acknowledge_with_comment(self, pbid: int, message: str, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **params) -> dict: # type: ignore[assignment] """ Acknowledge a model breach and add a comment in one call. @@ -235,7 +235,7 @@ def acknowledge_with_comment(self, pbid: int, message: str, timeout: Optional[Un "add_comment": comment_response } - def unacknowledge_with_comment(self, pbid: int, message: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params) -> dict: + def unacknowledge_with_comment(self, pbid: int, message: str, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **params) -> dict: # type: ignore[assignment] """ Unacknowledge a model breach and add a comment in one call. diff --git a/darktrace/dt_components.py b/darktrace/dt_components.py index 9d98bcd..2a80dd8 100644 --- a/darktrace/dt_components.py +++ b/darktrace/dt_components.py @@ -1,12 +1,12 @@ import requests from typing import Optional, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class Components(BaseEndpoint): def __init__(self, client): super().__init__(client) - def get(self, cid: Optional[int] = None, responsedata: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): + def get(self, cid: Optional[int] = None, responsedata: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **params): # type: ignore[assignment] """ Get information about model components. diff --git a/darktrace/dt_cves.py b/darktrace/dt_cves.py index 368a690..2d991d7 100644 --- a/darktrace/dt_cves.py +++ b/darktrace/dt_cves.py @@ -1,6 +1,6 @@ import requests from typing import Optional, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class CVEs(BaseEndpoint): def __init__(self, client): @@ -10,7 +10,7 @@ def get( self, did: Optional[int] = None, fulldevicedetails: Optional[bool] = None, - timeout: Optional[Union[float, Tuple[float, float]]] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, # type: ignore[assignment] **params ): """ diff --git a/darktrace/dt_details.py b/darktrace/dt_details.py index 57d1f9f..7da0dbd 100644 --- a/darktrace/dt_details.py +++ b/darktrace/dt_details.py @@ -1,6 +1,6 @@ import requests from typing import Optional, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class Details(BaseEndpoint): def __init__(self, client): @@ -31,7 +31,7 @@ def get( deduplicate: bool = False, fulldevicedetails: bool = False, responsedata: Optional[str] = None, - timeout: Optional[Union[float, Tuple[float, float]]] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, # type: ignore[assignment] **params ): """ diff --git a/darktrace/dt_deviceinfo.py b/darktrace/dt_deviceinfo.py index 68f73b5..61732e5 100644 --- a/darktrace/dt_deviceinfo.py +++ b/darktrace/dt_deviceinfo.py @@ -1,6 +1,6 @@ import requests from typing import Optional, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class DeviceInfo(BaseEndpoint): def __init__(self, client): @@ -17,7 +17,7 @@ def get( showallgraphdata: bool = True, similardevices: Optional[int] = None, intervalhours: int = 1, - timeout: Optional[Union[float, Tuple[float, float]]] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, # type: ignore[assignment] **params ): """ diff --git a/darktrace/dt_devices.py b/darktrace/dt_devices.py index cb9e5e3..f39f28d 100644 --- a/darktrace/dt_devices.py +++ b/darktrace/dt_devices.py @@ -1,7 +1,7 @@ import requests import json from typing import List, Dict, Any, Optional, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class Devices(BaseEndpoint): @@ -21,7 +21,7 @@ def get( responsedata: str = None, cloudsecurity: bool = None, saasfilter: Any = None, - timeout: Optional[Union[float, Tuple[float, float]]] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, # type: ignore[assignment] ): """ Update a single device. @@ -87,7 +87,7 @@ def get( response.raise_for_status() return response.json() - def update(self, did: int, timeout: Optional[Union[float, Tuple[float, float]]] = None, **kwargs) -> dict: + def update(self, did: int, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **kwargs) -> dict: # type: ignore[assignment] """Update device properties in Darktrace. Args: diff --git a/darktrace/dt_devicesearch.py b/darktrace/dt_devicesearch.py index 09d75cc..f89954d 100644 --- a/darktrace/dt_devicesearch.py +++ b/darktrace/dt_devicesearch.py @@ -1,6 +1,6 @@ import requests from typing import Optional, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class DeviceSearch(BaseEndpoint): @@ -37,7 +37,7 @@ def get( offset=None, responsedata=None, seensince=None, - timeout: Optional[Union[float, Tuple[float, float]]] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, # type: ignore[assignment] **kwargs, ): """ @@ -123,7 +123,7 @@ def get( response.raise_for_status() return response.json() - def get_tag(self, tag: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **kwargs): + def get_tag(self, tag: str, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **kwargs): # type: ignore[assignment] """ Search for devices with a specific tag. @@ -138,7 +138,7 @@ def get_tag(self, tag: str, timeout: Optional[Union[float, Tuple[float, float]]] query = f'tag:"{tag}"' return self.get(query=query, timeout=timeout, **kwargs) - def get_type(self, type: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **kwargs): + def get_type(self, type: str, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **kwargs): # type: ignore[assignment] """ Search for devices with a specific type. @@ -153,7 +153,7 @@ def get_type(self, type: str, timeout: Optional[Union[float, Tuple[float, float] query = f'type:"{type}"' return self.get(query=query, timeout=timeout, **kwargs) - def get_label(self, label: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **kwargs): + def get_label(self, label: str, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **kwargs): # type: ignore[assignment] """ Search for devices with a specific label. @@ -168,7 +168,7 @@ def get_label(self, label: str, timeout: Optional[Union[float, Tuple[float, floa query = f'label:"{label}"' return self.get(query=query, timeout=timeout, **kwargs) - def get_vendor(self, vendor: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **kwargs): + def get_vendor(self, vendor: str, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **kwargs): # type: ignore[assignment] """ Search for devices with a specific vendor. @@ -183,7 +183,7 @@ def get_vendor(self, vendor: str, timeout: Optional[Union[float, Tuple[float, fl query = f'vendor:"{vendor}"' return self.get(query=query, timeout=timeout, **kwargs) - def get_hostname(self, hostname: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **kwargs): + def get_hostname(self, hostname: str, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **kwargs): # type: ignore[assignment] """ Search for devices with a specific hostname. @@ -198,7 +198,7 @@ def get_hostname(self, hostname: str, timeout: Optional[Union[float, Tuple[float query = f'hostname:"{hostname}"' return self.get(query=query, timeout=timeout, **kwargs) - def get_ip(self, ip: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **kwargs): + def get_ip(self, ip: str, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **kwargs): # type: ignore[assignment] """ Search for devices with a specific IP address. @@ -213,7 +213,7 @@ def get_ip(self, ip: str, timeout: Optional[Union[float, Tuple[float, float]]] = query = f'ip:"{ip}"' return self.get(query=query, timeout=timeout, **kwargs) - def get_mac(self, mac: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **kwargs): + def get_mac(self, mac: str, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **kwargs): # type: ignore[assignment] """ Search for devices with a specific MAC address. diff --git a/darktrace/dt_devicesummary.py b/darktrace/dt_devicesummary.py index 869bccf..545a9b4 100644 --- a/darktrace/dt_devicesummary.py +++ b/darktrace/dt_devicesummary.py @@ -1,6 +1,6 @@ import requests from typing import Optional, Union, Tuple, Dict, Any -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class DeviceSummary(BaseEndpoint): """ @@ -44,7 +44,7 @@ def get( source: Optional[str] = None, status: Optional[str] = None, responsedata: Optional[str] = None, - timeout: Optional[Union[float, Tuple[float, float]]] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, # type: ignore[assignment] **kwargs ) -> Dict[str, Any]: """ diff --git a/darktrace/dt_email.py b/darktrace/dt_email.py index cb766c8..6938b8e 100644 --- a/darktrace/dt_email.py +++ b/darktrace/dt_email.py @@ -1,13 +1,13 @@ import requests import json from typing import Dict, Any, Optional, Union, List, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class DarktraceEmail(BaseEndpoint): def __init__(self, client): super().__init__(client) - def decode_link(self, link: str, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: + def decode_link(self, link: str, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> Dict[str, Any]: # type: ignore[assignment] """ Decode a link using the Darktrace/Email API. @@ -30,7 +30,7 @@ def decode_link(self, link: str, timeout: Optional[Union[float, Tuple[float, flo response.raise_for_status() return response.json() - def get_action_summary(self, days: Optional[int] = None, limit: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: + def get_action_summary(self, days: Optional[int] = None, limit: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> Dict[str, Any]: # type: ignore[assignment] """ Get action summary from Darktrace/Email API. @@ -58,7 +58,7 @@ def get_action_summary(self, days: Optional[int] = None, limit: Optional[int] = response.raise_for_status() return response.json() - def get_dash_stats(self, days: Optional[int] = None, limit: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: + def get_dash_stats(self, days: Optional[int] = None, limit: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> Dict[str, Any]: # type: ignore[assignment] """ Get dashboard stats from Darktrace/Email API. @@ -86,7 +86,7 @@ def get_dash_stats(self, days: Optional[int] = None, limit: Optional[int] = None response.raise_for_status() return response.json() - def get_data_loss(self, days: Optional[int] = None, limit: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: + def get_data_loss(self, days: Optional[int] = None, limit: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> Dict[str, Any]: # type: ignore[assignment] """ Get data loss information from Darktrace/Email API. @@ -114,7 +114,7 @@ def get_data_loss(self, days: Optional[int] = None, limit: Optional[int] = None, response.raise_for_status() return response.json() - def get_user_anomaly(self, days: Optional[int] = None, limit: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: + def get_user_anomaly(self, days: Optional[int] = None, limit: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> Dict[str, Any]: # type: ignore[assignment] """ Get user anomaly data from Darktrace/Email API. @@ -142,7 +142,7 @@ def get_user_anomaly(self, days: Optional[int] = None, limit: Optional[int] = No response.raise_for_status() return response.json() - def email_action(self, uuid: str, data: Dict[str, Any], timeout: Optional[Union[float, Tuple[float, float]]] = None): + def email_action(self, uuid: str, data: Dict[str, Any], timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET): # type: ignore[assignment] """Perform an action on an email by UUID in Darktrace/Email API.""" endpoint = f'/agemail/api/ep/api/v1.0/emails/{uuid}/action' url = f"{self.client.host}{endpoint}" @@ -156,7 +156,7 @@ def email_action(self, uuid: str, data: Dict[str, Any], timeout: Optional[Union[ response.raise_for_status() return response.json() - def get_email(self, uuid: str, include_headers: Optional[bool] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: + def get_email(self, uuid: str, include_headers: Optional[bool] = None, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> Dict[str, Any]: # type: ignore[assignment] """ Get email details by UUID from Darktrace/Email API. @@ -182,7 +182,7 @@ def get_email(self, uuid: str, include_headers: Optional[bool] = None, timeout: response.raise_for_status() return response.json() - def download_email(self, uuid: str, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> bytes: + def download_email(self, uuid: str, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> bytes: # type: ignore[assignment] """ Download an email by UUID from Darktrace/Email API. @@ -204,7 +204,7 @@ def download_email(self, uuid: str, timeout: Optional[Union[float, Tuple[float, response.raise_for_status() return response.content - def search_emails(self, data: Dict[str, Any], timeout: Optional[Union[float, Tuple[float, float]]] = None): + def search_emails(self, data: Dict[str, Any], timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET): # type: ignore[assignment] """Search emails in Darktrace/Email API.""" endpoint = '/agemail/api/ep/api/v1.0/emails/search' url = f"{self.client.host}{endpoint}" @@ -218,7 +218,7 @@ def search_emails(self, data: Dict[str, Any], timeout: Optional[Union[float, Tup response.raise_for_status() return response.json() - def get_tags(self, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: + def get_tags(self, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> Dict[str, Any]: # type: ignore[assignment] """ Get tags from Darktrace/Email API. @@ -239,7 +239,7 @@ def get_tags(self, timeout: Optional[Union[float, Tuple[float, float]]] = None) response.raise_for_status() return response.json() - def get_actions(self, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: + def get_actions(self, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> Dict[str, Any]: # type: ignore[assignment] """ Get actions from Darktrace/Email API. @@ -260,7 +260,7 @@ def get_actions(self, timeout: Optional[Union[float, Tuple[float, float]]] = Non response.raise_for_status() return response.json() - def get_filters(self, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: + def get_filters(self, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> Dict[str, Any]: # type: ignore[assignment] """ Get filters from Darktrace/Email API. @@ -281,7 +281,7 @@ def get_filters(self, timeout: Optional[Union[float, Tuple[float, float]]] = Non response.raise_for_status() return response.json() - def get_event_types(self, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: + def get_event_types(self, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> Dict[str, Any]: # type: ignore[assignment] """ Get audit event types from Darktrace/Email API. @@ -302,7 +302,7 @@ def get_event_types(self, timeout: Optional[Union[float, Tuple[float, float]]] = response.raise_for_status() return response.json() - def get_audit_events(self, event_type: Optional[str] = None, limit: Optional[int] = None, offset: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> Dict[str, Any]: + def get_audit_events(self, event_type: Optional[str] = None, limit: Optional[int] = None, offset: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> Dict[str, Any]: # type: ignore[assignment] """ Get audit events from Darktrace/Email API. diff --git a/darktrace/dt_endpointdetails.py b/darktrace/dt_endpointdetails.py index e7cedcd..d949dc2 100644 --- a/darktrace/dt_endpointdetails.py +++ b/darktrace/dt_endpointdetails.py @@ -1,6 +1,6 @@ import requests from typing import Optional, Any, Dict, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class EndpointDetails(BaseEndpoint): def __init__(self, client): diff --git a/darktrace/dt_enums.py b/darktrace/dt_enums.py index 21bbc41..78d8566 100644 --- a/darktrace/dt_enums.py +++ b/darktrace/dt_enums.py @@ -1,6 +1,6 @@ import requests from typing import Optional, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class Enums(BaseEndpoint): """ @@ -11,7 +11,7 @@ class Enums(BaseEndpoint): def __init__(self, client): super().__init__(client) - def get(self, responsedata: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): + def get(self, responsedata: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **params): # type: ignore[assignment] """ Get enum values for all types or restrict to a specific field/object. diff --git a/darktrace/dt_filtertypes.py b/darktrace/dt_filtertypes.py index 94dac50..4322fa0 100644 --- a/darktrace/dt_filtertypes.py +++ b/darktrace/dt_filtertypes.py @@ -1,6 +1,6 @@ import requests from typing import Optional, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class FilterTypes(BaseEndpoint): """ @@ -22,7 +22,7 @@ class FilterTypes(BaseEndpoint): def __init__(self, client): super().__init__(client) - def get(self, responsedata: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): + def get(self, responsedata: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **params): # type: ignore[assignment] """ Get all filter types or restrict to a specific field/object. diff --git a/darktrace/dt_intelfeed.py b/darktrace/dt_intelfeed.py index ef6cfbe..458d14d 100644 --- a/darktrace/dt_intelfeed.py +++ b/darktrace/dt_intelfeed.py @@ -1,7 +1,7 @@ import requests import json from typing import Optional, List, Dict, Any, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class IntelFeed(BaseEndpoint): @@ -27,7 +27,7 @@ def __init__(self, client): def get(self, sources: Optional[bool] = None, source: Optional[str] = None, fulldetails: Optional[bool] = None, responsedata: Optional[str] = None, - timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): + timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **params): # type: ignore[assignment] """ Get the intelfeed list, sources, or detailed entries. @@ -77,7 +77,7 @@ def update(self, add_entry: Optional[str] = None, add_list: Optional[List[str]] expiry: Optional[str] = None, is_hostname: bool = False, remove_entry: Optional[str] = None, remove_all: bool = False, enable_antigena: bool = False, - timeout: Optional[Union[float, Tuple[float, float]]] = None): + timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET): # type: ignore[assignment] """Update the intel feed (watched domains) in Darktrace. Args: diff --git a/darktrace/dt_mbcomments.py b/darktrace/dt_mbcomments.py index 55b3926..891e82a 100644 --- a/darktrace/dt_mbcomments.py +++ b/darktrace/dt_mbcomments.py @@ -1,7 +1,7 @@ import requests import json from typing import Optional, Dict, Any, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class MBComments(BaseEndpoint): def __init__(self, client): @@ -14,7 +14,7 @@ def get(self, responsedata: Optional[str] = None, count: Optional[int] = None, pbid: Optional[int] = None, - timeout: Optional[Union[float, Tuple[float, float]]] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, # type: ignore[assignment] **params ): """ @@ -54,7 +54,7 @@ def get(self, response.raise_for_status() return response.json() - def post(self, breach_id: str, comment: str, timeout: Optional[Union[float, Tuple[float, float]]] = None, **params): + def post(self, breach_id: str, comment: str, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, **params): # type: ignore[assignment] """Add a comment to a model breach. Args: diff --git a/darktrace/dt_metricdata.py b/darktrace/dt_metricdata.py index 3ea7294..91e2253 100644 --- a/darktrace/dt_metricdata.py +++ b/darktrace/dt_metricdata.py @@ -1,6 +1,6 @@ import requests from typing import Optional, List, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class MetricData(BaseEndpoint): def __init__(self, client): @@ -26,7 +26,7 @@ def get( breachtimes: Optional[bool] = None, fulldevicedetails: Optional[bool] = None, devices: Optional[List[str]] = None, - timeout: Optional[Union[float, Tuple[float, float]]] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, # type: ignore[assignment] **params ): """ diff --git a/darktrace/dt_metrics.py b/darktrace/dt_metrics.py index 971488e..ea4ed38 100644 --- a/darktrace/dt_metrics.py +++ b/darktrace/dt_metrics.py @@ -1,6 +1,6 @@ import requests from typing import Optional, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class Metrics(BaseEndpoint): def __init__(self, client): @@ -10,7 +10,7 @@ def get( self, metric_id: Optional[int] = None, responsedata: Optional[str] = None, - timeout: Optional[Union[float, Tuple[float, float]]] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, # type: ignore[assignment] **params ): """ diff --git a/darktrace/dt_models.py b/darktrace/dt_models.py index d2492b4..d6a394c 100644 --- a/darktrace/dt_models.py +++ b/darktrace/dt_models.py @@ -1,12 +1,12 @@ import requests from typing import Optional, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class Models(BaseEndpoint): def __init__(self, client): super().__init__(client) - def get(self, uuid: Optional[str] = None, responsedata: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None): + def get(self, uuid: Optional[str] = None, responsedata: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET): # type: ignore[assignment] """ Get model information from Darktrace. diff --git a/darktrace/dt_network.py b/darktrace/dt_network.py index 7346960..2b7cf67 100644 --- a/darktrace/dt_network.py +++ b/darktrace/dt_network.py @@ -1,6 +1,6 @@ import requests from typing import Optional, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class Network(BaseEndpoint): def __init__(self, client): diff --git a/darktrace/dt_pcaps.py b/darktrace/dt_pcaps.py index 832bd01..6f7bda2 100644 --- a/darktrace/dt_pcaps.py +++ b/darktrace/dt_pcaps.py @@ -1,12 +1,12 @@ import requests from typing import Optional, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class PCAPs(BaseEndpoint): def __init__(self, client): super().__init__(client) - def get(self, pcap_id: Optional[str] = None, responsedata: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None): + def get(self, pcap_id: Optional[str] = None, responsedata: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET): # type: ignore[assignment] """ Retrieve PCAP information or download a specific PCAP file from Darktrace. @@ -30,7 +30,7 @@ def get(self, pcap_id: Optional[str] = None, responsedata: Optional[str] = None, # Return JSON if possible, else return raw content (for PCAP file download) return response.json() if 'application/json' in response.headers.get('Content-Type', '') else response.content - def create(self, ip1: str, start: int, end: int, ip2: Optional[str] = None, port1: Optional[int] = None, port2: Optional[int] = None, protocol: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None): + def create(self, ip1: str, start: int, end: int, ip2: Optional[str] = None, port1: Optional[int] = None, port2: Optional[int] = None, protocol: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET): # type: ignore[assignment] """ Create a new PCAP capture request in Darktrace. diff --git a/darktrace/dt_similardevices.py b/darktrace/dt_similardevices.py index 6a7920a..bf53af5 100644 --- a/darktrace/dt_similardevices.py +++ b/darktrace/dt_similardevices.py @@ -1,6 +1,6 @@ import requests from typing import Optional, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class SimilarDevices(BaseEndpoint): def __init__(self, client): @@ -13,7 +13,7 @@ def get( fulldevicedetails: Optional[bool] = None, token: Optional[str] = None, responsedata: Optional[str] = None, - timeout: Optional[Union[float, Tuple[float, float]]] = None, + timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET, # type: ignore[assignment] **kwargs ): """ diff --git a/darktrace/dt_status.py b/darktrace/dt_status.py index 68476d4..faf8bab 100644 --- a/darktrace/dt_status.py +++ b/darktrace/dt_status.py @@ -1,6 +1,6 @@ import requests from typing import Optional, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class Status(BaseEndpoint): def __init__(self, client): diff --git a/darktrace/dt_subnets.py b/darktrace/dt_subnets.py index 95629f6..c5ebf8c 100644 --- a/darktrace/dt_subnets.py +++ b/darktrace/dt_subnets.py @@ -1,6 +1,6 @@ import requests from typing import Optional, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class Subnets(BaseEndpoint): def __init__(self, client): diff --git a/darktrace/dt_summarystatistics.py b/darktrace/dt_summarystatistics.py index 4f31e41..41a4004 100644 --- a/darktrace/dt_summarystatistics.py +++ b/darktrace/dt_summarystatistics.py @@ -1,6 +1,6 @@ import requests from typing import Optional, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class SummaryStatistics(BaseEndpoint): def __init__(self, client): diff --git a/darktrace/dt_tags.py b/darktrace/dt_tags.py index 635990c..93328d4 100644 --- a/darktrace/dt_tags.py +++ b/darktrace/dt_tags.py @@ -1,6 +1,6 @@ import requests from typing import Optional, Union, Tuple -from .dt_utils import debug_print, BaseEndpoint +from .dt_utils import debug_print, BaseEndpoint, _UNSET class Tags(BaseEndpoint): @@ -42,7 +42,7 @@ def get(self, response.raise_for_status() return response.json() - def create(self, name: str, color: Optional[int] = None, description: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None): + def create(self, name: str, color: Optional[int] = None, description: Optional[str] = None, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET): # type: ignore[assignment] """ Create a new tag in Darktrace. @@ -70,7 +70,7 @@ def create(self, name: str, color: Optional[int] = None, description: Optional[s response.raise_for_status() return response.json() - def delete(self, tag_id: str, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: + def delete(self, tag_id: str, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> dict: # type: ignore[assignment] """ Delete a tag by tag ID (tid). @@ -93,7 +93,7 @@ def delete(self, tag_id: str, timeout: Optional[Union[float, Tuple[float, float] #TAGS/ENTITIES ENDPOINT - def get_entities(self, did: Optional[int] = None, tag: Optional[str] = None, responsedata: Optional[str] = None, fulldevicedetails: Optional[bool] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None): + def get_entities(self, did: Optional[int] = None, tag: Optional[str] = None, responsedata: Optional[str] = None, fulldevicedetails: Optional[bool] = None, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET): # type: ignore[assignment] """ Get tags for a device or devices for a tag via /tags/entities. @@ -125,7 +125,7 @@ def get_entities(self, did: Optional[int] = None, tag: Optional[str] = None, res response.raise_for_status() return response.json() - def post_entities(self, did: int, tag: str, duration: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None): + def post_entities(self, did: int, tag: str, duration: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET): # type: ignore[assignment] """ Add a tag to a device via /tags/entities (POST, form-encoded). @@ -150,7 +150,7 @@ def post_entities(self, did: int, tag: str, duration: Optional[int] = None, time response.raise_for_status() return response.json() - def delete_entities(self, did: int, tag: str, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: + def delete_entities(self, did: int, tag: str, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> dict: # type: ignore[assignment] """ Remove a tag from a device via /tags/entities (DELETE). @@ -173,7 +173,7 @@ def delete_entities(self, did: int, tag: str, timeout: Optional[Union[float, Tup return response.json() # /tags/[tid]/entities ENDPOINT - def get_tag_entities(self, tid: int, responsedata: Optional[str] = None, fulldevicedetails: Optional[bool] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None): + def get_tag_entities(self, tid: int, responsedata: Optional[str] = None, fulldevicedetails: Optional[bool] = None, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET): # type: ignore[assignment] """ Get entities (devices or credentials) associated with a specific tag via /tags/[tid]/entities (GET). @@ -200,7 +200,7 @@ def get_tag_entities(self, tid: int, responsedata: Optional[str] = None, fulldev response.raise_for_status() return response.json() - def post_tag_entities(self, tid: int, entityType: str, entityValue, expiryDuration: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = None): + def post_tag_entities(self, tid: int, entityType: str, entityValue, expiryDuration: Optional[int] = None, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET): # type: ignore[assignment] """ Add a tag to one or more entities (device or credential) via /tags/[tid]/entities (POST, JSON body). @@ -226,7 +226,7 @@ def post_tag_entities(self, tid: int, entityType: str, entityValue, expiryDurati response.raise_for_status() return response.json() - def delete_tag_entity(self, tid: int, teid: int, timeout: Optional[Union[float, Tuple[float, float]]] = None) -> dict: + def delete_tag_entity(self, tid: int, teid: int, timeout: Optional[Union[float, Tuple[float, float]]] = _UNSET) -> dict: # type: ignore[assignment] """ Remove a tag from an entity via /tags/[tid]/entities/[teid] (DELETE). diff --git a/darktrace/dt_utils.py b/darktrace/dt_utils.py index bf1dac1..8716898 100644 --- a/darktrace/dt_utils.py +++ b/darktrace/dt_utils.py @@ -5,19 +5,28 @@ # Type alias for timeout parameter - can be None, float, or tuple of (connect, read) TimeoutType = Optional[Union[float, Tuple[float, float]]] +# Sentinel value for unset timeout - allows distinguishing between +# "not specified" (use client default) and "explicitly None" (no timeout) +_UNSET = object() + def debug_print(message: str, debug: bool = False): if debug: print(f"DEBUG: {message}") class BaseEndpoint: """Base class for all Darktrace API endpoint modules.""" - + def __init__(self, client): self.client = client - - def _resolve_timeout(self, timeout: TimeoutType) -> TimeoutType: - """Resolve timeout value, using client default if not specified.""" - if timeout is not None: + + def _resolve_timeout(self, timeout: TimeoutType = _UNSET) -> TimeoutType: # type: ignore[assignment] + """Resolve timeout value, using client default if not specified. + + Args: + timeout: Per-request timeout. _UNSET (default) uses client.timeout. + None means no timeout. Float or tuple sets specific timeout. + """ + if timeout is not _UNSET: return timeout return getattr(self.client, 'timeout', None) diff --git a/tests/test_timeout.py b/tests/test_timeout.py index 310013c..15770fa 100644 --- a/tests/test_timeout.py +++ b/tests/test_timeout.py @@ -104,21 +104,36 @@ def test_per_request_override_tuple(self, mock_response): call_kwargs = mock_get.call_args[1] assert call_kwargs.get("timeout") == (10.0, 60.0) - def test_per_request_none_uses_client_default(self, mock_response): - """Per-request None should use client default.""" + def test_per_request_none_disables_timeout(self, mock_response): + """Per-request None should disable timeout (no timeout).""" client = DarktraceClient( host="https://test.darktrace.com", public_token="test_public", private_token="test_private", - timeout=30.0, + timeout=30.0, # Client default ) with patch("darktrace.dt_devices.requests.get", return_value=mock_response) as mock_get: - # Explicitly passing None should use client default + # Explicitly passing None should disable timeout client.devices.get(timeout=None) - + call_kwargs = mock_get.call_args[1] - assert call_kwargs.get("timeout") == 30.0 + assert call_kwargs.get("timeout") is None # No timeout + + def test_per_request_unset_uses_client_default(self, mock_response): + """Not passing timeout should use client default.""" + client = DarktraceClient( + host="https://test.darktrace.com", + public_token="test_public", + private_token="test_private", + timeout=30.0, + ) + + with patch("darktrace.dt_devices.requests.get", return_value=mock_response) as mock_get: + client.devices.get() # No timeout argument + + call_kwargs = mock_get.call_args[1] + assert call_kwargs.get("timeout") == 30.0 # Uses client default class TestTimeoutType: