pangea-sdk 3.8.0b1__py3-none-any.whl → 5.4.0b1__py3-none-any.whl

pangea/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "3.8.0beta1"
+__version__ = "5.4.0beta1"
 
 from pangea.asyncio.request import PangeaRequestAsync
 from pangea.config import PangeaConfig

pangea/asyncio/file_uploader.py

@@ -28,7 +28,7 @@ class FileUploaderAsync:
     ) -> None:
         if transfer_method == TransferMethod.PUT_URL:
             files = [("file", ("filename", file, "application/octet-stream"))]
-            await self._request.put_presigned_url(url=url, files=files)  # type: ignore[arg-type]
+            await self._request.put_presigned_url(url=url, files=files)
         elif transfer_method == TransferMethod.POST_URL:
             files = [("file", ("filename", file, "application/octet-stream"))]
             await self._request.post_presigned_url(url=url, data=file_details, files=files)  # type: ignore[arg-type]

pangea/asyncio/request.py

@@ -1,20 +1,25 @@
 # Copyright 2022 Pangea Cyber Corporation
 # Author: Pangea Cyber Corporation
+from __future__ import annotations
 
 import asyncio
 import json
-import os
 import time
-from typing import Dict, List, Optional, Tuple, Type, Union
+from typing import Dict, List, Optional, Sequence, Tuple, Type, Union, cast
 
 import aiohttp
 from aiohttp import FormData
+from pydantic import BaseModel
+from pydantic_core import to_jsonable_python
+from typing_extensions import Any, TypeVar
 
 import pangea.exceptions as pe
 from pangea.request import MultipartResponse, PangeaRequestBase
 from pangea.response import AttachedFile, PangeaResponse, PangeaResponseResult, ResponseStatus, TransferMethod
 from pangea.utils import default_encoder
 
+TResult = TypeVar("TResult", bound=PangeaResponseResult)
+
 
 class PangeaRequestAsync(PangeaRequestBase):
     """An object that makes direct calls to Pangea Service APIs.
@@ -28,12 +33,12 @@ class PangeaRequestAsync(PangeaRequestBase):
     async def post(
         self,
         endpoint: str,
-        result_class: Type[PangeaResponseResult],
-        data: Union[str, Dict] = {},
-        files: List[Tuple] = [],
+        result_class: Type[TResult],
+        data: str | BaseModel | dict[str, Any] | None = None,
+        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:
@@ -44,17 +49,27 @@ class PangeaRequestAsync(PangeaRequestBase):
             PangeaResponse which contains the response in its entirety and
                various properties to retrieve individual fields
         """
+
+        if isinstance(data, BaseModel):
+            data = data.model_dump(exclude_none=True)
+
+        if data is None:
+            data = {}
+
+        # Normalize.
+        data = cast(dict[str, Any], to_jsonable_python(data))
+
         if url is None:
             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 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 files and type(data) is dict and (transfer_method == TransferMethod.POST_URL.value):
             requests_response = await self._full_post_presigned_url(
@@ -91,9 +106,7 @@ class PangeaRequestAsync(PangeaRequestBase):
 
         return self._check_response(pangea_response)
 
-    async def get(
-        self, path: str, result_class: Type[PangeaResponseResult], check_response: bool = True
-    ) -> PangeaResponse[Type[PangeaResponseResult]]:
+    async 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:
@@ -110,7 +123,7 @@ class PangeaRequestAsync(PangeaRequestBase):
 
         async with self.session.get(url, headers=self._headers()) as requests_response:
             await self._check_http_errors(requests_response)
-            pangea_response = PangeaResponse(  # type: ignore[var-annotated]
+            pangea_response = PangeaResponse(
                 requests_response, result_class=result_class, json=await requests_response.json()
             )
 
@@ -131,11 +144,11 @@ class PangeaRequestAsync(PangeaRequestBase):
             raise pe.ServiceTemporarilyUnavailable(await resp.json())
 
     async def poll_result_by_id(
-        self, request_id: str, result_class: Union[Type[PangeaResponseResult], Type[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 await self.get(path, result_class, check_response=check_response)  # type: ignore[arg-type]
+        return await self.get(path, result_class, check_response=check_response)
 
     async def poll_result_once(self, response: PangeaResponse, check_response: bool = True):
         request_id = response.request_id
@@ -160,7 +173,7 @@ class PangeaRequestAsync(PangeaRequestBase):
         if resp.status < 200 or resp.status >= 300:
             raise pe.PresignedUploadError(f"presigned POST failure: {resp.status}", await resp.text())
 
-    async def put_presigned_url(self, url: str, files: List[Tuple]):
+    async def put_presigned_url(self, url: str, files: Sequence[Tuple]):
         # Send put request with file as body
         resp = await self._http_put(url=url, files=files)
         self.logger.debug(
@@ -173,7 +186,20 @@ class PangeaRequestAsync(PangeaRequestBase):
         if resp.status < 200 or resp.status >= 300:
             raise pe.PresignedUploadError(f"presigned PUT failure: {resp.status}", await resp.text())
 
-    async def download_file(self, url: str, filename: Optional[str] = None) -> AttachedFile:
+    async def download_file(self, url: str, filename: str | None = None) -> AttachedFile:
+        """
+        Download file
+
+        Download a file from the specified URL and save it with the given
+        filename.
+
+        Args:
+            url: URL of the file to download
+            filename: Name to save the downloaded file as. If not provided, the
+              filename will be determined from the Content-Disposition header or
+              the URL.
+        """
+
         self.logger.debug(
             json.dumps(
                 {
@@ -209,16 +235,16 @@ class PangeaRequestAsync(PangeaRequestBase):
                 )
 
                 return AttachedFile(filename=filename, file=await response.read(), content_type=content_type)
-            else:
-                raise pe.DownloadFileError(f"Failed to download file. Status: {response.status}", await response.text())
+            raise pe.DownloadFileError(f"Failed to download file. Status: {response.status}", await response.text())
 
-    async def _get_pangea_json(self, reader: aiohttp.MultipartReader) -> Optional[Dict]:
+    async def _get_pangea_json(self, reader: aiohttp.multipart.MultipartResponseWrapper) -> Optional[Dict[str, Any]]:
         # Iterate through parts
         async for part in reader:
-            return await part.json()
+            if isinstance(part, aiohttp.BodyPartReader):
+                return await part.json()
         return None
 
-    async def _get_attached_files(self, reader: aiohttp.MultipartReader) -> List[AttachedFile]:
+    async def _get_attached_files(self, reader: aiohttp.multipart.MultipartResponseWrapper) -> List[AttachedFile]:
         files = []
         i = 0
 
@@ -229,7 +255,7 @@ class PangeaRequestAsync(PangeaRequestBase):
             if name is None:
                 name = f"default_file_name_{i}"
                 i += 1
-            files.append(AttachedFile(name, await part.read(), content_type))
+            files.append(AttachedFile(name, await part.read(), content_type))  # type: ignore[union-attr]
 
         return files
 
@@ -237,13 +263,12 @@ class PangeaRequestAsync(PangeaRequestBase):
         # Parse the multipart response
         multipart_reader = aiohttp.MultipartReader.from_response(resp)
 
-        pangea_json = await self._get_pangea_json(multipart_reader)  # type: ignore[arg-type]
+        pangea_json = await self._get_pangea_json(multipart_reader)
         self.logger.debug(
             json.dumps({"service": self.service, "action": "multipart response", "response": pangea_json})
         )
 
-        multipart_reader = multipart_reader.__aiter__()
-        attached_files = await self._get_attached_files(multipart_reader)  # type: ignore[arg-type]
+        attached_files = await self._get_attached_files(multipart_reader)
         return MultipartResponse(pangea_json, attached_files)  # type: ignore[arg-type]
 
     async def _http_post(
@@ -251,7 +276,7 @@ class PangeaRequestAsync(PangeaRequestBase):
         url: str,
         headers: Dict = {},
         data: Union[str, Dict] = {},
-        files: List[Tuple] = [],
+        files: Optional[List[Tuple]] = [],
         presigned_url_post: bool = False,
     ) -> aiohttp.ClientResponse:
         if files:
@@ -276,7 +301,7 @@ class PangeaRequestAsync(PangeaRequestBase):
     async def _http_put(
         self,
         url: str,
-        files: List[Tuple],
+        files: Sequence[Tuple],
         headers: Dict = {},
     ) -> aiohttp.ClientResponse:
         self.logger.debug(
@@ -328,9 +353,7 @@ class PangeaRequestAsync(PangeaRequestBase):
         # Receive 202
         return await self._poll_presigned_url(accepted_exception.response)
 
-    async def _poll_presigned_url(
-        self, response: PangeaResponse[Type[PangeaResponseResult]]
-    ) -> PangeaResponse[Type[PangeaResponseResult]]:
+    async def _poll_presigned_url(self, response: PangeaResponse[TResult]) -> PangeaResponse[TResult]:
         if response.http_status != 202:
             raise AttributeError("Response should be 202")
 
@@ -373,8 +396,7 @@ class PangeaRequestAsync(PangeaRequestBase):
 
         if loop_resp.accepted_result is not None and not loop_resp.accepted_result.has_upload_url:
             return loop_resp
-        else:
-            raise loop_exc
+        raise loop_exc
 
     async def _handle_queued_result(self, response: PangeaResponse) -> PangeaResponse:
         if self._queued_retry_enabled and response.http_status == 202:

pangea/asyncio/services/__init__.py

@@ -1,8 +1,12 @@
+from .ai_guard import AIGuardAsync
 from .audit import AuditAsync
 from .authn import AuthNAsync
+from .authz import AuthZAsync
 from .embargo import EmbargoAsync
 from .file_scan import FileScanAsync
 from .intel import DomainIntelAsync, FileIntelAsync, IpIntelAsync, UrlIntelAsync, UserIntelAsync
+from .prompt_guard import PromptGuardAsync
 from .redact import RedactAsync
+from .sanitize import SanitizeAsync
 from .share import ShareAsync
 from .vault import VaultAsync

pangea/asyncio/services/ai_guard.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from pangea.asyncio.services.base import ServiceBaseAsync
+from pangea.config import PangeaConfig
+from pangea.response import PangeaResponse
+from pangea.services.ai_guard import TextGuardResult
+
+
+class AIGuardAsync(ServiceBaseAsync):
+    """AI Guard service client.
+
+    Provides methods to interact with Pangea's AI Guard service.
+
+    Examples:
+        from pangea import PangeaConfig
+        from pangea.asyncio.services import AIGuardAsync
+
+        config = PangeaConfig(domain="aws.us.pangea.cloud")
+        ai_guard = AIGuardAsync(token="pangea_token", config=config)
+    """
+
+    service_name = "ai-guard"
+
+    def __init__(
+        self, token: str, config: PangeaConfig | None = None, logger_name: str = "pangea", config_id: str | None = None
+    ) -> None:
+        """
+        AI Guard service client.
+
+        Initializes a new AI Guard client.
+
+        Args:
+            token: Pangea API token.
+            config: Pangea service configuration.
+            logger_name: Logger name.
+            config_id: Configuration ID.
+
+        Examples:
+            from pangea import PangeaConfig
+            from pangea.asyncio.services import AIGuardAsync
+
+            config = PangeaConfig(domain="aws.us.pangea.cloud")
+            ai_guard = AIGuardAsync(token="pangea_token", config=config)
+        """
+
+        super().__init__(token, config, logger_name, config_id)
+
+    async def guard_text(
+        self,
+        text: str,
+        *,
+        recipe: str = "pangea_prompt_guard",
+        debug: bool = False,
+    ) -> PangeaResponse[TextGuardResult]:
+        """
+        Text guard (Beta)
+
+        Guard text.
+
+        How to install a [Beta release](https://pangea.cloud/docs/sdk/python/#beta-releases).
+
+        OperationId: ai_guard_post_v1beta_text_guard
+
+        Args:
+            text: Text.
+            recipe: Recipe.
+            debug: Debug.
+
+        Examples:
+            response = await ai_guard.guard_text("text")
+        """
+
+        return await self.request.post(
+            "v1beta/text/guard", TextGuardResult, data={"text": text, "recipe": recipe, "debug": debug}
+        )