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

pangea/response.py

@@ -2,14 +2,38 @@
 # Author: Pangea Cyber Corporation
 import datetime
 import enum
-from typing import Any, Dict, Generic, List, Optional, Type, TypeVar, Union
+import os
+from typing import Any, Dict, Generic, List, Optional, Type, Union
 
 import aiohttp
 import requests
-from pangea.utils import format_datetime
 from pydantic import BaseModel
+from typing_extensions import TypeVar
+
+from pangea.utils import format_datetime
+
+
+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"
 
-T = TypeVar("T")
+        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):
@@ -17,6 +41,7 @@ class TransferMethod(str, enum.Enum):
     POST_URL = "post-url"
     PUT_URL = "put-url"
     SOURCE_URL = "source-url"
+    DEST_URL = "dest-url"
 
     def __str__(self):
         return str(self.value)
@@ -117,22 +142,34 @@ 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, default=PangeaResponseResult)
 
 
 class PangeaResponse(Generic[T], ResponseHeader):
@@ -141,17 +178,26 @@ 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: Type[T] = PangeaResponseResult  # type: ignore[assignment]
     _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: Type[T],
+        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]
+            self.result_class(**self.raw_result)
             if self.raw_result is not None and issubclass(self.result_class, PangeaResponseResult) and self.success
             else None
         )

pangea/services/__init__.py

@@ -1,5 +1,6 @@
 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

pangea/services/audit/audit.py

@@ -1,16 +1,23 @@
 # Copyright 2022 Pangea Cyber Corporation
 # Author: Pangea Cyber Corporation
+from __future__ import annotations
+
 import datetime
 import json
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, Union
 
 import pangea.exceptions as pexc
-from pangea.response import PangeaResponse
+from pangea.config import PangeaConfig
+from pangea.response import PangeaResponse, PangeaResponseResult
 from pangea.services.audit.exceptions import AuditException, EventCorruption
 from pangea.services.audit.models import (
+    DownloadFormat,
+    DownloadRequest,
+    DownloadResult,
     Event,
     EventEnvelope,
     EventVerification,
+    ExportRequest,
     LogBulkRequest,
     LogBulkResult,
     LogEvent,
@@ -48,9 +55,9 @@ from pangea.utils import canonicalize_nested_json
 
 class AuditBase:
     def __init__(
-        self, private_key_file: str = "", public_key_info: Dict[str, str] = {}, tenant_id: Optional[str] = None
-    ):
-        self.pub_roots: Dict[int, Root] = {}
+        self, private_key_file: str = "", public_key_info: dict[str, str] = {}, tenant_id: str | None = None
+    ) -> None:
+        self.pub_roots: Dict[int, PublishedRoot] = {}
         self.buffer_data: Optional[str] = None
         self.signer: Optional[Signer] = Signer(private_key_file) if private_key_file else None
         self.public_key_info = public_key_info
@@ -182,8 +189,6 @@ class AuditBase:
         unpublished_root = response.result.unpublished_root  # type: ignore[union-attr]
 
         if verify_consistency:
-            self.update_published_roots(response.result)  # type: ignore[arg-type]
-
             for search_event in response.result.events:  # type: ignore[union-attr]
                 # verify membership proofs
                 if self.can_verify_membership_proof(search_event):
@@ -201,22 +206,9 @@ class AuditBase:
 
         return response
 
-    def update_published_roots(self, result: SearchOutput):
-        """Fetches series of published root hashes from Arweave
-
-        This is used for subsequent calls to verify_consistency_proof(). Root hashes
-        are published on [Arweave](https://arweave.net).
-
-        Args:
-            result (SearchOutput): PangeaResponse object from previous call to audit.search()
-
-        Raises:
-            AuditException: If an audit based api exception happens
-            PangeaAPIException: If an API Error happens
-        """
-
+    def _get_tree_sizes_and_roots(self, result: SearchResultOutput) -> Tuple[Set[int], Dict[int, PublishedRoot]]:
         if not result.root:
-            return
+            return set(), {}
 
         tree_sizes = set()
         for search_event in result.events:
@@ -230,22 +222,11 @@ class AuditBase:
         tree_sizes.difference_update(self.pub_roots.keys())
 
         if tree_sizes:
-            arweave_roots = get_arweave_published_roots(result.root.tree_name, list(tree_sizes))  # + [result.count])
+            arweave_roots = get_arweave_published_roots(result.root.tree_name, list(tree_sizes))
         else:
             arweave_roots = {}
 
-        # fill the missing roots from the server (if allowed)
-        for tree_size in tree_sizes:
-            pub_root = None
-            if tree_size in arweave_roots:
-                pub_root = PublishedRoot(**arweave_roots[tree_size].dict(exclude_none=True))
-                pub_root.source = RootSource.ARWEAVE
-            elif self.allow_server_roots:
-                resp = self.root(tree_size=tree_size)  # type: ignore[attr-defined]
-                if resp.success:
-                    pub_root = PublishedRoot(**resp.result.data.dict(exclude_none=True))
-                    pub_root.source = RootSource.PANGEA
-            self.pub_roots[tree_size] = pub_root  # type: ignore[assignment]
+        return tree_sizes, arweave_roots
 
     def can_verify_membership_proof(self, event: SearchEvent) -> bool:
         """
@@ -303,7 +284,7 @@ class AuditBase:
         """
         return event.published and event.leaf_index is not None and event.leaf_index >= 0  # type: ignore[return-value]
 
-    def verify_consistency_proof(self, pub_roots: Dict[int, Root], event: SearchEvent) -> bool:
+    def verify_consistency_proof(self, pub_roots: Dict[int, PublishedRoot], event: SearchEvent) -> bool:
         """
         Verify consistency proof
 
@@ -329,7 +310,7 @@ class AuditBase:
             return False
 
         if not self.allow_server_roots and (
-            curr_root.source != RootSource.ARWEAVE or prev_root.source != RootSource.ARWEAVE  # type: ignore[attr-defined]
+            curr_root.source != RootSource.ARWEAVE or prev_root.source != RootSource.ARWEAVE
         ):
             return False
 
@@ -355,7 +336,9 @@ class AuditBase:
         if audit_envelope and audit_envelope.signature and public_key:
             v = Verifier()
             verification = v.verify_signature(
-                audit_envelope.signature, canonicalize_event(audit_envelope.event), public_key  # type: ignore[arg-type]
+                audit_envelope.signature,
+                canonicalize_event(Event(**audit_envelope.event)),
+                public_key,
             )
             if verification is not None:
                 return EventVerification.PASS if verification else EventVerification.FAIL
@@ -397,14 +380,32 @@ class Audit(ServiceBase, AuditBase):
 
     def __init__(
         self,
-        token,
-        config=None,
+        token: str,
+        config: PangeaConfig | None = None,
         private_key_file: str = "",
-        public_key_info: Dict[str, str] = {},
-        tenant_id: Optional[str] = None,
-        logger_name="pangea",
-        config_id: Optional[str] = None,
-    ):
+        public_key_info: dict[str, str] = {},
+        tenant_id: str | None = None,
+        logger_name: str = "pangea",
+        config_id: str | None = None,
+    ) -> None:
+        """
+        Audit client
+
+        Initializes a new Audit client.
+
+        Args:
+            token: Pangea API token.
+            config: Configuration.
+            private_key_file: Private key filepath.
+            public_key_info: Public key information.
+            tenant_id: Tenant ID.
+            logger_name: Logger name.
+            config_id: Configuration ID.
+
+        Examples:
+             config = PangeaConfig(domain="pangea_domain")
+             audit = Audit(token="pangea_token", config=config)
+        """
         # FIXME: Temporary check to deprecate config_id from PangeaConfig.
         # Delete it when deprecate PangeaConfig.config_id
         if config_id and config is not None and config.config_id is not None:
@@ -518,8 +519,8 @@ class Audit(ServiceBase, AuditBase):
         """
 
         input = self._get_log_request(event, sign_local=sign_local, verify=verify, verbose=verbose)
-        response = self.request.post("v1/log", LogResult, data=input.dict(exclude_none=True))
-        if response.success:
+        response: PangeaResponse[LogResult] = self.request.post("v1/log", LogResult, data=input.dict(exclude_none=True))
+        if response.success and response.result is not None:
             self._process_log_result(response.result, verify=verify)
         return response
 
@@ -557,9 +558,11 @@ class Audit(ServiceBase, AuditBase):
         """
 
         input = self._get_log_request(events, sign_local=sign_local, verify=False, verbose=verbose)
-        response = self.request.post("v2/log", LogBulkResult, data=input.dict(exclude_none=True))
+        response: PangeaResponse[LogBulkResult] = self.request.post(
+            "v2/log", LogBulkResult, data=input.dict(exclude_none=True)
+        )
 
-        if response.success:
+        if response.success and response.result is not None:
             for result in response.result.results:
                 self._process_log_result(result, verify=True)
         return response
@@ -599,13 +602,13 @@ class Audit(ServiceBase, AuditBase):
 
         try:
             # Calling to v2 methods will return always a 202.
-            response = self.request.post(
+            response: PangeaResponse[LogBulkResult] = self.request.post(
                 "v2/log_async", LogBulkResult, data=input.dict(exclude_none=True), poll_result=False
             )
         except pexc.AcceptedRequestException as e:
             return e.response
 
-        if response.success:
+        if response.success and response.result is not None:
             for result in response.result.results:
                 self._process_log_result(result, verify=True)
         return response
@@ -619,7 +622,7 @@ class Audit(ServiceBase, AuditBase):
         end: Optional[Union[datetime.datetime, str]] = None,
         limit: Optional[int] = None,
         max_results: Optional[int] = None,
-        search_restriction: Optional[dict] = None,
+        search_restriction: Optional[Dict[str, Sequence[str]]] = None,
         verbose: Optional[bool] = None,
         verify_consistency: bool = False,
         verify_events: bool = True,
@@ -648,7 +651,7 @@ class Audit(ServiceBase, AuditBase):
             end (datetime, optional): An RFC-3339 formatted timestamp, or relative time adjustment from the current time.
             limit (int, optional): Optional[int] = None,
             max_results (int, optional): Maximum number of results to return.
-            search_restriction (dict, optional): A list of keys to restrict the search results to. Useful for partitioning data available to the query string.
+            search_restriction (Dict[str, Sequence[str]], optional): A list of keys to restrict the search results to. Useful for partitioning data available to the query string.
             verbose (bool, optional): If true, response include root and membership and consistency proofs.
             verify_consistency (bool): True to verify logs consistency
             verify_events (bool): True to verify hash events and signatures
@@ -687,7 +690,11 @@ class Audit(ServiceBase, AuditBase):
             verbose=verbose,
         )
 
-        response = self.request.post("v1/search", SearchOutput, data=input.dict(exclude_none=True))
+        response: PangeaResponse[SearchOutput] = self.request.post(
+            "v1/search", SearchOutput, data=input.dict(exclude_none=True)
+        )
+        if verify_consistency and response.result is not None:
+            self.update_published_roots(response.result)
         return self.handle_search_response(response, verify_consistency, verify_events)
 
     def results(
@@ -695,6 +702,7 @@ class Audit(ServiceBase, AuditBase):
         id: str,
         limit: Optional[int] = 20,
         offset: Optional[int] = 0,
+        assert_search_restriction: Optional[Dict[str, Sequence[str]]] = None,
         verify_consistency: bool = False,
         verify_events: bool = True,
     ) -> PangeaResponse[SearchResultOutput]:
@@ -709,6 +717,7 @@ class Audit(ServiceBase, AuditBase):
             id (string): the id of a search action, found in `response.result.id`
             limit (integer, optional): the maximum number of results to return, default is 20
             offset (integer, optional): the position of the first result to return, default is 0
+            assert_search_restriction (Dict[str, Sequence[str]], optional): Assert the requested search results were queried with the exact same search restrictions, to ensure the results comply to the expected restrictions.
             verify_consistency (bool): True to verify logs consistency
             verify_events (bool): True to verify hash events and signatures
         Raises:
@@ -733,10 +742,116 @@ class Audit(ServiceBase, AuditBase):
             id=id,
             limit=limit,
             offset=offset,
+            assert_search_restriction=assert_search_restriction,
         )
-        response = self.request.post("v1/results", SearchResultOutput, data=input.dict(exclude_none=True))
+        response: PangeaResponse[SearchResultOutput] = self.request.post(
+            "v1/results", SearchResultOutput, data=input.dict(exclude_none=True)
+        )
+        if verify_consistency and response.result is not None:
+            self.update_published_roots(response.result)
         return self.handle_results_response(response, verify_consistency, verify_events)
 
+    def export(
+        self,
+        *,
+        format: DownloadFormat = DownloadFormat.CSV,
+        start: Optional[datetime.datetime] = None,
+        end: Optional[datetime.datetime] = None,
+        order: Optional[SearchOrder] = None,
+        order_by: Optional[str] = None,
+        verbose: bool = True,
+    ) -> PangeaResponse[PangeaResponseResult]:
+        """
+        Export from the audit log
+
+        Bulk export of data from the Secure Audit Log, with optional filtering.
+
+        OperationId: audit_post_v1_export
+
+        Args:
+            format: Format for the records.
+            start: The start of the time range to perform the search on.
+            end: The end of the time range to perform the search on. If omitted,
+                then all records up to the latest will be searched.
+            order: Specify the sort order of the response.
+            order_by: Name of column to sort the results by.
+            verbose: Whether or not to include the root hash of the tree and the
+                membership proof for each record.
+
+        Raises:
+            AuditException: If an audit based api exception happens
+            PangeaAPIException: If an API Error happens
+
+        Examples:
+            export_res = audit.export(verbose=False)
+
+            # Export may take several dozens of minutes, so polling for the result
+            # should be done in a loop. That is omitted here for brevity's sake.
+            try:
+                audit.poll_result(request_id=export_res.request_id)
+            except AcceptedRequestException:
+                # Retry later.
+
+            # Download the result when it's ready.
+            download_res = audit.download_results(request_id=export_res.request_id)
+            download_res.result.dest_url
+            # => https://pangea-runtime.s3.amazonaws.com/audit/xxxxx/search_results_[...]
+        """
+        input = ExportRequest(
+            format=format,
+            start=start,
+            end=end,
+            order=order,
+            order_by=order_by,
+            verbose=verbose,
+        )
+        try:
+            return self.request.post(
+                "v1/export", PangeaResponseResult, data=input.dict(exclude_none=True), poll_result=False
+            )
+        except pexc.AcceptedRequestException as e:
+            return e.response
+
+    def log_stream(self, data: dict) -> PangeaResponse[PangeaResponseResult]:
+        """
+        Log streaming endpoint
+
+        This API allows 3rd party vendors (like Auth0) to stream events to this
+        endpoint where the structure of the payload varies across different
+        vendors.
+
+        OperationId: audit_post_v1_log_stream
+
+        Args:
+            data: Event data. The exact schema of this will vary by vendor.
+
+        Raises:
+            AuditException: If an audit based api exception happens
+            PangeaAPIException: If an API Error happens
+
+        Examples:
+            data = {
+                "logs": [
+                    {
+                        "log_id": "some log ID",
+                        "data": {
+                            "date": "2024-03-29T17:26:50.193Z",
+                            "type": "sapi",
+                            "description": "Create a log stream",
+                            "client_id": "some client ID",
+                            "ip": "127.0.0.1",
+                            "user_agent": "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0",
+                            "user_id": "some user ID",
+                        },
+                    }
+                    # ...
+                ]
+            }
+
+            response = audit.log_stream(data)
+        """
+        return self.request.post("v1/log_stream", PangeaResponseResult, data=data)
+
     def root(self, tree_size: Optional[int] = None) -> PangeaResponse[RootResult]:
         """
         Tamperproof verification
@@ -760,3 +875,74 @@ class Audit(ServiceBase, AuditBase):
         """
         input = RootRequest(tree_size=tree_size)
         return self.request.post("v1/root", RootResult, data=input.dict(exclude_none=True))
+
+    def download_results(
+        self,
+        result_id: Optional[str] = None,
+        format: DownloadFormat = DownloadFormat.CSV,
+        request_id: Optional[str] = None,
+    ) -> PangeaResponse[DownloadResult]:
+        """
+        Download search results
+
+        Get all search results as a compressed (gzip) CSV file.
+
+        OperationId: audit_post_v1_download_results
+
+        Args:
+            result_id: ID returned by the search API.
+            format: Format for the records.
+            request_id: ID returned by the export API.
+
+        Returns:
+            URL where search results can be downloaded.
+
+        Raises:
+            AuditException: If an Audit-based API exception occurs.
+            PangeaAPIException: If an API exception occurs.
+
+        Examples:
+            response = audit.download_results(
+                result_id="pas_[...]",
+                format=DownloadFormat.JSON,
+            )
+        """
+
+        if request_id is None and result_id is None:
+            raise ValueError("must pass one of `request_id` or `result_id`")
+
+        input = DownloadRequest(request_id=request_id, result_id=result_id, format=format)
+        return self.request.post("v1/download_results", DownloadResult, data=input.dict(exclude_none=True))
+
+    def update_published_roots(self, result: SearchResultOutput):
+        """Fetches series of published root hashes from Arweave
+
+        This is used for subsequent calls to verify_consistency_proof(). Root hashes
+        are published on [Arweave](https://arweave.net).
+
+        Args:
+            result (SearchResultOutput): Result object from previous call to Audit.search() or Audit.results()
+
+        Raises:
+            AuditException: If an audit based api exception happens
+            PangeaAPIException: If an API Error happens
+        """
+
+        if not result.root:
+            return
+
+        tree_sizes, arweave_roots = self._get_tree_sizes_and_roots(result)
+
+        # fill the missing roots from the server (if allowed)
+        for tree_size in tree_sizes:
+            pub_root = None
+            if tree_size in arweave_roots:
+                pub_root = PublishedRoot(**arweave_roots[tree_size].dict(exclude_none=True))
+                pub_root.source = RootSource.ARWEAVE
+            elif self.allow_server_roots:
+                resp = self.root(tree_size=tree_size)
+                if resp.success and resp.result is not None:
+                    pub_root = PublishedRoot(**resp.result.data.dict(exclude_none=True))
+                    pub_root.source = RootSource.PANGEA
+            if pub_root is not None:
+                self.pub_roots[tree_size] = pub_root

pangea/services/audit/exceptions.py

@@ -1,6 +1,5 @@
 from pangea.exceptions import PangeaException
-
-from .models import EventEnvelope
+from pangea.services.audit.models import EventEnvelope
 
 
 # Audit SDK Specific Exceptions