pangea-sdk 6.1.1__py3-none-any.whl → 6.2.0b2__py3-none-any.whl

pangea/request.py

@@ -1,19 +1,24 @@
 # Copyright 2022 Pangea Cyber Corporation
 # Author: Pangea Cyber Corporation
+
+# TODO: Modernize.
+# ruff: noqa: UP006, UP035
+
 from __future__ import annotations
 
 import copy
 import json
 import logging
 import time
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, Type, Union, cast
+from collections.abc import Iterable, Mapping
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union, cast, overload
 
 import requests
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter
 from pydantic_core import to_jsonable_python
 from requests.adapters import HTTPAdapter, Retry
 from requests_toolbelt import MultipartDecoder  # type: ignore[import-untyped]
-from typing_extensions import TypeVar
+from typing_extensions import Literal, TypeAlias, TypeVar, override
 from yarl import URL
 
 import pangea
@@ -26,11 +31,24 @@ if TYPE_CHECKING:
     import aiohttp
 
 
+_Data: TypeAlias = Union[Iterable[bytes], str, bytes, list[tuple[Any, Any]]]
+
+_FileName: TypeAlias = Union[str, None]
+_FileContent: TypeAlias = Union[str, bytes]
+_FileContentType: TypeAlias = str
+_FileCustomHeaders: TypeAlias = Mapping[str, str]
+_FileSpecTuple2: TypeAlias = tuple[_FileName, _FileContent]
+_FileSpecTuple3: TypeAlias = tuple[_FileName, _FileContent, _FileContentType]
+_FileSpecTuple4: TypeAlias = tuple[_FileName, _FileContent, _FileContentType, _FileCustomHeaders]
+_FileSpec: TypeAlias = Union[_FileContent, _FileSpecTuple2, _FileSpecTuple3, _FileSpecTuple4]
+_Files: TypeAlias = Union[Mapping[str, _FileSpec], Iterable[tuple[str, _FileSpec]]]
+
+
 class MultipartResponse:
     pangea_json: Dict[str, str]
     attached_files: List = []
 
-    def __init__(self, pangea_json: Dict[str, str], attached_files: List = []):
+    def __init__(self, pangea_json: dict[str, str], attached_files: list = []):  # noqa: B006
         self.pangea_json = pangea_json
         self.attached_files = attached_files
 
@@ -115,6 +133,7 @@ class PangeaRequestBase:
         headers = {
             "User-Agent": self._user_agent,
             "Authorization": f"Bearer {self.token}",
+            "Content-Type": "application/json",
         }
 
         # We want to ignore previous headers if user tried to set them, so we will overwrite them.
@@ -184,6 +203,9 @@ class PangeaRequestBase:
             raise pe.AcceptedRequestException(response)
         raise pe.PangeaAPIException(f"{summary} ", response)
 
+    def _init_session(self) -> requests.Session | aiohttp.ClientSession:
+        raise NotImplementedError
+
 
 TResult = TypeVar("TResult", bound=PangeaResponseResult)
 
@@ -200,26 +222,88 @@ class PangeaRequest(PangeaRequestBase):
     def __del__(self) -> None:
         self.session.close()
 
+    def delete(self, endpoint: str) -> None:
+        """
+        Makes a DELETE call to a Pangea endpoint.
+
+        Args:
+            endpoint: The Pangea API endpoint.
+        """
+
+        url = self._url(endpoint)
+
+        self.logger.debug(
+            json.dumps({"service": self.service, "action": "delete", "url": url}, default=default_encoder)
+        )
+
+        requests_response = self._http_delete(url, headers=self._headers())
+        self._check_http_errors(requests_response)
+
+    @overload
     def post(
         self,
         endpoint: str,
         result_class: Type[TResult],
-        data: str | BaseModel | dict[str, Any] | None = None,
-        files: Optional[List[Tuple]] = None,
+        data: str | BaseModel | Mapping[str, Any] | None = None,
+        files: Optional[list[Tuple]] = None,
         poll_result: bool = True,
         url: Optional[str] = None,
+        *,
+        pangea_response: Literal[True] = True,
     ) -> PangeaResponse[TResult]:
-        """Makes the POST call to a Pangea Service endpoint.
+        """
+        Makes the POST call to a Pangea Service endpoint.
 
         Args:
-            endpoint(str): The Pangea Service API endpoint.
-            data(dict): The POST body payload object
+            endpoint: The Pangea Service API endpoint.
+            data: The POST body payload object
 
         Returns:
             PangeaResponse which contains the response in its entirety and
                various properties to retrieve individual fields
         """
 
+    @overload
+    def post(
+        self,
+        endpoint: str,
+        result_class: Type[TResult],
+        data: str | BaseModel | Mapping[str, Any] | None = None,
+        files: Optional[list[Tuple]] = None,
+        poll_result: bool = True,
+        url: Optional[str] = None,
+        *,
+        pangea_response: Literal[False],
+    ) -> TResult:
+        """
+        Makes the POST call to a Pangea Service endpoint.
+
+        Args:
+            endpoint: The Pangea Service API endpoint.
+            data: The POST body payload object
+        """
+
+    def post(
+        self,
+        endpoint: str,
+        result_class: Type[TResult],
+        data: str | BaseModel | Mapping[str, Any] | None = None,
+        files: Optional[list[Tuple]] = None,
+        poll_result: bool = True,
+        url: Optional[str] = None,
+        *,
+        pangea_response: bool = True,
+    ) -> PangeaResponse[TResult] | TResult:
+        """
+        Makes a POST call to a Pangea Service endpoint.
+
+        Args:
+            endpoint: The Pangea Service API endpoint.
+            data: The POST body payload object
+            pangea_response: Whether or not the response body follows Pangea's
+              standard response schema
+        """
+
         if isinstance(data, BaseModel):
             data = data.model_dump(exclude_none=True)
 
@@ -250,15 +334,20 @@ class PangeaRequest(PangeaRequestBase):
                 endpoint, result_class=result_class, data=data, files=files
             )
         else:
-            requests_response = self._http_post(
-                url, headers=self._headers(), data=data, files=files, multipart_post=True
-            )
+            headers = self._headers()
+            if transfer_method == TransferMethod.MULTIPART.value:
+                del headers["Content-Type"]
+            requests_response = self._http_post(url, headers=headers, data=data, files=files)
 
         self._check_http_errors(requests_response)
 
+        if not pangea_response:
+            type_adapter = TypeAdapter(result_class)
+            return type_adapter.validate_python(requests_response.json())
+
         if "multipart/form-data" in requests_response.headers.get("content-type", ""):
             multipart_response = self._process_multipart_response(requests_response)
-            pangea_response: PangeaResponse = PangeaResponse(
+            pangea_response_obj: PangeaResponse = PangeaResponse(
                 requests_response,
                 result_class=result_class,
                 json=multipart_response.pangea_json,
@@ -271,14 +360,14 @@ class PangeaRequest(PangeaRequestBase):
                     json.dumps({"service": self.service, "action": "post", "url": url, "response": json_resp})
                 )
 
-                pangea_response = PangeaResponse(requests_response, result_class=result_class, json=json_resp)
+                pangea_response_obj = PangeaResponse(requests_response, result_class=result_class, json=json_resp)
             except requests.exceptions.JSONDecodeError as e:
-                raise pe.PangeaException(f"Failed to decode json response. {e}. Body: {requests_response.text}")
+                raise pe.PangeaException(f"Failed to decode json response. {e}. Body: {requests_response.text}") from e
 
         if poll_result:
-            pangea_response = self._handle_queued_result(pangea_response)
+            pangea_response_obj = self._handle_queued_result(pangea_response_obj)
 
-        return self._check_response(pangea_response)
+        return self._check_response(pangea_response_obj)
 
     def _get_pangea_json(self, decoder: MultipartDecoder) -> Optional[Dict]:
         # Iterate through parts
@@ -321,12 +410,20 @@ class PangeaRequest(PangeaRequestBase):
         if resp.status_code == 503:
             raise pe.ServiceTemporarilyUnavailable(resp.json())
 
+    def _http_delete(
+        self,
+        url: str,
+        *,
+        headers: Mapping[str, str | bytes | None] = {},
+    ) -> requests.Response:
+        return self.session.delete(url, headers=headers)
+
     def _http_post(
         self,
         url: str,
-        headers: Dict = {},
-        data: Union[str, Dict] = {},
-        files: Optional[List[Tuple]] = None,
+        headers: Mapping[str, str] = {},
+        data: str | dict[Any, Any] = {},  # noqa: B006
+        files: _Files | None = None,
         multipart_post: bool = True,
     ) -> requests.Response:
         data_send, files = self._http_post_process(data=data, files=files, multipart_post=multipart_post)
@@ -334,31 +431,25 @@ class PangeaRequest(PangeaRequestBase):
 
     def _http_post_process(
         self,
-        data: Union[str, Dict] = {},
-        files: Optional[Sequence[Tuple[str, Tuple[Any, str, str]]]] = None,
+        data: str | dict[Any, Any] = {},  # noqa: B006
+        files: _Files | None = None,
         multipart_post: bool = True,
-    ):
+    ) -> tuple[_Data | None, _Files | None]:
         if files:
             if multipart_post is True:
                 data_send: str = json.dumps(data, default=default_encoder) if isinstance(data, dict) else data
-                multi = [("request", (None, data_send, "application/json"))]
-                multi.extend(files)
-                files = multi
-                return None, files
-            # Post to presigned url as form
-            data_send: list = []  # type: ignore[no-redef]
-            for k, v in data.items():  # type: ignore[union-attr]
-                data_send.append((k, v))  # type: ignore[attr-defined]
-            # When posting to presigned url, file key should be 'file'
-            files = {  # type: ignore[assignment]
-                "file": files[0][1],
-            }
-            return data_send, files
+                multi: list[tuple[str, _FileSpec]] = [("request", (None, data_send, "application/json")), *files]
+                return None, multi
+
+            # Post to presigned URL as form.
+            # When posting to presigned URL, file key should be 'file'.
+            assert isinstance(data, dict)
+            assert isinstance(files, list)
+            return [(k, v) for k, v in data.items()], {"file": files[0][1]}
+
         data_send = json.dumps(data, default=default_encoder) if isinstance(data, dict) else data
         return data_send, None
 
-        return data, files
-
     def _handle_queued_result(self, response: PangeaResponse[TResult]) -> PangeaResponse[TResult]:
         if self._queued_retry_enabled and response.http_status == 202:
             self.logger.debug(
@@ -371,37 +462,98 @@ class PangeaRequest(PangeaRequestBase):
 
         return response
 
-    def get(self, path: str, result_class: Type[TResult], check_response: bool = True) -> PangeaResponse[TResult]:
-        """Makes the GET call to a Pangea Service endpoint.
+    @overload
+    def get(
+        self,
+        path: str,
+        result_class: Type[TResult],
+        check_response: bool = True,
+        *,
+        params: (
+            Mapping[str | bytes | int | float, str | bytes | int | float | Iterable[str | bytes | int | float] | None]
+            | None
+        ) = None,
+        pangea_response: Literal[True] = True,
+    ) -> PangeaResponse[TResult]:
+        """
+        Makes the GET call to a Pangea Service endpoint.
 
         Args:
-            endpoint(str): The Pangea Service API endpoint.
-            path(str): Additional URL path
+            path: Additional URL path
+            params: Dictionary of querystring data to attach to the request
 
         Returns:
             PangeaResponse which contains the response in its entirety and
                various properties to retrieve individual fields
         """
 
+    @overload
+    def get(
+        self,
+        path: str,
+        result_class: Type[TResult],
+        check_response: bool = True,
+        *,
+        params: (
+            Mapping[str | bytes | int | float, str | bytes | int | float | Iterable[str | bytes | int | float] | None]
+            | None
+        ) = None,
+        pangea_response: Literal[False] = False,
+    ) -> TResult:
+        """
+        Makes the GET call to a Pangea Service endpoint.
+
+        Args:
+            path: Additional URL path
+            params: Dictionary of querystring data to attach to the request
+        """
+
+    def get(
+        self,
+        path: str,
+        result_class: Type[TResult],
+        check_response: bool = True,
+        *,
+        params: (
+            Mapping[str | bytes | int | float, str | bytes | int | float | Iterable[str | bytes | int | float] | None]
+            | None
+        ) = None,
+        pangea_response: bool = True,
+    ) -> PangeaResponse[TResult] | TResult:
+        """
+        Makes the GET call to a Pangea Service endpoint.
+
+        Args:
+            path: Additional URL path
+            params: Dictionary of querystring data to attach to the request
+            pangea_response: Whether or not the response body follows Pangea's
+              standard response schema
+        """
+
         url = self._url(path)
         self.logger.debug(json.dumps({"service": self.service, "action": "get", "url": url}))
-        requests_response = self.session.get(url, headers=self._headers())
+        requests_response = self.session.get(url, params=params, headers=self._headers())
         self._check_http_errors(requests_response)
-        pangea_response: PangeaResponse = PangeaResponse(
+
+        if not pangea_response:
+            type_adapter = TypeAdapter(result_class)
+            return type_adapter.validate_python(requests_response.json())
+
+        pangea_response_obj: PangeaResponse = PangeaResponse(
             requests_response, result_class=result_class, json=requests_response.json()
         )
 
         self.logger.debug(
             json.dumps(
-                {"service": self.service, "action": "get", "url": url, "response": pangea_response.json},
+                {"service": self.service, "action": "get", "url": url, "response": pangea_response_obj.json},
                 default=default_encoder,
             )
         )
 
         if check_response is False:
-            return pangea_response
+            return pangea_response_obj
 
-        return self._check_response(pangea_response)
+        return self._check_response(pangea_response_obj)
 
     def download_file(self, url: str, filename: str | None = None) -> AttachedFile:
         """
@@ -477,7 +629,7 @@ class PangeaRequest(PangeaRequestBase):
         self,
         endpoint: str,
         result_class: Type[PangeaResponseResult],
-        data: Union[str, Dict] = {},
+        data: Union[str, Mapping[str, Any]] = {},
     ) -> PangeaResponse:
         # Send request
         try:
@@ -520,8 +672,8 @@ class PangeaRequest(PangeaRequestBase):
     def _http_put(
         self,
         url: str,
-        files: List[Tuple],
-        headers: Dict = {},
+        files: list[Tuple],
+        headers: Mapping[str, str] = {},
     ) -> requests.Response:
         self.logger.debug(
             json.dumps({"service": self.service, "action": "http_put", "url": url}, default=default_encoder)
@@ -533,7 +685,7 @@ class PangeaRequest(PangeaRequestBase):
         self,
         endpoint: str,
         result_class: Type[PangeaResponseResult],
-        data: Union[str, Dict] = {},
+        data: Union[str, Mapping[str, Any]] = {},
         files: Optional[List[Tuple]] = None,
     ):
         if files is None or len(files) == 0:
@@ -603,7 +755,7 @@ class PangeaRequest(PangeaRequestBase):
                         {"service": self.service, "action": "poll_presigned_url", "step": "exit", "cause": {str(e)}}
                     )
                 )
-                raise pe.PresignedURLException("Failed to pull Presigned URL", loop_resp, e)
+                raise pe.PresignedURLException("Failed to pull Presigned URL", loop_resp, e) from e
 
         self.logger.debug(json.dumps({"service": self.service, "action": "poll_presigned_url", "step": "exit"}))
 
@@ -611,6 +763,7 @@ class PangeaRequest(PangeaRequestBase):
             return loop_resp
         raise loop_exc
 
+    @override
     def _init_session(self) -> requests.Session:
         retry_config = Retry(
             total=self.config.request_retries,

pangea/response.py

@@ -1,19 +1,25 @@
 # Copyright 2022 Pangea Cyber Corporation
 # Author: Pangea Cyber Corporation
+
+# TODO: Modernize.
+# ruff: noqa: UP006, UP035
+
+from __future__ import annotations
+
 import datetime
 import enum
 import os
-from typing import Any, Dict, Generic, List, Optional, Type, Union
+from typing import Annotated, Any, Dict, Generic, List, Optional, Type, Union
 
 import aiohttp
 import requests
 from pydantic import BaseModel, ConfigDict, PlainSerializer
-from typing_extensions import Annotated, TypeVar
+from typing_extensions import TypeVar
 
 from pangea.utils import format_datetime
 
 
-class AttachedFile(object):
+class AttachedFile:
     filename: str
     file: bytes
     content_type: str
@@ -40,10 +46,7 @@ class AttachedFile(object):
         base_name, ext = os.path.splitext(file_path)
         counter = 1
         while os.path.exists(file_path):
-            if ext:
-                file_path = f"{base_name}_{counter}{ext}"
-            else:
-                file_path = f"{base_name}_{counter}"
+            file_path = f"{base_name}_{counter}{ext}" if ext else f"{base_name}_{counter}"
             counter += 1
         return file_path
 
@@ -199,16 +202,16 @@ class PangeaResponse(ResponseHeader, Generic[T]):
     accepted_result: Optional[AcceptedResult] = None
     result_class: Type[T] = PangeaResponseResult  # type: ignore[assignment]
     _json: Any
-    attached_files: List[AttachedFile] = []
+    attached_files: list[AttachedFile] = []
 
     def __init__(
         self,
         response: requests.Response,
         result_class: Type[T],
         json: dict,
-        attached_files: List[AttachedFile] = [],
+        attached_files: list[AttachedFile] = [],  # noqa: B006
     ):
-        super(PangeaResponse, self).__init__(**json)
+        super().__init__(**json)
         self._json = json
         self.raw_response = response
         self.raw_result = self._json["result"]
@@ -241,10 +244,10 @@ class PangeaResponse(ResponseHeader, Generic[T]):
     @property
     def http_status(self) -> int:  # type: ignore[return]
         if self.raw_response:
-            if type(self.raw_response) == aiohttp.ClientResponse:
+            if isinstance(self.raw_response, aiohttp.ClientResponse):
                 return self.raw_response.status
             else:
-                return self.raw_response.status_code  # type: ignore[union-attr]
+                return self.raw_response.status_code
 
     @property
     def url(self) -> str:

pangea/services/__init__.py

@@ -1,3 +1,5 @@
+# ruff: noqa: F401
+
 from .ai_guard import AIGuard
 from .audit.audit import Audit
 from .authn.authn import AuthN
@@ -5,6 +7,7 @@ from .authz import AuthZ
 from .embargo import Embargo
 from .file_scan import FileScan
 from .intel import DomainIntel, FileIntel, IpIntel, UrlIntel, UserIntel
+from .management import Management
 from .prompt_guard import PromptGuard
 from .redact import Redact
 from .sanitize import Sanitize