pangea-sdk 3.8.0__py3-none-any.whl → 5.3.0__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/file_uploader.py

@@ -0,0 +1,35 @@
+# Copyright 2022 Pangea Cyber Corporation
+# Author: Pangea Cyber Corporation
+import io
+import logging
+from typing import Dict, Optional
+
+from pangea.request import PangeaConfig, PangeaRequest
+from pangea.response import TransferMethod
+
+
+class FileUploader:
+    def __init__(self):
+        self.logger = logging.getLogger("pangea")
+        self._request = PangeaRequest(
+            config=PangeaConfig(),
+            token="",
+            service="FileUploader",
+            logger=self.logger,
+        )
+
+    def upload_file(
+        self,
+        url: str,
+        file: io.BufferedReader,
+        transfer_method: TransferMethod = TransferMethod.PUT_URL,
+        file_details: Optional[Dict] = None,
+    ):
+        if transfer_method == TransferMethod.PUT_URL:
+            files = [("file", ("filename", file, "application/octet-stream"))]
+            self._request.put_presigned_url(url=url, files=files)
+        elif transfer_method == TransferMethod.POST_URL:
+            files = [("file", ("filename", file, "application/octet-stream"))]
+            self._request.post_presigned_url(url=url, data=file_details, files=files)
+        else:
+            raise ValueError(f"Transfer method not supported: {transfer_method}")

pangea/request.py

@@ -1,15 +1,17 @@
 # Copyright 2022 Pangea Cyber Corporation
 # Author: Pangea Cyber Corporation
+from __future__ import annotations
 
 import copy
 import json
 import logging
 import time
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, Type, Union
 
 import requests
+from pydantic import BaseModel
 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
@@ -22,7 +24,7 @@ if TYPE_CHECKING:
     import aiohttp
 
 
-class MultipartResponse(object):
+class MultipartResponse:
     pangea_json: Dict[str, str]
     attached_files: List = []
 
@@ -31,7 +33,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
     ):
@@ -129,8 +131,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]
@@ -157,40 +158,40 @@ 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, default=PangeaResponseResult)
+TResult = TypeVar("TResult", bound=PangeaResponseResult)
 
 
 class PangeaRequest(PangeaRequestBase):
@@ -209,7 +210,7 @@ class PangeaRequest(PangeaRequestBase):
         self,
         endpoint: str,
         result_class: Type[TResult],
-        data: Union[str, Dict] = {},
+        data: str | BaseModel | dict[str, Any] | None = None,
         files: Optional[List[Tuple]] = None,
         poll_result: bool = True,
         url: Optional[str] = None,
@@ -224,6 +225,13 @@ 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 = {}
+
         if url is None:
             url = self._url(endpoint)
 
@@ -324,28 +332,29 @@ 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
 
@@ -393,7 +402,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(
                 {
@@ -429,8 +451,7 @@ 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: Type[TResult], check_response: bool = True
@@ -459,10 +480,8 @@ class PangeaRequest(PangeaRequestBase):
     ) -> PangeaResponse:
         # 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)
-
+            # This should return 202 (AcceptedRequestException) at least zero size file is sent
+            return self.post(endpoint=endpoint, result_class=result_class, data=data, poll_result=False)
         except pe.AcceptedRequestException as e:
             accepted_exception = e
         except Exception as e:
@@ -520,6 +539,9 @@ class PangeaRequest(PangeaRequestBase):
             raise AttributeError("files attribute should have at least 1 file")
 
         response = self.request_presigned_url(endpoint=endpoint, result_class=result_class, data=data)
+
+        if response.success:  # This should only happen when uploading a zero bytes file
+            return response.raw_response
         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:
@@ -586,8 +608,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

@@ -7,8 +7,8 @@ from typing import Any, Dict, Generic, List, Optional, Type, Union
 
 import aiohttp
 import requests
-from pydantic import BaseModel
-from typing_extensions import TypeVar
+from pydantic import BaseModel, ConfigDict, PlainSerializer
+from typing_extensions import Annotated, TypeVar
 
 from pangea.utils import format_datetime
 
@@ -28,6 +28,7 @@ class AttachedFile(object):
             filename = self.filename if self.filename else "default_save_filename"
 
         filepath = os.path.join(dest_folder, filename)
+        filepath = self._find_available_file(filepath)
         directory = os.path.dirname(filepath)
         if not os.path.exists(directory):
             os.makedirs(directory)
@@ -35,12 +36,37 @@ class AttachedFile(object):
         with open(filepath, "wb") as file:
             file.write(self.file)
 
+    def _find_available_file(self, file_path):
+        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}"
+            counter += 1
+        return file_path
+
 
 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):
@@ -50,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):
@@ -169,10 +188,10 @@ class ResponseHeader(APIResponseModel):
     """
 
 
-T = TypeVar("T", bound=PangeaResponseResult, default=PangeaResponseResult)
+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
@@ -229,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

@@ -5,4 +5,6 @@ from .embargo import Embargo
 from .file_scan import FileScan
 from .intel import DomainIntel, FileIntel, IpIntel, UrlIntel, UserIntel
 from .redact import Redact
+from .sanitize import Sanitize
+from .share.share import Share
 from .vault.vault import Vault