pangea-sdk 3.6.1__py3-none-any.whl → 3.8.0__py3-none-any.whl

pangea/asyncio/services/vault.py

@@ -1,8 +1,12 @@
 # Copyright 2022 Pangea Cyber Corporation
 # Author: Pangea Cyber Corporation
+from __future__ import annotations
+
 import datetime
 from typing import Dict, Optional, Union
 
+from pangea.asyncio.services.base import ServiceBaseAsync
+from pangea.config import PangeaConfig
 from pangea.response import PangeaResponse
 from pangea.services.vault.models.asymmetric import (
     AsymmetricGenerateRequest,
@@ -47,8 +51,8 @@ from pangea.services.vault.models.common import (
     StateChangeRequest,
     StateChangeResult,
     SymmetricAlgorithm,
-    TDict,
     Tags,
+    TDict,
     UpdateRequest,
     UpdateResult,
 )
@@ -69,8 +73,6 @@ from pangea.services.vault.models.symmetric import (
     SymmetricStoreResult,
 )
 
-from .base import ServiceBaseAsync
-
 
 class VaultAsync(ServiceBaseAsync):
     """Vault service client.
@@ -92,17 +94,31 @@ class VaultAsync(ServiceBaseAsync):
         vault_config = PangeaConfig(domain="pangea.cloud")
 
         # Setup Pangea Vault service
-        vault = Vault(token=PANGEA_VAULT_TOKEN, config=audit_config)
+        vault = Vault(token=PANGEA_VAULT_TOKEN, config=vault_config)
     """
 
     service_name = "vault"
 
     def __init__(
         self,
-        token,
-        config=None,
-        logger_name="pangea",
-    ):
+        token: str,
+        config: PangeaConfig | None = None,
+        logger_name: str = "pangea",
+    ) -> None:
+        """
+        Vault client
+
+        Initializes a new Vault client.
+
+        Args:
+            token: Pangea API token.
+            config: Configuration.
+            logger_name: Logger name.
+
+        Examples:
+             config = PangeaConfig(domain="pangea_domain")
+             vault = VaultAsync(token="pangea_token", config=config)
+        """
         super().__init__(token, config, logger_name)
 
     # Delete endpoint
@@ -859,8 +875,8 @@ class VaultAsync(ServiceBaseAsync):
 
                 Default is `deactivated`.
             public_key (EncodedPublicKey, optional): The public key (in PEM format)
-            private_key: (EncodedPrivateKey, optional): The private key (in PEM format)
-            key: (EncodedSymmetricKey, optional): The key material (in base64)
+            private_key (EncodedPrivateKey, optional): The private key (in PEM format)
+            key (EncodedSymmetricKey, optional): The key material (in base64)
 
         Raises:
             PangeaAPIException: If an API Error happens
@@ -1222,7 +1238,7 @@ class VaultAsync(ServiceBaseAsync):
             )
         """
 
-        input = EncryptStructuredRequest(
+        input: EncryptStructuredRequest[TDict] = EncryptStructuredRequest(
             id=id, structured_data=structured_data, filter=filter, version=version, additional_data=additional_data
         )
         return await self.request.post(
@@ -1271,7 +1287,7 @@ class VaultAsync(ServiceBaseAsync):
             )
         """
 
-        input = EncryptStructuredRequest(
+        input: EncryptStructuredRequest[TDict] = EncryptStructuredRequest(
             id=id, structured_data=structured_data, filter=filter, version=version, additional_data=additional_data
         )
         return await self.request.post(

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

@@ -5,16 +5,30 @@ import copy
 import json
 import logging
 import time
-from typing import Dict, List, Optional, Tuple, Type, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type, Union
+
+import requests
+from requests.adapters import HTTPAdapter, Retry
+from requests_toolbelt import MultipartDecoder  # type: ignore
+from typing_extensions import TypeVar
 
-import aiohttp
 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
+
+if TYPE_CHECKING:
+    import aiohttp
+
+
+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 +44,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 +72,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 +125,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 +174,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 +182,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:
@@ -165,6 +190,9 @@ class PangeaRequestBase(object):
         raise pe.PangeaAPIException(f"{summary} ", response)
 
 
+TResult = TypeVar("TResult", bound=PangeaResponseResult, default=PangeaResponseResult)
+
+
 class PangeaRequest(PangeaRequestBase):
     """An object that makes direct calls to Pangea Service APIs.
 
@@ -180,12 +208,12 @@ class PangeaRequest(PangeaRequestBase):
     def post(
         self,
         endpoint: str,
-        result_class: Type[PangeaResponseResult],
+        result_class: Type[TResult],
         data: Union[str, Dict] = {},
         files: Optional[List[Tuple]] = None,
         poll_result: bool = True,
         url: Optional[str] = None,
-    ) -> PangeaResponse:
+    ) -> PangeaResponse[TResult]:
         """Makes the POST call to a Pangea Service endpoint.
 
         Args:
@@ -200,13 +228,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 +246,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 +349,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[TResult]) -> PangeaResponse[TResult]:
+        if self._queued_retry_enabled and response.http_status == 202:
             self.logger.debug(
                 json.dumps(
                     {"service": self.service, "action": "poll_result", "response": response.json},
@@ -320,7 +361,7 @@ class PangeaRequest(PangeaRequestBase):
 
         return response
 
-    def get(self, path: str, result_class: Type[PangeaResponseResult], check_response: bool = True) -> PangeaResponse:
+    def get(self, path: str, result_class: Type[TResult], check_response: bool = True) -> PangeaResponse[TResult]:
         """Makes the GET call to a Pangea Service endpoint.
 
         Args:
@@ -336,7 +377,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,20 +393,61 @@ 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: Type[TResult], check_response: bool = True
+    ) -> PangeaResponse[TResult]:
         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]
+        return self.get(path, result_class, check_response=check_response)
 
-    def poll_result_once(self, response: PangeaResponse, check_response: bool = True):
+    def poll_result_once(
+        self, response: PangeaResponse[TResult], check_response: bool = True
+    ) -> PangeaResponse[TResult]:
         request_id = response.request_id
         if not request_id:
             raise pe.PangeaException("Poll result error: response did not include a 'request_id'")
 
         if response.status != ResponseStatus.ACCEPTED.value:
-            raise pe.PangeaException("Response already proccesed")
+            raise pe.PangeaException("Response already processed")
 
         return self.poll_result_by_id(request_id, response.result_class, check_response=check_response)
 
@@ -422,7 +506,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 +516,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[TResult]) -> PangeaResponse[TResult]:
         retry_count = 1
         start = time.time()
 
@@ -453,7 +543,7 @@ 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[TResult]) -> PangeaResponse[TResult]:
         if response.http_status != 202:
             raise AttributeError("Response should be 202")