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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (52) hide show
  1. pangea/__init__.py +1 -1
  2. pangea/asyncio/file_uploader.py +1 -1
  3. pangea/asyncio/request.py +56 -34
  4. pangea/asyncio/services/__init__.py +4 -0
  5. pangea/asyncio/services/ai_guard.py +75 -0
  6. pangea/asyncio/services/audit.py +192 -31
  7. pangea/asyncio/services/authn.py +187 -109
  8. pangea/asyncio/services/authz.py +285 -0
  9. pangea/asyncio/services/base.py +21 -2
  10. pangea/asyncio/services/embargo.py +2 -2
  11. pangea/asyncio/services/file_scan.py +24 -9
  12. pangea/asyncio/services/intel.py +108 -34
  13. pangea/asyncio/services/prompt_guard.py +73 -0
  14. pangea/asyncio/services/redact.py +72 -4
  15. pangea/asyncio/services/sanitize.py +217 -0
  16. pangea/asyncio/services/share.py +246 -73
  17. pangea/asyncio/services/vault.py +1710 -750
  18. pangea/crypto/rsa.py +135 -0
  19. pangea/deep_verify.py +7 -1
  20. pangea/dump_audit.py +9 -8
  21. pangea/request.py +87 -59
  22. pangea/response.py +49 -31
  23. pangea/services/__init__.py +4 -0
  24. pangea/services/ai_guard.py +128 -0
  25. pangea/services/audit/audit.py +205 -42
  26. pangea/services/audit/models.py +56 -8
  27. pangea/services/audit/signing.py +6 -5
  28. pangea/services/audit/util.py +3 -3
  29. pangea/services/authn/authn.py +140 -70
  30. pangea/services/authn/models.py +167 -11
  31. pangea/services/authz.py +400 -0
  32. pangea/services/base.py +39 -8
  33. pangea/services/embargo.py +2 -2
  34. pangea/services/file_scan.py +32 -15
  35. pangea/services/intel.py +157 -32
  36. pangea/services/prompt_guard.py +83 -0
  37. pangea/services/redact.py +152 -4
  38. pangea/services/sanitize.py +371 -0
  39. pangea/services/share/share.py +683 -107
  40. pangea/services/vault/models/asymmetric.py +120 -18
  41. pangea/services/vault/models/common.py +439 -141
  42. pangea/services/vault/models/keys.py +94 -0
  43. pangea/services/vault/models/secret.py +27 -3
  44. pangea/services/vault/models/symmetric.py +68 -22
  45. pangea/services/vault/vault.py +1690 -749
  46. pangea/tools.py +6 -7
  47. pangea/utils.py +16 -27
  48. pangea/verify_audit.py +270 -83
  49. {pangea_sdk-3.8.0b1.dist-info → pangea_sdk-5.4.0b1.dist-info}/METADATA +43 -35
  50. pangea_sdk-5.4.0b1.dist-info/RECORD +60 -0
  51. {pangea_sdk-3.8.0b1.dist-info → pangea_sdk-5.4.0b1.dist-info}/WHEEL +1 -1
  52. pangea_sdk-3.8.0b1.dist-info/RECORD +0 -50
@@ -0,0 +1,128 @@
1
+ from __future__ import annotations
2
+
3
+ from typing import Any, Dict, Generic, List, Optional, TypeVar
4
+
5
+ from pangea.config import PangeaConfig
6
+ from pangea.response import APIResponseModel, PangeaResponse, PangeaResponseResult
7
+ from pangea.services.base import ServiceBase
8
+
9
+
10
+ class AnalyzerResponse(APIResponseModel):
11
+ analyzer: str
12
+ confidence: float
13
+
14
+
15
+ class PromptInjectionResult(APIResponseModel):
16
+ analyzer_responses: List[AnalyzerResponse]
17
+ """Triggered prompt injection analyzers."""
18
+
19
+
20
+ class PiiEntity(APIResponseModel):
21
+ type: str
22
+ value: str
23
+ redacted: bool
24
+ start_pos: Optional[int] = None
25
+
26
+
27
+ class PiiEntityResult(APIResponseModel):
28
+ entities: List[PiiEntity]
29
+
30
+
31
+ class MaliciousEntity(APIResponseModel):
32
+ type: str
33
+ value: str
34
+ redacted: Optional[bool] = None
35
+ start_pos: Optional[int] = None
36
+ raw: Optional[Dict[str, Any]] = None
37
+
38
+
39
+ class MaliciousEntityResult(APIResponseModel):
40
+ entities: List[MaliciousEntity]
41
+
42
+
43
+ T = TypeVar("T")
44
+
45
+
46
+ class TextGuardDetector(APIResponseModel, Generic[T]):
47
+ detected: bool
48
+ data: Optional[T] = None
49
+
50
+
51
+ class TextGuardDetectors(APIResponseModel):
52
+ prompt_injection: Optional[TextGuardDetector[PromptInjectionResult]] = None
53
+ pii_entity: Optional[TextGuardDetector[PiiEntityResult]] = None
54
+ malicious_entity: Optional[TextGuardDetector[MaliciousEntityResult]] = None
55
+
56
+
57
+ class TextGuardResult(PangeaResponseResult):
58
+ detectors: TextGuardDetectors
59
+ prompt: str
60
+
61
+
62
+ class AIGuard(ServiceBase):
63
+ """AI Guard service client.
64
+
65
+ Provides methods to interact with Pangea's AI Guard service.
66
+
67
+ Examples:
68
+ from pangea import PangeaConfig
69
+ from pangea.services import AIGuard
70
+
71
+ config = PangeaConfig(domain="aws.us.pangea.cloud")
72
+ ai_guard = AIGuard(token="pangea_token", config=config)
73
+ """
74
+
75
+ service_name = "ai-guard"
76
+
77
+ def __init__(
78
+ self, token: str, config: PangeaConfig | None = None, logger_name: str = "pangea", config_id: str | None = None
79
+ ) -> None:
80
+ """
81
+ AI Guard service client.
82
+
83
+ Initializes a new AI Guard client.
84
+
85
+ Args:
86
+ token: Pangea API token.
87
+ config: Pangea service configuration.
88
+ logger_name: Logger name.
89
+ config_id: Configuration ID.
90
+
91
+ Examples:
92
+ from pangea import PangeaConfig
93
+ from pangea.services import AIGuard
94
+
95
+ config = PangeaConfig(domain="aws.us.pangea.cloud")
96
+ ai_guard = AIGuard(token="pangea_token", config=config)
97
+ """
98
+
99
+ super().__init__(token, config, logger_name, config_id)
100
+
101
+ def guard_text(
102
+ self,
103
+ text: str,
104
+ *,
105
+ recipe: str = "pangea_prompt_guard",
106
+ debug: bool = False,
107
+ ) -> PangeaResponse[TextGuardResult]:
108
+ """
109
+ Text guard (Beta)
110
+
111
+ Guard text.
112
+
113
+ How to install a [Beta release](https://pangea.cloud/docs/sdk/python/#beta-releases).
114
+
115
+ OperationId: ai_guard_post_v1beta_text_guard
116
+
117
+ Args:
118
+ text: Text.
119
+ recipe: Recipe.
120
+ debug: Debug.
121
+
122
+ Examples:
123
+ response = ai_guard.guard_text("text")
124
+ """
125
+
126
+ return self.request.post(
127
+ "v1beta/text/guard", TextGuardResult, data={"text": text, "recipe": recipe, "debug": debug}
128
+ )
@@ -1,11 +1,14 @@
1
1
  # Copyright 2022 Pangea Cyber Corporation
2
2
  # Author: Pangea Cyber Corporation
3
+ from __future__ import annotations
4
+
3
5
  import datetime
4
6
  import json
5
- from typing import Any, Dict, List, Optional, Set, Tuple, Union
7
+ from typing import Any, Dict, Iterable, List, Optional, Sequence, Set, Tuple, Union
6
8
 
7
9
  import pangea.exceptions as pexc
8
- from pangea.response import PangeaResponse
10
+ from pangea.config import PangeaConfig
11
+ from pangea.response import PangeaResponse, PangeaResponseResult
9
12
  from pangea.services.audit.exceptions import AuditException, EventCorruption
10
13
  from pangea.services.audit.models import (
11
14
  DownloadFormat,
@@ -14,6 +17,7 @@ from pangea.services.audit.models import (
14
17
  Event,
15
18
  EventEnvelope,
16
19
  EventVerification,
20
+ ExportRequest,
17
21
  LogBulkRequest,
18
22
  LogBulkResult,
19
23
  LogEvent,
@@ -51,8 +55,8 @@ from pangea.utils import canonicalize_nested_json
51
55
 
52
56
  class AuditBase:
53
57
  def __init__(
54
- self, private_key_file: str = "", public_key_info: Dict[str, str] = {}, tenant_id: Optional[str] = None
55
- ):
58
+ self, private_key_file: str = "", public_key_info: dict[str, str] = {}, tenant_id: str | None = None
59
+ ) -> None:
56
60
  self.pub_roots: Dict[int, PublishedRoot] = {}
57
61
  self.buffer_data: Optional[str] = None
58
62
  self.signer: Optional[Signer] = Signer(private_key_file) if private_key_file else None
@@ -195,7 +199,7 @@ class AuditBase:
195
199
 
196
200
  # verify consistency proofs
197
201
  if self.can_verify_consistency_proof(search_event):
198
- if self.verify_consistency_proof(self.pub_roots, search_event):
202
+ if search_event.leaf_index is not None and self.verify_consistency_proof(search_event.leaf_index):
199
203
  search_event.consistency_verification = EventVerification.PASS
200
204
  else:
201
205
  search_event.consistency_verification = EventVerification.FAIL
@@ -218,7 +222,7 @@ class AuditBase:
218
222
  tree_sizes.difference_update(self.pub_roots.keys())
219
223
 
220
224
  if tree_sizes:
221
- arweave_roots = get_arweave_published_roots(result.root.tree_name, list(tree_sizes))
225
+ arweave_roots = get_arweave_published_roots(result.root.tree_name, tree_sizes)
222
226
  else:
223
227
  arweave_roots = {}
224
228
 
@@ -280,7 +284,7 @@ class AuditBase:
280
284
  """
281
285
  return event.published and event.leaf_index is not None and event.leaf_index >= 0 # type: ignore[return-value]
282
286
 
283
- def verify_consistency_proof(self, pub_roots: Dict[int, PublishedRoot], event: SearchEvent) -> bool:
287
+ def verify_consistency_proof(self, tree_size: int) -> bool:
284
288
  """
285
289
  Verify consistency proof
286
290
 
@@ -289,18 +293,17 @@ class AuditBase:
289
293
  Read more at: [What is a consistency proof?](https://pangea.cloud/docs/audit/merkle-trees#what-is-a-consistency-proof)
290
294
 
291
295
  Args:
292
- pub_roots (dict[int, Root]): list of published root hashes across time
293
- event (SearchEvent): Audit event to be verified.
296
+ leaf_index (int): The tree size of the root to be verified.
294
297
 
295
298
  Returns:
296
299
  bool: True if consistency proof is verified, False otherwise.
297
300
  """
298
301
 
299
- if event.leaf_index == 0:
302
+ if tree_size == 0:
300
303
  return True
301
304
 
302
- curr_root = pub_roots.get(event.leaf_index + 1) # type: ignore[operator]
303
- prev_root = pub_roots.get(event.leaf_index) # type: ignore[arg-type]
305
+ curr_root = self.pub_roots.get(tree_size + 1)
306
+ prev_root = self.pub_roots.get(tree_size)
304
307
 
305
308
  if not curr_root or not prev_root:
306
309
  return False
@@ -310,9 +313,12 @@ class AuditBase:
310
313
  ):
311
314
  return False
312
315
 
316
+ if curr_root.consistency_proof is None:
317
+ return False
318
+
313
319
  curr_root_hash = decode_hash(curr_root.root_hash)
314
320
  prev_root_hash = decode_hash(prev_root.root_hash)
315
- proof = decode_consistency_proof(curr_root.consistency_proof) # type: ignore[arg-type]
321
+ proof = decode_consistency_proof(curr_root.consistency_proof)
316
322
 
317
323
  return verify_consistency_proof(curr_root_hash, prev_root_hash, proof)
318
324
 
@@ -332,7 +338,9 @@ class AuditBase:
332
338
  if audit_envelope and audit_envelope.signature and public_key:
333
339
  v = Verifier()
334
340
  verification = v.verify_signature(
335
- audit_envelope.signature, canonicalize_event(audit_envelope.event), public_key # type: ignore[arg-type]
341
+ audit_envelope.signature,
342
+ canonicalize_event(Event(**audit_envelope.event)),
343
+ public_key,
336
344
  )
337
345
  if verification is not None:
338
346
  return EventVerification.PASS if verification else EventVerification.FAIL
@@ -374,14 +382,32 @@ class Audit(ServiceBase, AuditBase):
374
382
 
375
383
  def __init__(
376
384
  self,
377
- token,
378
- config=None,
385
+ token: str,
386
+ config: PangeaConfig | None = None,
379
387
  private_key_file: str = "",
380
- public_key_info: Dict[str, str] = {},
381
- tenant_id: Optional[str] = None,
382
- logger_name="pangea",
383
- config_id: Optional[str] = None,
384
- ):
388
+ public_key_info: dict[str, str] = {},
389
+ tenant_id: str | None = None,
390
+ logger_name: str = "pangea",
391
+ config_id: str | None = None,
392
+ ) -> None:
393
+ """
394
+ Audit client
395
+
396
+ Initializes a new Audit client.
397
+
398
+ Args:
399
+ token: Pangea API token.
400
+ config: Configuration.
401
+ private_key_file: Private key filepath.
402
+ public_key_info: Public key information.
403
+ tenant_id: Tenant ID.
404
+ logger_name: Logger name.
405
+ config_id: Configuration ID.
406
+
407
+ Examples:
408
+ config = PangeaConfig(domain="pangea_domain")
409
+ audit = Audit(token="pangea_token", config=config)
410
+ """
385
411
  # FIXME: Temporary check to deprecate config_id from PangeaConfig.
386
412
  # Delete it when deprecate PangeaConfig.config_id
387
413
  if config_id and config is not None and config.config_id is not None:
@@ -466,7 +492,7 @@ class Audit(ServiceBase, AuditBase):
466
492
  verbose: Optional[bool] = None,
467
493
  ) -> PangeaResponse[LogResult]:
468
494
  """
469
- Log an entry
495
+ Log an event
470
496
 
471
497
  Create a log entry in the Secure Audit Log.
472
498
 
@@ -475,6 +501,7 @@ class Audit(ServiceBase, AuditBase):
475
501
  verify (bool, optional): True to verify logs consistency after response.
476
502
  sign_local (bool, optional): True to sign event with local key.
477
503
  verbose (bool, optional): True to get a more verbose response.
504
+
478
505
  Raises:
479
506
  AuditException: If an audit based api exception happens
480
507
  PangeaAPIException: If an API Error happens
@@ -485,17 +512,13 @@ class Audit(ServiceBase, AuditBase):
485
512
  Available response fields can be found in our [API documentation](https://pangea.cloud/docs/api/audit#/v1/log).
486
513
 
487
514
  Examples:
488
- try:
489
- log_response = audit.log({"message": "hello world"}, verbose=True)
490
- print(f"Response. Hash: {log_response.result.hash}")
491
- except pe.PangeaAPIException as e:
492
- print(f"Request Error: {e.response.summary}")
493
- for err in e.errors:
494
- print(f"\\t{err.detail} \\n")
515
+ response = audit.log_event({"message": "hello world"}, verbose=True)
495
516
  """
496
517
 
497
518
  input = self._get_log_request(event, sign_local=sign_local, verify=verify, verbose=verbose)
498
- response: PangeaResponse[LogResult] = self.request.post("v1/log", LogResult, data=input.dict(exclude_none=True))
519
+ response: PangeaResponse[LogResult] = self.request.post(
520
+ "v1/log", LogResult, data=input.model_dump(exclude_none=True)
521
+ )
499
522
  if response.success and response.result is not None:
500
523
  self._process_log_result(response.result, verify=verify)
501
524
  return response
@@ -535,7 +558,7 @@ class Audit(ServiceBase, AuditBase):
535
558
 
536
559
  input = self._get_log_request(events, sign_local=sign_local, verify=False, verbose=verbose)
537
560
  response: PangeaResponse[LogBulkResult] = self.request.post(
538
- "v2/log", LogBulkResult, data=input.dict(exclude_none=True)
561
+ "v2/log", LogBulkResult, data=input.model_dump(exclude_none=True)
539
562
  )
540
563
 
541
564
  if response.success and response.result is not None:
@@ -579,7 +602,7 @@ class Audit(ServiceBase, AuditBase):
579
602
  try:
580
603
  # Calling to v2 methods will return always a 202.
581
604
  response: PangeaResponse[LogBulkResult] = self.request.post(
582
- "v2/log_async", LogBulkResult, data=input.dict(exclude_none=True), poll_result=False
605
+ "v2/log_async", LogBulkResult, data=input.model_dump(exclude_none=True), poll_result=False
583
606
  )
584
607
  except pexc.AcceptedRequestException as e:
585
608
  return e.response
@@ -598,10 +621,11 @@ class Audit(ServiceBase, AuditBase):
598
621
  end: Optional[Union[datetime.datetime, str]] = None,
599
622
  limit: Optional[int] = None,
600
623
  max_results: Optional[int] = None,
601
- search_restriction: Optional[dict] = None,
624
+ search_restriction: Optional[Dict[str, Sequence[str]]] = None,
602
625
  verbose: Optional[bool] = None,
603
626
  verify_consistency: bool = False,
604
627
  verify_events: bool = True,
628
+ return_context: Optional[bool] = None,
605
629
  ) -> PangeaResponse[SearchOutput]:
606
630
  """
607
631
  Search the log
@@ -627,10 +651,11 @@ class Audit(ServiceBase, AuditBase):
627
651
  end (datetime, optional): An RFC-3339 formatted timestamp, or relative time adjustment from the current time.
628
652
  limit (int, optional): Optional[int] = None,
629
653
  max_results (int, optional): Maximum number of results to return.
630
- search_restriction (dict, optional): A list of keys to restrict the search results to. Useful for partitioning data available to the query string.
654
+ 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.
631
655
  verbose (bool, optional): If true, response include root and membership and consistency proofs.
632
656
  verify_consistency (bool): True to verify logs consistency
633
657
  verify_events (bool): True to verify hash events and signatures
658
+ return_context (bool): Return the context data needed to decrypt secure audit events that have been redacted with format preserving encryption.
634
659
 
635
660
  Raises:
636
661
  AuditException: If an audit based api exception happens
@@ -664,10 +689,11 @@ class Audit(ServiceBase, AuditBase):
664
689
  max_results=max_results,
665
690
  search_restriction=search_restriction,
666
691
  verbose=verbose,
692
+ return_context=return_context,
667
693
  )
668
694
 
669
695
  response: PangeaResponse[SearchOutput] = self.request.post(
670
- "v1/search", SearchOutput, data=input.dict(exclude_none=True)
696
+ "v1/search", SearchOutput, data=input.model_dump(exclude_none=True)
671
697
  )
672
698
  if verify_consistency and response.result is not None:
673
699
  self.update_published_roots(response.result)
@@ -678,8 +704,10 @@ class Audit(ServiceBase, AuditBase):
678
704
  id: str,
679
705
  limit: Optional[int] = 20,
680
706
  offset: Optional[int] = 0,
707
+ assert_search_restriction: Optional[Dict[str, Sequence[str]]] = None,
681
708
  verify_consistency: bool = False,
682
709
  verify_events: bool = True,
710
+ return_context: Optional[bool] = None,
683
711
  ) -> PangeaResponse[SearchResultOutput]:
684
712
  """
685
713
  Results of a search
@@ -692,8 +720,10 @@ class Audit(ServiceBase, AuditBase):
692
720
  id (string): the id of a search action, found in `response.result.id`
693
721
  limit (integer, optional): the maximum number of results to return, default is 20
694
722
  offset (integer, optional): the position of the first result to return, default is 0
723
+ 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.
695
724
  verify_consistency (bool): True to verify logs consistency
696
725
  verify_events (bool): True to verify hash events and signatures
726
+ return_context (bool): Return the context data needed to decrypt secure audit events that have been redacted with format preserving encryption.
697
727
  Raises:
698
728
  AuditException: If an audit based api exception happens
699
729
  PangeaAPIException: If an API Error happens
@@ -716,14 +746,117 @@ class Audit(ServiceBase, AuditBase):
716
746
  id=id,
717
747
  limit=limit,
718
748
  offset=offset,
749
+ assert_search_restriction=assert_search_restriction,
750
+ return_context=return_context,
719
751
  )
720
752
  response: PangeaResponse[SearchResultOutput] = self.request.post(
721
- "v1/results", SearchResultOutput, data=input.dict(exclude_none=True)
753
+ "v1/results", SearchResultOutput, data=input.model_dump(exclude_none=True)
722
754
  )
723
755
  if verify_consistency and response.result is not None:
724
756
  self.update_published_roots(response.result)
725
757
  return self.handle_results_response(response, verify_consistency, verify_events)
726
758
 
759
+ def export(
760
+ self,
761
+ *,
762
+ format: DownloadFormat = DownloadFormat.CSV,
763
+ start: Optional[datetime.datetime] = None,
764
+ end: Optional[datetime.datetime] = None,
765
+ order: Optional[SearchOrder] = None,
766
+ order_by: Optional[str] = None,
767
+ verbose: bool = True,
768
+ ) -> PangeaResponse[PangeaResponseResult]:
769
+ """
770
+ Export from the audit log
771
+
772
+ Bulk export of data from the Secure Audit Log, with optional filtering.
773
+
774
+ OperationId: audit_post_v1_export
775
+
776
+ Args:
777
+ format: Format for the records.
778
+ start: The start of the time range to perform the search on.
779
+ end: The end of the time range to perform the search on. If omitted,
780
+ then all records up to the latest will be searched.
781
+ order: Specify the sort order of the response.
782
+ order_by: Name of column to sort the results by.
783
+ verbose: Whether or not to include the root hash of the tree and the
784
+ membership proof for each record.
785
+
786
+ Raises:
787
+ AuditException: If an audit based api exception happens
788
+ PangeaAPIException: If an API Error happens
789
+
790
+ Examples:
791
+ export_res = audit.export(verbose=False)
792
+
793
+ # Export may take several dozens of minutes, so polling for the result
794
+ # should be done in a loop. That is omitted here for brevity's sake.
795
+ try:
796
+ audit.poll_result(request_id=export_res.request_id)
797
+ except AcceptedRequestException:
798
+ # Retry later.
799
+
800
+ # Download the result when it's ready.
801
+ download_res = audit.download_results(request_id=export_res.request_id)
802
+ download_res.result.dest_url
803
+ # => https://pangea-runtime.s3.amazonaws.com/audit/xxxxx/search_results_[...]
804
+ """
805
+ input = ExportRequest(
806
+ format=format,
807
+ start=start,
808
+ end=end,
809
+ order=order,
810
+ order_by=order_by,
811
+ verbose=verbose,
812
+ )
813
+ try:
814
+ return self.request.post(
815
+ "v1/export", PangeaResponseResult, data=input.model_dump(exclude_none=True), poll_result=False
816
+ )
817
+ except pexc.AcceptedRequestException as e:
818
+ return e.response
819
+
820
+ def log_stream(self, data: dict) -> PangeaResponse[PangeaResponseResult]:
821
+ """
822
+ Log streaming endpoint
823
+
824
+ This API allows 3rd party vendors (like Auth0) to stream events to this
825
+ endpoint where the structure of the payload varies across different
826
+ vendors.
827
+
828
+ OperationId: audit_post_v1_log_stream
829
+
830
+ Args:
831
+ data: Event data. The exact schema of this will vary by vendor.
832
+
833
+ Raises:
834
+ AuditException: If an audit based api exception happens
835
+ PangeaAPIException: If an API Error happens
836
+
837
+ Examples:
838
+ data = {
839
+ "logs": [
840
+ {
841
+ "log_id": "some log ID",
842
+ "data": {
843
+ "date": "2024-03-29T17:26:50.193Z",
844
+ "type": "sapi",
845
+ "description": "Create a log stream",
846
+ "client_id": "some client ID",
847
+ "ip": "127.0.0.1",
848
+ "user_agent": "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0",
849
+ "user_id": "some user ID",
850
+ },
851
+ }
852
+ # ...
853
+ ]
854
+ }
855
+
856
+ response = audit.log_stream(data)
857
+ """
858
+ return self.request.post("v1/log_stream", PangeaResponseResult, data=data)
859
+
727
860
  def root(self, tree_size: Optional[int] = None) -> PangeaResponse[RootResult]:
728
861
  """
729
862
  Tamperproof verification
@@ -746,10 +879,14 @@ class Audit(ServiceBase, AuditBase):
746
879
  response = audit.root(tree_size=7)
747
880
  """
748
881
  input = RootRequest(tree_size=tree_size)
749
- return self.request.post("v1/root", RootResult, data=input.dict(exclude_none=True))
882
+ return self.request.post("v1/root", RootResult, data=input.model_dump(exclude_none=True))
750
883
 
751
884
  def download_results(
752
- self, result_id: str, format: Optional[DownloadFormat] = None
885
+ self,
886
+ result_id: Optional[str] = None,
887
+ format: DownloadFormat = DownloadFormat.CSV,
888
+ request_id: Optional[str] = None,
889
+ return_context: Optional[bool] = None,
753
890
  ) -> PangeaResponse[DownloadResult]:
754
891
  """
755
892
  Download search results
@@ -761,6 +898,8 @@ class Audit(ServiceBase, AuditBase):
761
898
  Args:
762
899
  result_id: ID returned by the search API.
763
900
  format: Format for the records.
901
+ request_id: ID returned by the export API.
902
+ return_context (bool): Return the context data needed to decrypt secure audit events that have been redacted with format preserving encryption.
764
903
 
765
904
  Returns:
766
905
  URL where search results can be downloaded.
@@ -776,8 +915,13 @@ class Audit(ServiceBase, AuditBase):
776
915
  )
777
916
  """
778
917
 
779
- input = DownloadRequest(result_id=result_id, format=format)
780
- return self.request.post("v1/download_results", DownloadResult, data=input.dict(exclude_none=True))
918
+ if request_id is None and result_id is None:
919
+ raise ValueError("must pass one of `request_id` or `result_id`")
920
+
921
+ input = DownloadRequest(
922
+ request_id=request_id, result_id=result_id, format=format, return_context=return_context
923
+ )
924
+ return self.request.post("v1/download_results", DownloadResult, data=input.model_dump(exclude_none=True))
781
925
 
782
926
  def update_published_roots(self, result: SearchResultOutput):
783
927
  """Fetches series of published root hashes from Arweave
@@ -802,12 +946,31 @@ class Audit(ServiceBase, AuditBase):
802
946
  for tree_size in tree_sizes:
803
947
  pub_root = None
804
948
  if tree_size in arweave_roots:
805
- pub_root = PublishedRoot(**arweave_roots[tree_size].dict(exclude_none=True))
949
+ pub_root = PublishedRoot(**arweave_roots[tree_size].model_dump(exclude_none=True))
806
950
  pub_root.source = RootSource.ARWEAVE
807
951
  elif self.allow_server_roots:
808
952
  resp = self.root(tree_size=tree_size)
809
953
  if resp.success and resp.result is not None:
810
- pub_root = PublishedRoot(**resp.result.data.dict(exclude_none=True))
954
+ pub_root = PublishedRoot(**resp.result.data.model_dump(exclude_none=True))
811
955
  pub_root.source = RootSource.PANGEA
812
956
  if pub_root is not None:
813
957
  self.pub_roots[tree_size] = pub_root
958
+
959
+ self._fix_consistency_proofs(tree_sizes)
960
+
961
+ def _fix_consistency_proofs(self, tree_sizes: Iterable[int]) -> None:
962
+ # on very rare occasions, the consistency proof in Arweave may be wrong
963
+ # override it with the proof from pangea (not the root hash, just the proof)
964
+ for tree_size in tree_sizes:
965
+ if tree_size not in self.pub_roots or tree_size - 1 not in self.pub_roots:
966
+ continue
967
+
968
+ if self.pub_roots[tree_size].source == RootSource.PANGEA:
969
+ continue
970
+
971
+ if self.verify_consistency_proof(tree_size):
972
+ continue
973
+
974
+ resp = self.root(tree_size=tree_size)
975
+ if resp.success and resp.result is not None and resp.result.data is not None:
976
+ self.pub_roots[tree_size].consistency_proof = resp.result.data.consistency_proof