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

pangea/crypto/rsa.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+import base64
+from typing import TYPE_CHECKING
+
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import padding, rsa
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+
+from pangea.services.vault.models.common import ExportEncryptionAlgorithm
+from pangea.services.vault.models.symmetric import SymmetricKeyEncryptionAlgorithm
+
+if TYPE_CHECKING:
+    from pangea.services.vault.models.common import ExportResult
+
+
+def generate_key_pair() -> tuple[rsa.RSAPrivateKey, rsa.RSAPublicKey]:
+    # Generate a 4096-bit RSA key pair
+    private_key = rsa.generate_private_key(
+        public_exponent=65537,
+        key_size=4096,
+    )
+
+    # Extract the public key from the private key
+    public_key = private_key.public_key()
+    return private_key, public_key
+
+
+def decrypt_sha512(private_key: rsa.RSAPrivateKey, encrypted_message: bytes) -> bytes:
+    # Decrypt the message using the private key and OAEP padding
+    return private_key.decrypt(
+        encrypted_message,
+        padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA512()), algorithm=hashes.SHA512(), label=None),
+    )
+
+
+def encrypt_sha512(public_key: rsa.RSAPublicKey, message: bytes) -> bytes:
+    # Encrypt the message using the public key and OAEP padding
+    return public_key.encrypt(
+        message, padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA512()), algorithm=hashes.SHA512(), label=None)
+    )
+
+
+def private_key_to_pem(private_key: rsa.RSAPrivateKey) -> bytes:
+    # Serialize private key to PEM format
+    return private_key.private_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PrivateFormat.TraditionalOpenSSL,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+
+
+def public_key_to_pem(public_key: rsa.RSAPublicKey) -> bytes:
+    # Serialize public key to PEM format
+    return public_key.public_bytes(
+        encoding=serialization.Encoding.PEM, format=serialization.PublicFormat.SubjectPublicKeyInfo
+    )
+
+
+_AES_GCM_IV_SIZE = 12
+"""Standard nonce size for GCM."""
+
+_KEY_LENGTH = 32
+"""AES-256 key length in bytes."""
+
+
+def kem_decrypt(
+    private_key: rsa.RSAPrivateKey,
+    iv: bytes,
+    ciphertext: bytes,
+    symmetric_algorithm: str,
+    asymmetric_algorithm: str,
+    encrypted_salt: bytes,
+    password: str,
+    iteration_count: int,
+    hash_algorithm: str,
+) -> str:
+    if symmetric_algorithm.casefold() != SymmetricKeyEncryptionAlgorithm.AES_GCM_256.value.casefold():
+        raise NotImplementedError(f"Unsupported symmetric algorithm: {symmetric_algorithm}")
+
+    if asymmetric_algorithm != ExportEncryptionAlgorithm.RSA_NO_PADDING_4096_KEM:
+        raise NotImplementedError(f"Unsupported asymmetric algorithm: {asymmetric_algorithm}")
+
+    if hash_algorithm.casefold() != "SHA512".casefold():
+        raise NotImplementedError(f"Unsupported hash algorithm: {hash_algorithm}")
+
+    # No-padding RSA decryption.
+    n = private_key.private_numbers().public_numbers.n
+    salt = pow(
+        int.from_bytes(encrypted_salt, byteorder="big"),
+        private_key.private_numbers().d,
+        n,
+    ).to_bytes(n.bit_length() // 8, byteorder="big")
+
+    kdf = PBKDF2HMAC(
+        algorithm=hashes.SHA512(), length=_KEY_LENGTH, salt=salt, iterations=iteration_count, backend=default_backend()
+    )
+    symmetric_key = kdf.derive(password.encode("utf-8"))
+
+    decrypted = AESGCM(symmetric_key).decrypt(nonce=iv, data=ciphertext, associated_data=None)
+
+    return decrypted.decode("ascii")
+
+
+def kem_decrypt_export_result(*, result: ExportResult, password: str, private_key: rsa.RSAPrivateKey) -> str:
+    """Decrypt the exported result of a KEM operation."""
+    cipher_encoded = result.private_key or result.key
+    if not cipher_encoded:
+        raise TypeError("`private_key` or `key` should be set.")
+
+    assert result.encrypted_salt
+    assert result.symmetric_algorithm
+    assert result.asymmetric_algorithm
+    assert result.iteration_count
+    assert result.hash_algorithm
+
+    cipher_with_iv = base64.b64decode(cipher_encoded)
+    encrypted_salt = base64.b64decode(result.encrypted_salt)
+
+    iv = cipher_with_iv[:_AES_GCM_IV_SIZE]
+    cipher = cipher_with_iv[_AES_GCM_IV_SIZE:]
+
+    return kem_decrypt(
+        private_key=private_key,
+        iv=iv,
+        ciphertext=cipher,
+        password=password,
+        encrypted_salt=encrypted_salt,
+        symmetric_algorithm=result.symmetric_algorithm,
+        asymmetric_algorithm=result.asymmetric_algorithm,
+        iteration_count=result.iteration_count,
+        hash_algorithm=result.hash_algorithm,
+    )

pangea/deep_verify.py

@@ -263,8 +263,14 @@ def main():
         audit = init_audit(args.token, args.domain)
         errors = deep_verify(audit, args.file)
 
-        print("\n\nTotal errors:")
+        print("\n\nWarnings:")
+        val = errors["not_persisted"]
+        print(f"\tnot_persisted: {val}")
+
+        print("\nTotal errors:")
         for key, val in errors.items():
+            if key == "not_persisted":
+                continue
             print(f"\t{key.title()}: {val}")
         print()
 

pangea/dump_audit.py

@@ -19,7 +19,7 @@ from pangea.utils import default_encoder
 
 
 def dump_event(output: io.TextIOWrapper, row: SearchEvent, resp: PangeaResponse[SearchOutput]):
-    row_data = filter_deep_none(row.dict())
+    row_data = filter_deep_none(row.model_dump())
     if resp.result and resp.result.root:
         row_data["tree_size"] = resp.result.root.size
     output.write(json.dumps(row_data, default=default_encoder) + "\n")
@@ -63,11 +63,12 @@ def dump_before(audit: Audit, output: io.TextIOWrapper, start: datetime) -> int:
     cnt = 0
     if search_res.result and search_res.result.count > 0:
         leaf_index = search_res.result.events[0].leaf_index
-        for row in reversed(search_res.result.events):
-            if row.leaf_index != leaf_index:
-                break
-            dump_event(output, row, search_res)
-            cnt += 1
+        if leaf_index is not None:
+            for row in reversed(search_res.result.events):
+                if row.leaf_index != leaf_index:
+                    break
+                dump_event(output, row, search_res)
+                cnt += 1
     print(f"Dumping before... {cnt} events")
     return cnt
 
@@ -89,7 +90,7 @@ def dump_after(audit: Audit, output: io.TextIOWrapper, start: datetime, last_eve
     cnt = 0
     if search_res.result and search_res.result.count > 0:
         leaf_index = search_res.result.events[0].leaf_index
-        if leaf_index == last_leaf_index:
+        if leaf_index is not None and leaf_index == last_leaf_index:
             start_idx: int = 1 if last_event_hash == search_res.result.events[0].hash else 0
             for row in search_res.result.events[start_idx:]:
                 if row.leaf_index != leaf_index:
@@ -124,7 +125,7 @@ def dump_page(
     msg = f"Dumping... {search_res.result.count} events"
 
     if search_res.result.count <= 1:
-        return end, 0  # type: ignore[return-value]
+        return end, 0, True, "", 0
 
     offset = 0
     result_id = search_res.result.id

pangea/request.py

@@ -1,16 +1,19 @@
 # Copyright 2022 Pangea Cyber Corporation
 # Author: Pangea Cyber Corporation
+from __future__ import annotations
 
 import copy
 import json
 import logging
 import time
-from typing import Dict, List, Optional, Tuple, Type, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, Type, Union, cast
 
-import aiohttp
 import requests
+from pydantic import BaseModel
+from pydantic_core import to_jsonable_python
 from requests.adapters import HTTPAdapter, Retry
-from requests_toolbelt import MultipartDecoder  # type: ignore
+from requests_toolbelt import MultipartDecoder  # type: ignore[import-untyped]
+from typing_extensions import TypeVar
 
 import pangea
 import pangea.exceptions as pe
@@ -18,8 +21,11 @@ from pangea.config import PangeaConfig
 from pangea.response import AttachedFile, PangeaResponse, PangeaResponseResult, ResponseStatus, TransferMethod
 from pangea.utils import default_encoder
 
+if TYPE_CHECKING:
+    import aiohttp
 
-class MultipartResponse(object):
+
+class MultipartResponse:
     pangea_json: Dict[str, str]
     attached_files: List = []
 
@@ -28,7 +34,7 @@ class MultipartResponse(object):
         self.attached_files = attached_files
 
 
-class PangeaRequestBase(object):
+class PangeaRequestBase:
     def __init__(
         self, config: PangeaConfig, token: str, service: str, logger: logging.Logger, config_id: Optional[str] = None
     ):
@@ -126,8 +132,7 @@ class PangeaRequestBase(object):
         filename_parts = content_disposition.split("name=")
         if len(filename_parts) > 1:
             return filename_parts[1].split(";")[0].strip('"')
-        else:
-            return None
+        return None
 
     def _get_filename_from_url(self, url: str) -> Optional[str]:
         return url.split("/")[-1].split("?")[0]
@@ -154,39 +159,42 @@ class PangeaRequestBase(object):
 
         if status == ResponseStatus.VALIDATION_ERR.value:
             raise pe.ValidationException(summary, response)
-        elif status == ResponseStatus.TOO_MANY_REQUESTS.value:
+        if status == ResponseStatus.TOO_MANY_REQUESTS.value:
             raise pe.RateLimitException(summary, response)
-        elif status == ResponseStatus.NO_CREDIT.value:
+        if status == ResponseStatus.NO_CREDIT.value:
             raise pe.NoCreditException(summary, response)
-        elif status == ResponseStatus.UNAUTHORIZED.value:
+        if status == ResponseStatus.UNAUTHORIZED.value:
             raise pe.UnauthorizedException(self.service, response)
-        elif status == ResponseStatus.SERVICE_NOT_ENABLED.value:
+        if status == ResponseStatus.SERVICE_NOT_ENABLED.value:
             raise pe.ServiceNotEnabledException(self.service, response)
-        elif status == ResponseStatus.PROVIDER_ERR.value:
+        if status == ResponseStatus.PROVIDER_ERR.value:
             raise pe.ProviderErrorException(summary, response)
-        elif status in (ResponseStatus.MISSING_CONFIG_ID_SCOPE.value, ResponseStatus.MISSING_CONFIG_ID.value):
+        if status in (ResponseStatus.MISSING_CONFIG_ID_SCOPE.value, ResponseStatus.MISSING_CONFIG_ID.value):
             raise pe.MissingConfigID(self.service, response)
-        elif status == ResponseStatus.SERVICE_NOT_AVAILABLE.value:
+        if status == ResponseStatus.SERVICE_NOT_AVAILABLE.value:
             raise pe.ServiceNotAvailableException(summary, response)
-        elif status == ResponseStatus.TREE_NOT_FOUND.value:
+        if status == ResponseStatus.TREE_NOT_FOUND.value:
             raise pe.TreeNotFoundException(summary, response)
-        elif status == ResponseStatus.IP_NOT_FOUND.value:
+        if status == ResponseStatus.IP_NOT_FOUND.value:
             raise pe.IPNotFoundException(summary, response)
-        elif status == ResponseStatus.BAD_OFFSET.value:
+        if status == ResponseStatus.BAD_OFFSET.value:
             raise pe.BadOffsetException(summary, response)
-        elif status == ResponseStatus.FORBIDDEN_VAULT_OPERATION.value:
+        if status == ResponseStatus.FORBIDDEN_VAULT_OPERATION.value:
             raise pe.ForbiddenVaultOperation(summary, response)
-        elif status == ResponseStatus.VAULT_ITEM_NOT_FOUND.value:
+        if status == ResponseStatus.VAULT_ITEM_NOT_FOUND.value:
             raise pe.VaultItemNotFound(summary, response)
-        elif status == ResponseStatus.NOT_FOUND.value:
-            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:
+        if status == ResponseStatus.NOT_FOUND.value:
+            raise pe.NotFound(str(response.raw_response.url) if response.raw_response is not None else "", response)
+        if status == ResponseStatus.INTERNAL_SERVER_ERROR.value:
             raise pe.InternalServerError(response)
-        elif status == ResponseStatus.ACCEPTED.value:
+        if status == ResponseStatus.ACCEPTED.value:
             raise pe.AcceptedRequestException(response)
         raise pe.PangeaAPIException(f"{summary} ", response)
 
 
+TResult = TypeVar("TResult", bound=PangeaResponseResult)
+
+
 class PangeaRequest(PangeaRequestBase):
     """An object that makes direct calls to Pangea Service APIs.
 
@@ -202,12 +210,12 @@ class PangeaRequest(PangeaRequestBase):
     def post(
         self,
         endpoint: str,
-        result_class: Type[PangeaResponseResult],
-        data: Union[str, Dict] = {},
+        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:
@@ -218,6 +226,16 @@ class PangeaRequest(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)
 
@@ -318,32 +336,33 @@ class PangeaRequest(PangeaRequestBase):
         return self.session.post(url, headers=headers, data=data_send, files=files)
 
     def _http_post_process(
-        self, data: Union[str, Dict] = {}, files: Optional[List[Tuple]] = None, multipart_post: bool = True
+        self,
+        data: Union[str, Dict] = {},
+        files: Optional[Sequence[Tuple[str, Tuple[Any, str, str]]]] = None,
+        multipart_post: bool = True,
     ):
         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)  # type: ignore[arg-type]
-                files = multi  # type: ignore[assignment]
+                multi.extend(files)
+                files = multi
                 return None, files
-            else:
-                # 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
-        else:
-            data_send = json.dumps(data, default=default_encoder) if isinstance(data, dict) else data
-            return data_send, None
+            # 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
+        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) -> PangeaResponse[Type[PangeaResponseResult]]:
+    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(
@@ -355,7 +374,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:
@@ -387,7 +406,20 @@ class PangeaRequest(PangeaRequestBase):
 
         return self._check_response(pangea_response)
 
-    def download_file(self, url: str, filename: Optional[str] = None) -> AttachedFile:
+    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(
                 {
@@ -423,25 +455,24 @@ class PangeaRequest(PangeaRequestBase):
                 )
             )
             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)
+        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], 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 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
-    ) -> PangeaResponse[Type[PangeaResponseResult]]:
+        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)
 
@@ -526,7 +557,7 @@ class PangeaRequest(PangeaRequestBase):
         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[Type[PangeaResponseResult]]:
+    def _poll_result_retry(self, response: PangeaResponse[TResult]) -> PangeaResponse[TResult]:
         retry_count = 1
         start = time.time()
 
@@ -538,9 +569,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[Type[PangeaResponseResult]]
-    ) -> PangeaResponse[Type[PangeaResponseResult]]:
+    def _poll_presigned_url(self, response: PangeaResponse[TResult]) -> PangeaResponse[TResult]:
         if response.http_status != 202:
             raise AttributeError("Response should be 202")
 
@@ -583,8 +612,7 @@ class PangeaRequest(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
 
     def _init_session(self) -> requests.Session:
         retry_config = Retry(

pangea/response.py

@@ -3,16 +3,15 @@
 import datetime
 import enum
 import os
-from typing import Any, Dict, Generic, List, Optional, Type, TypeVar, Union
+from typing import Any, Dict, Generic, List, Optional, Type, Union
 
 import aiohttp
 import requests
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict, PlainSerializer
+from typing_extensions import Annotated, TypeVar
 
 from pangea.utils import format_datetime
 
-T = TypeVar("T")
-
 
 class AttachedFile(object):
     filename: str
@@ -50,10 +49,24 @@ class AttachedFile(object):
 
 
 class TransferMethod(str, enum.Enum):
+    """Transfer methods for uploading file data."""
+
     MULTIPART = "multipart"
     POST_URL = "post-url"
     PUT_URL = "put-url"
     SOURCE_URL = "source-url"
+    """
+    A `source-url` is a caller-specified URL where the Pangea APIs can fetch the
+    contents of the input file. When calling a Pangea API with a
+    `transfer_method` of `source-url`, you must also specify a `source_url`
+    input parameter that provides a URL to the input file. The source URL can be
+    a presigned URL created by the caller, and it will be used to download the
+    content of the input file. The `source-url` transfer method is useful when
+    you already have a file in your storage and can provide a URL from which
+    Pangea API can fetch the input file—there is no need to transfer it to
+    Pangea with a separate POST or PUT request.
+    """
+
     DEST_URL = "dest-url"
 
     def __str__(self):
@@ -63,24 +76,17 @@ class TransferMethod(str, enum.Enum):
         return str(self.value)
 
 
+PangeaDateTime = Annotated[datetime.datetime, PlainSerializer(format_datetime)]
+
+
 # API response should accept arbitrary fields to make them accept possible new parameters
 class APIResponseModel(BaseModel):
-    class Config:
-        arbitrary_types_allowed = True
-        # allow parameters despite they are not declared in model. Make SDK accept server new parameters
-        extra = "allow"
+    model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
 
 
 # API request models doesn't not allow arbitrary fields
 class APIRequestModel(BaseModel):
-    class Config:
-        arbitrary_types_allowed = True
-        extra = (
-            "allow"  # allow parameters despite they are not declared in model. Make SDK accept server new parameters
-        )
-        json_encoders = {
-            datetime.datetime: format_datetime,
-        }
+    model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
 
 
 class PangeaResponseResult(APIResponseModel):
@@ -155,38 +161,50 @@ class ResponseStatus(str, enum.Enum):
 
 
 class ResponseHeader(APIResponseModel):
-    """
-    Pangea response API header.
-
-    Arguments:
-    request_id -- The request ID.
-    request_time -- The time the request was issued, ISO8601.
-    response_time -- The time the response was issued, ISO8601.
-    status -- Pangea response status
-    summary -- The summary of the response.
-    """
+    """Pangea response API header."""
 
     request_id: str
+    """A unique identifier assigned to each request made to the API."""
+
     request_time: str
+    """
+    Timestamp indicating the exact moment when a request is made to the API.
+    """
+
     response_time: str
+    """
+    Duration it takes for the API to process a request and generate a response.
+    """
+
     status: str
+    """
+    Represents the status or outcome of the API request.
+    """
+
     summary: str
+    """
+    Provides a concise and brief overview of the purpose or primary objective of
+    the API endpoint.
+    """
+
+
+T = TypeVar("T", bound=PangeaResponseResult)
 
 
-class PangeaResponse(Generic[T], ResponseHeader):
+class PangeaResponse(ResponseHeader, Generic[T]):
     raw_result: Optional[Dict[str, Any]] = None
     raw_response: Optional[Union[requests.Response, aiohttp.ClientResponse]] = None
     result: Optional[T] = None
     pangea_error: Optional[PangeaError] = None
     accepted_result: Optional[AcceptedResult] = None
-    result_class: Union[Type[PangeaResponseResult], Type[dict]] = PangeaResponseResult
+    result_class: Type[T] = PangeaResponseResult  # type: ignore[assignment]
     _json: Any
     attached_files: List[AttachedFile] = []
 
     def __init__(
         self,
         response: requests.Response,
-        result_class: Union[Type[PangeaResponseResult], Type[dict]],
+        result_class: Type[T],
         json: dict,
         attached_files: List[AttachedFile] = [],
     ):
@@ -198,7 +216,7 @@ class PangeaResponse(Generic[T], ResponseHeader):
         self.attached_files = attached_files
 
         self.result = (
-            self.result_class(**self.raw_result)  # type: ignore[assignment]
+            self.result_class(**self.raw_result)
             if self.raw_result is not None and issubclass(self.result_class, PangeaResponseResult) and self.success
             else None
         )
@@ -230,4 +248,4 @@ class PangeaResponse(Generic[T], ResponseHeader):
 
     @property
     def url(self) -> str:
-        return str(self.raw_response.url)  # type: ignore[arg-type,union-attr]
+        return str(self.raw_response.url)  # type: ignore[union-attr]

pangea/services/__init__.py

@@ -1,8 +1,12 @@
+from .ai_guard import AIGuard
 from .audit.audit import Audit
 from .authn.authn import AuthN
+from .authz import AuthZ
 from .embargo import Embargo
 from .file_scan import FileScan
 from .intel import DomainIntel, FileIntel, IpIntel, UrlIntel, UserIntel
+from .prompt_guard import PromptGuard
 from .redact import Redact
+from .sanitize import Sanitize
 from .share.share import Share
 from .vault.vault import Vault