diff --git a/pyproject.toml b/pyproject.toml index 94797bb..5b38b8d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "pyhuntress" -version = "0.2.16" +version = "0.2.17" authors = [ { name="Peter Annabel", email="peter.annabel@gmail.com" }, ] diff --git a/src/pyhuntress/clients/huntress_client.py b/src/pyhuntress/clients/huntress_client.py index 0254cb1..f2b094d 100644 --- a/src/pyhuntress/clients/huntress_client.py +++ b/src/pyhuntress/clients/huntress_client.py @@ -21,6 +21,7 @@ from pyhuntress.exceptions import ( PermissionsFailedException, ServerError, TooManyRequestsException, + RequestTimedOutException, ) if TYPE_CHECKING: @@ -77,6 +78,7 @@ class HuntressClient(ABC): method, url, headers=headers, + timeout=180, data=data, params=cast(dict[str, Any], params or {}), stream=stream, @@ -86,6 +88,7 @@ class HuntressClient(ABC): method, url, headers=headers, + timeout=180, params=cast(dict[str, Any], params or {}), stream=stream, ) @@ -130,5 +133,7 @@ class HuntressClient(ABC): return self._make_request(method, url, data, params, headers, retry_count) raise Timeout(response=response) raise ServerError(response) + if response.status_code == 504: + raise RequestTimedOutException(response) return response diff --git a/src/pyhuntress/exceptions.py b/src/pyhuntress/exceptions.py index 4f81af3..04421f7 100644 --- a/src/pyhuntress/exceptions.py +++ b/src/pyhuntress/exceptions.py @@ -87,3 +87,7 @@ class ServerError(HuntressException): class ObjectExistsError(HuntressException): _code_explanation = "Object Exists" _error_suggestion = "This resource already exists." + +class RequestTimedOutException(HuntressException): + _code_explanation = "Request Timed Out" + _error_suggestion = "This resource is possibly busy or not responding. Try again later or use a smaller request." diff --git a/src/pyhuntress/models/managedsat/__init__.py b/src/pyhuntress/models/managedsat/__init__.py index c25a72f..d90a02d 100644 --- a/src/pyhuntress/models/managedsat/__init__.py +++ b/src/pyhuntress/models/managedsat/__init__.py @@ -120,7 +120,7 @@ class SATPhishingCampaigns(HuntressModel): ] | None = Field(default=None, alias="Type") id: str | None = Field(default=None, alias="Id") attributes: dict[str, Any] | None = Field(default=None, alias="Attributes") - relationships: dict[str, dict[str, dict[str, str]]] | None = Field(default=None, alias="Relationships") + relationships: dict[str, dict[str, Any]] | None = Field(default=None, alias="Relationships") links: dict[str, str] | None = Field(default=None, alias="Links") meta: dict[str, dict[str, int]] | None = Field(default=None, alias="Meta") @@ -131,7 +131,7 @@ class SATPhishingScenarios(HuntressModel): ] | None = Field(default=None, alias="Type") id: str | None = Field(default=None, alias="Id") attributes: dict[str, Any] | None = Field(default=None, alias="Attributes") - relationships: dict[str, dict[str, dict[str, str]]] | None = Field(default=None, alias="Relationships") + relationships: dict[str, dict[str, dict[str, Any]]] | None = Field(default=None, alias="Relationships") links: dict[str, str] | None = Field(default=None, alias="Links") meta: dict[str, dict[str, int]] | None = Field(default=None, alias="Meta")