pangea-sdk 3.7.0__py3-none-any.whl → 3.7.1__py3-none-any.whl

pangea/config.py

@@ -9,65 +9,57 @@ from typing import Optional
 class PangeaConfig:
     """Holds run time configuration information used by SDK components."""
 
+    domain: str = "aws.us.pangea.cloud"
     """
     Used to set Pangea domain (and port if needed), it should not include service subdomain
     just for particular use cases when environment = "local", domain could be set to an url including:
     scheme (http:// or https://), subdomain, domain and port.
-
     """
-    domain: str = "aws.us.pangea.cloud"
 
+    environment: str = "production"
     """
     Used to generate service url.
     It should be only 'production' or 'local' in cases of particular services that can run locally as Redact.
-
     """
-    environment: str = "production"
 
+    config_id: Optional[str] = None
     """
     Only used for services that support multiconfig (e.g.: Audit service)
 
     @deprecated("config_id will be deprecated from PangeaConfig. Set it on service initialization instead")
     """
-    config_id: Optional[str] = None
 
+    insecure: bool = False
     """
     Set to true to use plain http
-
     """
-    insecure: bool = False
 
+    request_retries: int = 3
     """
     Number of retries on the initial request
-
     """
-    request_retries: int = 3
 
+    request_backoff: float = 0.5
     """
     Backoff strategy passed to 'requests'
-
     """
-    request_backoff: float = 0.5
 
+    request_timeout: int = 5
     """
     Timeout used on initial request attempts
-
     """
-    request_timeout: int = 5
 
+    poll_result_timeout: int = 30
     """
     Timeout used to poll results after 202 (in secs)
-
     """
-    poll_result_timeout: int = 30
 
+    queued_retry_enabled: bool = True
     """
     Enable queued request retry support
     """
-    queued_retry_enabled: bool = True
 
+    custom_user_agent: Optional[str] = None
     """
     Extra user agent to be added to request user agent
-
     """
-    custom_user_agent: Optional[str] = None

pangea/dump_audit.py

@@ -10,6 +10,7 @@ from datetime import datetime
 from typing import Tuple
 
 import dateutil.parser
+
 from pangea.response import PangeaResponse
 from pangea.services import Audit
 from pangea.services.audit.models import SearchEvent, SearchOrder, SearchOrderBy, SearchOutput, SearchResultOutput

pangea/exceptions.py

@@ -30,6 +30,14 @@ class PresignedUploadError(PangeaException):
         self.body = body
 
 
+class DownloadFileError(PangeaException):
+    body: str
+
+    def __init__(self, message: str, body: str):
+        super().__init__(message)
+        self.body = body
+
+
 class PangeaAPIException(PangeaException):
     """Exceptions raised during API calls"""
 

pangea/request.py

@@ -8,13 +8,25 @@ import time
 from typing import Dict, List, Optional, Tuple, Type, Union
 
 import aiohttp
+import requests
+from requests.adapters import HTTPAdapter, Retry
+from requests.structures import CaseInsensitiveDict
+from requests_toolbelt import MultipartDecoder  # type: ignore
+
 import pangea
 import pangea.exceptions as pe
-import requests
 from pangea.config import PangeaConfig
-from pangea.response import PangeaResponse, PangeaResponseResult, ResponseStatus, TransferMethod
+from pangea.response import AttachedFile, PangeaResponse, PangeaResponseResult, ResponseStatus, TransferMethod
 from pangea.utils import default_encoder
-from requests.adapters import HTTPAdapter, Retry
+
+
+class MultipartResponse(object):
+    pangea_json: Dict[str, str]
+    attached_files: List = []
+
+    def __init__(self, pangea_json: Dict[str, str], attached_files: List = []):
+        self.pangea_json = pangea_json
+        self.attached_files = attached_files
 
 
 class PangeaRequestBase(object):
@@ -30,9 +42,10 @@ class PangeaRequestBase(object):
         self._queued_retry_enabled = config.queued_retry_enabled
 
         # Custom headers
-        self._extra_headers = {}  # type: ignore[var-annotated]
+        self._extra_headers: Dict = {}
         self._user_agent = ""
-        self.set_custom_user_agent(config.custom_user_agent)  # type: ignore[arg-type]
+
+        self.set_custom_user_agent(config.custom_user_agent)
         self._session: Optional[Union[requests.Session, aiohttp.ClientSession]] = None
 
         self.logger = logger
@@ -57,7 +70,7 @@ class PangeaRequestBase(object):
         if isinstance(headers, dict):
             self._extra_headers = headers
 
-    def set_custom_user_agent(self, user_agent: str):
+    def set_custom_user_agent(self, user_agent: Optional[str]):
         self.config.custom_user_agent = user_agent
         self._user_agent = f"pangea-python/{pangea.__version__}"
         if self.config.custom_user_agent:
@@ -110,6 +123,16 @@ class PangeaRequestBase(object):
         self._extra_headers.update(headers)
         return self._extra_headers
 
+    def _get_filename_from_content_disposition(self, content_disposition: str) -> Optional[str]:
+        filename_parts = content_disposition.split("name=")
+        if len(filename_parts) > 1:
+            return filename_parts[1].split(";")[0].strip('"')
+        else:
+            return None
+
+    def _get_filename_from_url(self, url: str) -> Optional[str]:
+        return url.split("/")[-1].split("?")[0]
+
     def _check_response(self, response: PangeaResponse) -> PangeaResponse:
         status = response.status
         summary = response.summary
@@ -149,7 +172,7 @@ class PangeaRequestBase(object):
         elif status == ResponseStatus.TREE_NOT_FOUND.value:
             raise pe.TreeNotFoundException(summary, response)
         elif status == ResponseStatus.IP_NOT_FOUND.value:
-            raise pe.IPNotFoundException(summary)  # type: ignore[call-arg]
+            raise pe.IPNotFoundException(summary, response)
         elif status == ResponseStatus.BAD_OFFSET.value:
             raise pe.BadOffsetException(summary, response)
         elif status == ResponseStatus.FORBIDDEN_VAULT_OPERATION.value:
@@ -157,7 +180,7 @@ class PangeaRequestBase(object):
         elif status == ResponseStatus.VAULT_ITEM_NOT_FOUND.value:
             raise pe.VaultItemNotFound(summary, response)
         elif status == ResponseStatus.NOT_FOUND.value:
-            raise pe.NotFound(response.raw_response.url if response.raw_response is not None else "", response)  # type: ignore[arg-type]
+            raise pe.NotFound(str(response.raw_response.url) if response.raw_response is not None else "", response)  # type: ignore[arg-type]
         elif status == ResponseStatus.INTERNAL_SERVER_ERROR.value:
             raise pe.InternalServerError(response)
         elif status == ResponseStatus.ACCEPTED.value:
@@ -200,13 +223,13 @@ class PangeaRequest(PangeaRequestBase):
             url = self._url(endpoint)
 
         # Set config ID if available
-        if self.config_id and data.get("config_id", None) is None:  # type: ignore[union-attr]
-            data["config_id"] = self.config_id  # type: ignore[index]
+        if self.config_id and isinstance(data, dict) and data.get("config_id", None) is None:
+            data["config_id"] = self.config_id
 
         self.logger.debug(
             json.dumps({"service": self.service, "action": "post", "url": url, "data": data}, default=default_encoder)
         )
-        transfer_method = data.get("transfer_method", None)  # type: ignore[union-attr]
+        transfer_method = data.get("transfer_method", None) if isinstance(data, dict) else None
 
         if files is not None and type(data) is dict and (transfer_method == TransferMethod.POST_URL.value):
             requests_response = self._full_post_presigned_url(
@@ -218,15 +241,68 @@ class PangeaRequest(PangeaRequestBase):
             )
 
         self._check_http_errors(requests_response)
-        json_resp = requests_response.json()
-        self.logger.debug(json.dumps({"service": self.service, "action": "post", "url": url, "response": json_resp}))
 
-        pangea_response = PangeaResponse(requests_response, result_class=result_class, json=json_resp)  # type: ignore[var-annotated]
+        if "multipart/form-data" in requests_response.headers.get("content-type", ""):
+            multipart_response = self._process_multipart_response(requests_response)
+            pangea_response: PangeaResponse = PangeaResponse(
+                requests_response,
+                result_class=result_class,
+                json=multipart_response.pangea_json,
+                attached_files=multipart_response.attached_files,
+            )
+        else:
+            try:
+                json_resp = requests_response.json()
+                self.logger.debug(
+                    json.dumps({"service": self.service, "action": "post", "url": url, "response": json_resp})
+                )
+
+                pangea_response = 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}")
+
         if poll_result:
             pangea_response = self._handle_queued_result(pangea_response)
 
         return self._check_response(pangea_response)
 
+    def _get_pangea_json(self, decoder: MultipartDecoder) -> Optional[Dict]:
+        # Iterate through parts
+        for i, part in enumerate(decoder.parts):
+            if i == 0:
+                json_str = part.content.decode("utf-8")
+                return json.loads(json_str)
+
+        return None
+
+    def _get_attached_files(self, decoder: MultipartDecoder) -> List[AttachedFile]:
+        files = []
+
+        for i, part in enumerate(decoder.parts):
+            content_type = part.headers.get(b"Content-Type", b"").decode("utf-8")
+            # if "application/octet-stream" in content_type:
+            if i > 0:
+                content_disposition = part.headers.get(b"Content-Disposition", b"").decode("utf-8")
+                name = self._get_filename_from_content_disposition(content_disposition)
+                if name is None:
+                    name = f"default_file_name_{i}"
+
+                files.append(AttachedFile(name, part.content, content_type))
+
+        return files
+
+    def _process_multipart_response(self, resp: requests.Response) -> MultipartResponse:
+        # Parse the multipart response
+        decoder = MultipartDecoder.from_response(resp)
+
+        pangea_json = self._get_pangea_json(decoder)
+        self.logger.debug(
+            json.dumps({"service": self.service, "action": "multipart response", "response": pangea_json})
+        )
+
+        attached_files = self._get_attached_files(decoder)
+        return MultipartResponse(pangea_json, attached_files)  # type: ignore
+
     def _check_http_errors(self, resp: requests.Response):
         if resp.status_code == 503:
             raise pe.ServiceTemporarilyUnavailable(resp.json())
@@ -268,48 +344,8 @@ class PangeaRequest(PangeaRequestBase):
 
         return data, files
 
-    def _post_presigned_url(
-        self,
-        endpoint: str,
-        result_class: Type[PangeaResponseResult],
-        data: Union[str, Dict] = {},
-        files: Optional[List[Tuple]] = None,
-    ):
-        if len(files) == 0:  # type: ignore[arg-type]
-            raise AttributeError("files attribute should have at least 1 file")
-
-        # Send request
-        try:
-            # This should return 202 (AcceptedRequestException)
-            resp = self.post(endpoint=endpoint, result_class=result_class, data=data, poll_result=False)
-            raise pe.PresignedURLException("Should return 202", resp)
-
-        except pe.AcceptedRequestException as e:
-            accepted_exception = e
-        except Exception as e:
-            raise e
-
-        # Receive 202 with accepted_status
-        result = self._poll_presigned_url(accepted_exception)  # type: ignore[arg-type]
-        data_to_presigned = result.accepted_status.upload_details  # type: ignore[attr-defined]
-        presigned_url = result.accepted_status.upload_url  # type: ignore[attr-defined]
-
-        # Send multipart request with file and upload_details as body
-        resp = self._http_post(url=presigned_url, data=data_to_presigned, files=files, multipart_post=False)  # type: ignore[assignment]
-        self.logger.debug(
-            json.dumps(
-                {"service": self.service, "action": "post presigned", "url": presigned_url, "response": resp.text},  # type: ignore[attr-defined]
-                default=default_encoder,
-            )
-        )
-
-        if resp.status_code < 200 or resp.status_code >= 300:  # type: ignore[attr-defined]
-            raise pe.PresignedUploadError(f"presigned POST failure: {resp.status_code}", resp.text)  # type: ignore[attr-defined]
-
-        return accepted_exception.response.raw_response
-
-    def _handle_queued_result(self, response: PangeaResponse) -> PangeaResponse:
-        if self._queued_retry_enabled and response.raw_response.status_code == 202:  # type: ignore[union-attr]
+    def _handle_queued_result(self, response: PangeaResponse) -> PangeaResponse[Type[PangeaResponseResult]]:
+        if self._queued_retry_enabled and response.http_status == 202:
             self.logger.debug(
                 json.dumps(
                     {"service": self.service, "action": "poll_result", "response": response.json},
@@ -336,7 +372,9 @@ class PangeaRequest(PangeaRequestBase):
         self.logger.debug(json.dumps({"service": self.service, "action": "get", "url": url}))
         requests_response = self.session.get(url, headers=self._headers())
         self._check_http_errors(requests_response)
-        pangea_response = PangeaResponse(requests_response, result_class=result_class, json=requests_response.json())  # type: ignore[var-annotated]
+        pangea_response: PangeaResponse = PangeaResponse(
+            requests_response, result_class=result_class, json=requests_response.json()
+        )
 
         self.logger.debug(
             json.dumps(
@@ -350,14 +388,55 @@ class PangeaRequest(PangeaRequestBase):
 
         return self._check_response(pangea_response)
 
+    def download_file(self, url: str, filename: Optional[str] = None) -> AttachedFile:
+        self.logger.debug(
+            json.dumps(
+                {
+                    "service": self.service,
+                    "action": "download_file",
+                    "url": url,
+                    "filename": filename,
+                    "status": "start",
+                }
+            )
+        )
+        response = self.session.get(url, headers={})
+        if response.status_code == 200:
+            if filename is None:
+                content_disposition = response.headers.get(b"Content-Disposition", b"").decode("utf-8")
+                filename = self._get_filename_from_content_disposition(content_disposition)
+                if filename is None:
+                    filename = self._get_filename_from_url(url)
+                    if filename is None:
+                        filename = "default_filename"
+
+            content_type = response.headers.get(b"Content-Type", b"").decode("utf-8")
+
+            self.logger.debug(
+                json.dumps(
+                    {
+                        "service": self.service,
+                        "action": "download_file",
+                        "url": url,
+                        "filename": filename,
+                        "status": "success",
+                    }
+                )
+            )
+            return AttachedFile(filename=filename, file=response.content, content_type=content_type)
+        else:
+            raise pe.DownloadFileError(f"Failed to download file. Status: {response.status_code}", response.text)
+
     def poll_result_by_id(
-        self, request_id: str, result_class: Union[Type[PangeaResponseResult], dict], check_response: bool = True
+        self, request_id: str, result_class: Union[Type[PangeaResponseResult], Type[dict]], check_response: bool = True
     ):
         path = self._get_poll_path(request_id)
         self.logger.debug(json.dumps({"service": self.service, "action": "poll_result_once", "url": path}))
         return self.get(path, result_class, check_response=check_response)  # type: ignore[arg-type]
 
-    def poll_result_once(self, response: PangeaResponse, check_response: bool = True):
+    def poll_result_once(
+        self, response: PangeaResponse, check_response: bool = True
+    ) -> PangeaResponse[Type[PangeaResponseResult]]:
         request_id = response.request_id
         if not request_id:
             raise pe.PangeaException("Poll result error: response did not include a 'request_id'")
@@ -422,7 +501,8 @@ class PangeaRequest(PangeaRequestBase):
         self.logger.debug(
             json.dumps({"service": self.service, "action": "http_put", "url": url}, default=default_encoder)
         )
-        return self.session.put(url, headers=headers, files=files)
+        _, value = files[0]
+        return self.session.put(url, headers=headers, data=value[1])
 
     def _full_post_presigned_url(
         self,
@@ -431,17 +511,22 @@ class PangeaRequest(PangeaRequestBase):
         data: Union[str, Dict] = {},
         files: Optional[List[Tuple]] = None,
     ):
-        if len(files) == 0:  # type: ignore[arg-type]
+        if files is None or len(files) == 0:
             raise AttributeError("files attribute should have at least 1 file")
 
         response = self.request_presigned_url(endpoint=endpoint, result_class=result_class, data=data)
-        data_to_presigned = response.accepted_result.post_form_data  # type: ignore[union-attr]
-        presigned_url = response.accepted_result.post_url  # type: ignore[union-attr]
+        if response.accepted_result is None:
+            raise pe.PangeaException("No accepted_result field when requesting presigned url")
+        if response.accepted_result.post_url is None:
+            raise pe.PresignedURLException("No presigned url", response)
+
+        data_to_presigned = response.accepted_result.post_form_data
+        presigned_url = response.accepted_result.post_url
 
-        self.post_presigned_url(url=presigned_url, data=data_to_presigned, files=files)  # type: ignore[arg-type]
+        self.post_presigned_url(url=presigned_url, data=data_to_presigned, files=files)
         return response.raw_response
 
-    def _poll_result_retry(self, response: PangeaResponse) -> PangeaResponse:
+    def _poll_result_retry(self, response: PangeaResponse) -> PangeaResponse[Type[PangeaResponseResult]]:
         retry_count = 1
         start = time.time()
 
@@ -453,7 +538,9 @@ class PangeaRequest(PangeaRequestBase):
         self.logger.debug(json.dumps({"service": self.service, "action": "poll_result_retry", "step": "exit"}))
         return self._check_response(response)
 
-    def _poll_presigned_url(self, response: PangeaResponse) -> PangeaResponse:
+    def _poll_presigned_url(
+        self, response: PangeaResponse[Type[PangeaResponseResult]]
+    ) -> PangeaResponse[Type[PangeaResponseResult]]:
         if response.http_status != 202:
             raise AttributeError("Response should be 202")
 

pangea/response.py

@@ -2,21 +2,47 @@
 # Author: Pangea Cyber Corporation
 import datetime
 import enum
+import os
 from typing import Any, Dict, Generic, List, Optional, Type, TypeVar, Union
 
 import aiohttp
 import requests
-from pangea.utils import format_datetime
 from pydantic import BaseModel
 
+from pangea.utils import format_datetime
+
 T = TypeVar("T")
 
 
+class AttachedFile(object):
+    filename: str
+    file: bytes
+    content_type: str
+
+    def __init__(self, filename: str, file: bytes, content_type: str):
+        self.filename = filename
+        self.file = file
+        self.content_type = content_type
+
+    def save(self, dest_folder: str = "./", filename: Optional[str] = None):
+        if filename is None:
+            filename = self.filename if self.filename else "default_save_filename"
+
+        filepath = os.path.join(dest_folder, filename)
+        directory = os.path.dirname(filepath)
+        if not os.path.exists(directory):
+            os.makedirs(directory)
+
+        with open(filepath, "wb") as file:
+            file.write(self.file)
+
+
 class TransferMethod(str, enum.Enum):
     MULTIPART = "multipart"
     POST_URL = "post-url"
     PUT_URL = "put-url"
     SOURCE_URL = "source-url"
+    DEST_URL = "dest-url"
 
     def __str__(self):
         return str(self.value)
@@ -141,15 +167,24 @@ class PangeaResponse(Generic[T], ResponseHeader):
     result: Optional[T] = None
     pangea_error: Optional[PangeaError] = None
     accepted_result: Optional[AcceptedResult] = None
-    result_class: Type[PangeaResponseResult] = PangeaResponseResult
+    result_class: Union[Type[PangeaResponseResult], Type[dict]] = PangeaResponseResult
     _json: Any
-
-    def __init__(self, response: requests.Response, result_class: Type[PangeaResponseResult], json: dict):
+    attached_files: List[AttachedFile] = []
+
+    def __init__(
+        self,
+        response: requests.Response,
+        result_class: Union[Type[PangeaResponseResult], Type[dict]],
+        json: dict,
+        attached_files: List[AttachedFile] = [],
+    ):
         super(PangeaResponse, self).__init__(**json)
         self._json = json
         self.raw_response = response
         self.raw_result = self._json["result"]
         self.result_class = result_class
+        self.attached_files = attached_files
+
         self.result = (
             self.result_class(**self.raw_result)  # type: ignore[assignment]
             if self.raw_result is not None and issubclass(self.result_class, PangeaResponseResult) and self.success