pangea-sdk 6.1.1__py3-none-any.whl → 6.2.0__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.
- pangea/__init__.py +9 -1
- pangea/asyncio/__init__.py +1 -0
- pangea/asyncio/file_uploader.py +4 -2
- pangea/asyncio/request.py +52 -17
- pangea/asyncio/services/__init__.py +2 -0
- pangea/asyncio/services/ai_guard.py +9 -12
- pangea/asyncio/services/audit.py +6 -1
- pangea/asyncio/services/authn.py +15 -4
- pangea/asyncio/services/base.py +4 -0
- pangea/asyncio/services/file_scan.py +7 -1
- pangea/asyncio/services/intel.py +26 -28
- pangea/asyncio/services/redact.py +4 -0
- pangea/asyncio/services/sanitize.py +5 -1
- pangea/asyncio/services/share.py +5 -1
- pangea/asyncio/services/vault.py +4 -0
- pangea/audit_logger.py +3 -1
- pangea/deep_verify.py +13 -13
- pangea/deprecated.py +1 -1
- pangea/dump_audit.py +2 -3
- pangea/exceptions.py +8 -5
- pangea/file_uploader.py +4 -0
- pangea/request.py +63 -47
- pangea/response.py +21 -18
- pangea/services/__init__.py +2 -0
- pangea/services/ai_guard.py +35 -24
- pangea/services/audit/audit.py +10 -7
- pangea/services/audit/models.py +71 -34
- pangea/services/audit/signing.py +1 -1
- pangea/services/audit/util.py +10 -10
- pangea/services/authn/authn.py +15 -4
- pangea/services/authn/models.py +10 -56
- pangea/services/authz.py +4 -0
- pangea/services/base.py +7 -4
- pangea/services/embargo.py +6 -0
- pangea/services/file_scan.py +7 -1
- pangea/services/intel.py +36 -19
- pangea/services/redact.py +4 -0
- pangea/services/sanitize.py +5 -1
- pangea/services/share/share.py +13 -7
- pangea/services/vault/models/asymmetric.py +4 -0
- pangea/services/vault/models/common.py +4 -0
- pangea/services/vault/models/symmetric.py +4 -0
- pangea/services/vault/vault.py +2 -4
- pangea/tools.py +13 -9
- pangea/utils.py +3 -5
- pangea/verify_audit.py +23 -27
- {pangea_sdk-6.1.1.dist-info → pangea_sdk-6.2.0.dist-info}/METADATA +36 -17
- pangea_sdk-6.2.0.dist-info/RECORD +60 -0
- pangea_sdk-6.1.1.dist-info/RECORD +0 -60
- {pangea_sdk-6.1.1.dist-info → pangea_sdk-6.2.0.dist-info}/WHEEL +0 -0
pangea/deep_verify.py
CHANGED
@@ -1,13 +1,16 @@
|
|
1
1
|
# Copyright 2022 Pangea Cyber Corporation
|
2
2
|
# Author: Pangea Cyber Corporation
|
3
3
|
|
4
|
+
from __future__ import annotations
|
5
|
+
|
4
6
|
import argparse
|
5
7
|
import io
|
6
8
|
import math
|
7
9
|
import os
|
8
10
|
import sys
|
11
|
+
from collections.abc import Iterator
|
9
12
|
from itertools import groupby
|
10
|
-
from typing import
|
13
|
+
from typing import Optional, TypedDict, Union
|
11
14
|
|
12
15
|
import pangea.services.audit.util as audit_util
|
13
16
|
from pangea.services import Audit
|
@@ -24,7 +27,7 @@ class Errors(TypedDict):
|
|
24
27
|
buffer_missing: int
|
25
28
|
|
26
29
|
|
27
|
-
root_hashes:
|
30
|
+
root_hashes: dict[int, str] = {}
|
28
31
|
|
29
32
|
|
30
33
|
def num_lines(f: io.TextIOWrapper) -> int:
|
@@ -124,10 +127,7 @@ def get_root_hash(audit: Audit, tree_size: int) -> str:
|
|
124
127
|
|
125
128
|
|
126
129
|
def print_error(msg: str, level: str = "error"):
|
127
|
-
if level == "warning"
|
128
|
-
dot = "🟡"
|
129
|
-
else:
|
130
|
-
dot = "🔴"
|
130
|
+
dot = "🟡" if level == "warning" else "🔴"
|
131
131
|
print(f"{dot} {msg:200s}")
|
132
132
|
|
133
133
|
|
@@ -148,14 +148,14 @@ def deep_verify(audit: Audit, file: io.TextIOWrapper) -> Errors:
|
|
148
148
|
}
|
149
149
|
|
150
150
|
events = file_events(root_hashes, file)
|
151
|
-
events_by_idx: Union[
|
151
|
+
events_by_idx: Union[list[Event], Iterator[Event]]
|
152
152
|
cold_indexes = SequenceFollower()
|
153
153
|
for leaf_index, events_by_idx in groupby(events, lambda event: event.get("leaf_index")):
|
154
154
|
events_by_idx = list(events_by_idx)
|
155
155
|
buffer_lines = (cnt, cnt + len(events_by_idx) - 1)
|
156
156
|
if leaf_index is None:
|
157
157
|
print_error(
|
158
|
-
f"Lines {buffer_lines[0]}-{buffer_lines[1]} ({buffer_lines[1]-buffer_lines[0]+1}): Buffer was not persisted"
|
158
|
+
f"Lines {buffer_lines[0]}-{buffer_lines[1]} ({buffer_lines[1] - buffer_lines[0] + 1}): Buffer was not persisted"
|
159
159
|
)
|
160
160
|
errors["not_persisted"] += len(events_by_idx)
|
161
161
|
cnt += len(events_by_idx)
|
@@ -164,8 +164,8 @@ def deep_verify(audit: Audit, file: io.TextIOWrapper) -> Errors:
|
|
164
164
|
cold_indexes.add(leaf_index)
|
165
165
|
|
166
166
|
cold_path_size: Optional[int] = None
|
167
|
-
hot_indexes:
|
168
|
-
for
|
167
|
+
hot_indexes: set[int] = set()
|
168
|
+
for _i, event in enumerate(events_by_idx):
|
169
169
|
cnt += 1
|
170
170
|
tree_size = get_tree_size(event)
|
171
171
|
if tree_size not in root_hashes:
|
@@ -203,11 +203,11 @@ def deep_verify(audit: Audit, file: io.TextIOWrapper) -> Errors:
|
|
203
203
|
errors["missing"] += len(hot_indexes_diff)
|
204
204
|
print(f"missing hot indexes: {hot_indexes_diff}")
|
205
205
|
print(f"hot_indexes: {hot_indexes} ")
|
206
|
-
print(
|
206
|
+
print("events:")
|
207
207
|
for e in events_by_idx:
|
208
208
|
print(e)
|
209
209
|
print_error(
|
210
|
-
f"Lines {buffer_lines[0]}-{buffer_lines[1]} ({buffer_lines[1]-buffer_lines[0]}), Buffer #{cold_idx}: {len(hot_indexes_diff)} event(s) missing"
|
210
|
+
f"Lines {buffer_lines[0]}-{buffer_lines[1]} ({buffer_lines[1] - buffer_lines[0]}), Buffer #{cold_idx}: {len(hot_indexes_diff)} event(s) missing"
|
211
211
|
)
|
212
212
|
|
213
213
|
cold_holes = cold_indexes.holes()
|
@@ -232,7 +232,7 @@ def create_parser():
|
|
232
232
|
"-f",
|
233
233
|
required=True,
|
234
234
|
type=argparse.FileType("r"),
|
235
|
-
help="Event input file. Must be a collection of
|
235
|
+
help="Event input file. Must be a collection of JSON Objects separated by newlines",
|
236
236
|
)
|
237
237
|
return parser
|
238
238
|
|
pangea/deprecated.py
CHANGED
pangea/dump_audit.py
CHANGED
@@ -7,7 +7,6 @@ import json
|
|
7
7
|
import os
|
8
8
|
import sys
|
9
9
|
from datetime import datetime
|
10
|
-
from typing import Tuple
|
11
10
|
|
12
11
|
import dateutil.parser
|
13
12
|
|
@@ -103,7 +102,7 @@ def dump_after(audit: Audit, output: io.TextIOWrapper, start: datetime, last_eve
|
|
103
102
|
|
104
103
|
def dump_page(
|
105
104
|
audit: Audit, output: io.TextIOWrapper, start: datetime, end: datetime, first: bool = False
|
106
|
-
) ->
|
105
|
+
) -> tuple[datetime, int, bool, str, int]:
|
107
106
|
PAGE_SIZE = 1000
|
108
107
|
print(start, end)
|
109
108
|
print("Dumping...")
|
@@ -175,7 +174,7 @@ def parse_args(parser: argparse.ArgumentParser):
|
|
175
174
|
raise ValueError("domain missing")
|
176
175
|
|
177
176
|
if args.output is None:
|
178
|
-
args.output = open(f"dump-{datetime.now().strftime('%Y%m%d%H%M%S')}.jsonl", "w")
|
177
|
+
args.output = open(f"dump-{datetime.now().strftime('%Y%m%d%H%M%S')}.jsonl", "w") # noqa: SIM115
|
179
178
|
|
180
179
|
args.start = make_aware_datetime(args.start)
|
181
180
|
args.end = make_aware_datetime(args.end)
|
pangea/exceptions.py
CHANGED
@@ -1,6 +1,9 @@
|
|
1
1
|
# Copyright 2022 Pangea Cyber Corporation
|
2
2
|
# Author: Pangea Cyber Corporation
|
3
3
|
|
4
|
+
# TODO: Modernize.
|
5
|
+
# ruff: noqa: UP006, UP035
|
6
|
+
|
4
7
|
from typing import List, Optional
|
5
8
|
|
6
9
|
from pangea.response import AcceptedResult, ErrorField, PangeaResponse
|
@@ -44,7 +47,7 @@ class PangeaAPIException(PangeaException):
|
|
44
47
|
response: PangeaResponse
|
45
48
|
|
46
49
|
def __init__(self, message: str, response: PangeaResponse):
|
47
|
-
super(
|
50
|
+
super().__init__(message)
|
48
51
|
self.response = response
|
49
52
|
|
50
53
|
@property
|
@@ -92,7 +95,7 @@ class UnauthorizedException(PangeaAPIException):
|
|
92
95
|
|
93
96
|
def __init__(self, service_name: str, response: PangeaResponse):
|
94
97
|
message = f"User is not authorized to access service {service_name}"
|
95
|
-
super(
|
98
|
+
super().__init__(message, response)
|
96
99
|
|
97
100
|
|
98
101
|
class NotFound(PangeaAPIException):
|
@@ -100,20 +103,20 @@ class NotFound(PangeaAPIException):
|
|
100
103
|
|
101
104
|
def __init__(self, url: str, response: PangeaResponse):
|
102
105
|
message = f"Resource url:'{url}' not found"
|
103
|
-
super(
|
106
|
+
super().__init__(message, response)
|
104
107
|
|
105
108
|
|
106
109
|
class ServiceNotEnabledException(PangeaAPIException):
|
107
110
|
def __init__(self, service_name: str, response: PangeaResponse):
|
108
111
|
message = f"{service_name} is not enabled. Go to console.pangea.cloud/service/{service_name} to enable"
|
109
|
-
super(
|
112
|
+
super().__init__(message, response)
|
110
113
|
|
111
114
|
|
112
115
|
class MissingConfigID(PangeaAPIException):
|
113
116
|
"""No config ID was provided in either token scopes or explicitly"""
|
114
117
|
|
115
118
|
def __init__(self, service_name: str, response: PangeaResponse):
|
116
|
-
super(
|
119
|
+
super().__init__(
|
117
120
|
f"Token did not contain a config scope for service {service_name}. Create a new token or provide a config ID explicitly in the service base",
|
118
121
|
response,
|
119
122
|
)
|
pangea/file_uploader.py
CHANGED
pangea/request.py
CHANGED
@@ -1,19 +1,24 @@
|
|
1
1
|
# Copyright 2022 Pangea Cyber Corporation
|
2
2
|
# Author: Pangea Cyber Corporation
|
3
|
+
|
4
|
+
# TODO: Modernize.
|
5
|
+
# ruff: noqa: UP006, UP035
|
6
|
+
|
3
7
|
from __future__ import annotations
|
4
8
|
|
5
9
|
import copy
|
6
10
|
import json
|
7
11
|
import logging
|
8
12
|
import time
|
9
|
-
from
|
13
|
+
from collections.abc import Iterable, Mapping
|
14
|
+
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union, cast
|
10
15
|
|
11
16
|
import requests
|
12
17
|
from pydantic import BaseModel
|
13
18
|
from pydantic_core import to_jsonable_python
|
14
19
|
from requests.adapters import HTTPAdapter, Retry
|
15
20
|
from requests_toolbelt import MultipartDecoder # type: ignore[import-untyped]
|
16
|
-
from typing_extensions import TypeVar
|
21
|
+
from typing_extensions import TypeAlias, TypeVar, override
|
17
22
|
from yarl import URL
|
18
23
|
|
19
24
|
import pangea
|
@@ -26,11 +31,26 @@ if TYPE_CHECKING:
|
|
26
31
|
import aiohttp
|
27
32
|
|
28
33
|
|
34
|
+
_Data: TypeAlias = Union[Iterable[bytes], str, bytes, list[tuple[Any, Any]]]
|
35
|
+
|
36
|
+
_FileName: TypeAlias = Union[str, None]
|
37
|
+
_FileContent: TypeAlias = Union[str, bytes]
|
38
|
+
_FileContentType: TypeAlias = str
|
39
|
+
_FileCustomHeaders: TypeAlias = Mapping[str, str]
|
40
|
+
_FileSpecTuple2: TypeAlias = tuple[_FileName, _FileContent]
|
41
|
+
_FileSpecTuple3: TypeAlias = tuple[_FileName, _FileContent, _FileContentType]
|
42
|
+
_FileSpecTuple4: TypeAlias = tuple[_FileName, _FileContent, _FileContentType, _FileCustomHeaders]
|
43
|
+
_FileSpec: TypeAlias = Union[_FileContent, _FileSpecTuple2, _FileSpecTuple3, _FileSpecTuple4]
|
44
|
+
_Files: TypeAlias = Union[Mapping[str, _FileSpec], Iterable[tuple[str, _FileSpec]]]
|
45
|
+
|
46
|
+
_HeadersUpdateMapping: TypeAlias = Mapping[str, str]
|
47
|
+
|
48
|
+
|
29
49
|
class MultipartResponse:
|
30
50
|
pangea_json: Dict[str, str]
|
31
51
|
attached_files: List = []
|
32
52
|
|
33
|
-
def __init__(self, pangea_json:
|
53
|
+
def __init__(self, pangea_json: dict[str, str], attached_files: list = []): # noqa: B006
|
34
54
|
self.pangea_json = pangea_json
|
35
55
|
self.attached_files = attached_files
|
36
56
|
|
@@ -48,7 +68,7 @@ class PangeaRequestBase:
|
|
48
68
|
self._queued_retry_enabled = config.queued_retry_enabled
|
49
69
|
|
50
70
|
# Custom headers
|
51
|
-
self._extra_headers:
|
71
|
+
self._extra_headers: _HeadersUpdateMapping = {}
|
52
72
|
self._user_agent = ""
|
53
73
|
|
54
74
|
self.set_custom_user_agent(config.custom_user_agent)
|
@@ -63,18 +83,17 @@ class PangeaRequestBase:
|
|
63
83
|
|
64
84
|
return self._session
|
65
85
|
|
66
|
-
def set_extra_headers(self, headers:
|
86
|
+
def set_extra_headers(self, headers: _HeadersUpdateMapping):
|
67
87
|
"""Sets any additional headers in the request.
|
68
88
|
|
69
89
|
Args:
|
70
|
-
headers
|
90
|
+
headers: key-value pairs containing extra headers to set
|
71
91
|
|
72
92
|
Example:
|
73
93
|
set_extra_headers({ "My-Header" : "foobar" })
|
74
94
|
"""
|
75
95
|
|
76
|
-
|
77
|
-
self._extra_headers = headers
|
96
|
+
self._extra_headers = headers
|
78
97
|
|
79
98
|
def set_custom_user_agent(self, user_agent: Optional[str]):
|
80
99
|
self.config.custom_user_agent = user_agent
|
@@ -111,16 +130,14 @@ class PangeaRequestBase:
|
|
111
130
|
url = URL(self.config.base_url_template.format(SERVICE_NAME=self.service))
|
112
131
|
return str(url / path)
|
113
132
|
|
114
|
-
def _headers(self) -> dict:
|
115
|
-
|
116
|
-
|
133
|
+
def _headers(self) -> dict[str, str]:
|
134
|
+
return {
|
135
|
+
**self._extra_headers,
|
117
136
|
"Authorization": f"Bearer {self.token}",
|
137
|
+
"Content-Type": "application/json",
|
138
|
+
"User-Agent": self._user_agent,
|
118
139
|
}
|
119
140
|
|
120
|
-
# We want to ignore previous headers if user tried to set them, so we will overwrite them.
|
121
|
-
self._extra_headers.update(headers)
|
122
|
-
return self._extra_headers
|
123
|
-
|
124
141
|
def _get_filename_from_content_disposition(self, content_disposition: str) -> Optional[str]:
|
125
142
|
filename_parts = content_disposition.split("name=")
|
126
143
|
if len(filename_parts) > 1:
|
@@ -184,6 +201,9 @@ class PangeaRequestBase:
|
|
184
201
|
raise pe.AcceptedRequestException(response)
|
185
202
|
raise pe.PangeaAPIException(f"{summary} ", response)
|
186
203
|
|
204
|
+
def _init_session(self) -> requests.Session | aiohttp.ClientSession:
|
205
|
+
raise NotImplementedError
|
206
|
+
|
187
207
|
|
188
208
|
TResult = TypeVar("TResult", bound=PangeaResponseResult)
|
189
209
|
|
@@ -204,8 +224,8 @@ class PangeaRequest(PangeaRequestBase):
|
|
204
224
|
self,
|
205
225
|
endpoint: str,
|
206
226
|
result_class: Type[TResult],
|
207
|
-
data: str | BaseModel |
|
208
|
-
files: Optional[
|
227
|
+
data: str | BaseModel | Mapping[str, Any] | None = None,
|
228
|
+
files: Optional[list[Tuple]] = None,
|
209
229
|
poll_result: bool = True,
|
210
230
|
url: Optional[str] = None,
|
211
231
|
) -> PangeaResponse[TResult]:
|
@@ -250,9 +270,10 @@ class PangeaRequest(PangeaRequestBase):
|
|
250
270
|
endpoint, result_class=result_class, data=data, files=files
|
251
271
|
)
|
252
272
|
else:
|
253
|
-
|
254
|
-
|
255
|
-
|
273
|
+
headers = self._headers()
|
274
|
+
if transfer_method == TransferMethod.MULTIPART.value:
|
275
|
+
del headers["Content-Type"]
|
276
|
+
requests_response = self._http_post(url, headers=headers, data=data, files=files)
|
256
277
|
|
257
278
|
self._check_http_errors(requests_response)
|
258
279
|
|
@@ -273,7 +294,7 @@ class PangeaRequest(PangeaRequestBase):
|
|
273
294
|
|
274
295
|
pangea_response = PangeaResponse(requests_response, result_class=result_class, json=json_resp)
|
275
296
|
except requests.exceptions.JSONDecodeError as e:
|
276
|
-
raise pe.PangeaException(f"Failed to decode json response. {e}. Body: {requests_response.text}")
|
297
|
+
raise pe.PangeaException(f"Failed to decode json response. {e}. Body: {requests_response.text}") from e
|
277
298
|
|
278
299
|
if poll_result:
|
279
300
|
pangea_response = self._handle_queued_result(pangea_response)
|
@@ -324,9 +345,9 @@ class PangeaRequest(PangeaRequestBase):
|
|
324
345
|
def _http_post(
|
325
346
|
self,
|
326
347
|
url: str,
|
327
|
-
headers:
|
328
|
-
data:
|
329
|
-
files:
|
348
|
+
headers: _HeadersUpdateMapping = {}, # noqa: B006
|
349
|
+
data: str | dict[Any, Any] = {}, # noqa: B006
|
350
|
+
files: _Files | None = None,
|
330
351
|
multipart_post: bool = True,
|
331
352
|
) -> requests.Response:
|
332
353
|
data_send, files = self._http_post_process(data=data, files=files, multipart_post=multipart_post)
|
@@ -334,31 +355,25 @@ class PangeaRequest(PangeaRequestBase):
|
|
334
355
|
|
335
356
|
def _http_post_process(
|
336
357
|
self,
|
337
|
-
data:
|
338
|
-
files:
|
358
|
+
data: str | dict[Any, Any] = {}, # noqa: B006
|
359
|
+
files: _Files | None = None,
|
339
360
|
multipart_post: bool = True,
|
340
|
-
):
|
361
|
+
) -> tuple[_Data | None, _Files | None]:
|
341
362
|
if files:
|
342
363
|
if multipart_post is True:
|
343
364
|
data_send: str = json.dumps(data, default=default_encoder) if isinstance(data, dict) else data
|
344
|
-
multi = [("request", (None, data_send, "application/json"))]
|
345
|
-
multi
|
346
|
-
|
347
|
-
|
348
|
-
#
|
349
|
-
|
350
|
-
|
351
|
-
|
352
|
-
|
353
|
-
files = { # type: ignore[assignment]
|
354
|
-
"file": files[0][1],
|
355
|
-
}
|
356
|
-
return data_send, files
|
365
|
+
multi: list[tuple[str, _FileSpec]] = [("request", (None, data_send, "application/json")), *files]
|
366
|
+
return None, multi
|
367
|
+
|
368
|
+
# Post to presigned URL as form.
|
369
|
+
# When posting to presigned URL, file key should be 'file'.
|
370
|
+
assert isinstance(data, dict)
|
371
|
+
assert isinstance(files, list)
|
372
|
+
return [(k, v) for k, v in data.items()], {"file": files[0][1]}
|
373
|
+
|
357
374
|
data_send = json.dumps(data, default=default_encoder) if isinstance(data, dict) else data
|
358
375
|
return data_send, None
|
359
376
|
|
360
|
-
return data, files
|
361
|
-
|
362
377
|
def _handle_queued_result(self, response: PangeaResponse[TResult]) -> PangeaResponse[TResult]:
|
363
378
|
if self._queued_retry_enabled and response.http_status == 202:
|
364
379
|
self.logger.debug(
|
@@ -477,7 +492,7 @@ class PangeaRequest(PangeaRequestBase):
|
|
477
492
|
self,
|
478
493
|
endpoint: str,
|
479
494
|
result_class: Type[PangeaResponseResult],
|
480
|
-
data: Union[str,
|
495
|
+
data: Union[str, Mapping[str, Any]] = {},
|
481
496
|
) -> PangeaResponse:
|
482
497
|
# Send request
|
483
498
|
try:
|
@@ -520,8 +535,8 @@ class PangeaRequest(PangeaRequestBase):
|
|
520
535
|
def _http_put(
|
521
536
|
self,
|
522
537
|
url: str,
|
523
|
-
files:
|
524
|
-
headers:
|
538
|
+
files: list[Tuple],
|
539
|
+
headers: Mapping[str, str] = {},
|
525
540
|
) -> requests.Response:
|
526
541
|
self.logger.debug(
|
527
542
|
json.dumps({"service": self.service, "action": "http_put", "url": url}, default=default_encoder)
|
@@ -533,7 +548,7 @@ class PangeaRequest(PangeaRequestBase):
|
|
533
548
|
self,
|
534
549
|
endpoint: str,
|
535
550
|
result_class: Type[PangeaResponseResult],
|
536
|
-
data: Union[str,
|
551
|
+
data: Union[str, Mapping[str, Any]] = {},
|
537
552
|
files: Optional[List[Tuple]] = None,
|
538
553
|
):
|
539
554
|
if files is None or len(files) == 0:
|
@@ -603,7 +618,7 @@ class PangeaRequest(PangeaRequestBase):
|
|
603
618
|
{"service": self.service, "action": "poll_presigned_url", "step": "exit", "cause": {str(e)}}
|
604
619
|
)
|
605
620
|
)
|
606
|
-
raise pe.PresignedURLException("Failed to pull Presigned URL", loop_resp, e)
|
621
|
+
raise pe.PresignedURLException("Failed to pull Presigned URL", loop_resp, e) from e
|
607
622
|
|
608
623
|
self.logger.debug(json.dumps({"service": self.service, "action": "poll_presigned_url", "step": "exit"}))
|
609
624
|
|
@@ -611,6 +626,7 @@ class PangeaRequest(PangeaRequestBase):
|
|
611
626
|
return loop_resp
|
612
627
|
raise loop_exc
|
613
628
|
|
629
|
+
@override
|
614
630
|
def _init_session(self) -> requests.Session:
|
615
631
|
retry_config = Retry(
|
616
632
|
total=self.config.request_retries,
|
pangea/response.py
CHANGED
@@ -1,19 +1,25 @@
|
|
1
1
|
# Copyright 2022 Pangea Cyber Corporation
|
2
2
|
# Author: Pangea Cyber Corporation
|
3
|
+
|
4
|
+
# TODO: Modernize.
|
5
|
+
# ruff: noqa: UP006, UP035
|
6
|
+
|
7
|
+
from __future__ import annotations
|
8
|
+
|
3
9
|
import datetime
|
4
10
|
import enum
|
5
11
|
import os
|
6
|
-
from typing import Any, Dict, Generic, List, Optional, Type, Union
|
12
|
+
from typing import Annotated, Any, Dict, Generic, List, Optional, Type, Union
|
7
13
|
|
8
14
|
import aiohttp
|
9
15
|
import requests
|
10
16
|
from pydantic import BaseModel, ConfigDict, PlainSerializer
|
11
|
-
from typing_extensions import
|
17
|
+
from typing_extensions import TypeVar
|
12
18
|
|
13
19
|
from pangea.utils import format_datetime
|
14
20
|
|
15
21
|
|
16
|
-
class AttachedFile
|
22
|
+
class AttachedFile:
|
17
23
|
filename: str
|
18
24
|
file: bytes
|
19
25
|
content_type: str
|
@@ -40,10 +46,7 @@ class AttachedFile(object):
|
|
40
46
|
base_name, ext = os.path.splitext(file_path)
|
41
47
|
counter = 1
|
42
48
|
while os.path.exists(file_path):
|
43
|
-
if ext
|
44
|
-
file_path = f"{base_name}_{counter}{ext}"
|
45
|
-
else:
|
46
|
-
file_path = f"{base_name}_{counter}"
|
49
|
+
file_path = f"{base_name}_{counter}{ext}" if ext else f"{base_name}_{counter}"
|
47
50
|
counter += 1
|
48
51
|
return file_path
|
49
52
|
|
@@ -79,14 +82,12 @@ class TransferMethod(str, enum.Enum):
|
|
79
82
|
PangeaDateTime = Annotated[datetime.datetime, PlainSerializer(format_datetime)]
|
80
83
|
|
81
84
|
|
82
|
-
|
83
|
-
|
84
|
-
model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
|
85
|
+
class APIRequestModel(BaseModel):
|
86
|
+
model_config = ConfigDict(extra="forbid")
|
85
87
|
|
86
88
|
|
87
|
-
|
88
|
-
|
89
|
-
model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
|
89
|
+
class APIResponseModel(BaseModel):
|
90
|
+
model_config = ConfigDict(extra="allow")
|
90
91
|
|
91
92
|
|
92
93
|
class PangeaResponseResult(APIResponseModel):
|
@@ -192,6 +193,8 @@ T = TypeVar("T", bound=PangeaResponseResult)
|
|
192
193
|
|
193
194
|
|
194
195
|
class PangeaResponse(ResponseHeader, Generic[T]):
|
196
|
+
model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
|
197
|
+
|
195
198
|
raw_result: Optional[Dict[str, Any]] = None
|
196
199
|
raw_response: Optional[Union[requests.Response, aiohttp.ClientResponse]] = None
|
197
200
|
result: Optional[T] = None
|
@@ -199,16 +202,16 @@ class PangeaResponse(ResponseHeader, Generic[T]):
|
|
199
202
|
accepted_result: Optional[AcceptedResult] = None
|
200
203
|
result_class: Type[T] = PangeaResponseResult # type: ignore[assignment]
|
201
204
|
_json: Any
|
202
|
-
attached_files:
|
205
|
+
attached_files: list[AttachedFile] = []
|
203
206
|
|
204
207
|
def __init__(
|
205
208
|
self,
|
206
209
|
response: requests.Response,
|
207
210
|
result_class: Type[T],
|
208
211
|
json: dict,
|
209
|
-
attached_files:
|
212
|
+
attached_files: list[AttachedFile] = [], # noqa: B006
|
210
213
|
):
|
211
|
-
super(
|
214
|
+
super().__init__(**json)
|
212
215
|
self._json = json
|
213
216
|
self.raw_response = response
|
214
217
|
self.raw_result = self._json["result"]
|
@@ -241,10 +244,10 @@ class PangeaResponse(ResponseHeader, Generic[T]):
|
|
241
244
|
@property
|
242
245
|
def http_status(self) -> int: # type: ignore[return]
|
243
246
|
if self.raw_response:
|
244
|
-
if
|
247
|
+
if isinstance(self.raw_response, aiohttp.ClientResponse):
|
245
248
|
return self.raw_response.status
|
246
249
|
else:
|
247
|
-
return self.raw_response.status_code
|
250
|
+
return self.raw_response.status_code
|
248
251
|
|
249
252
|
@property
|
250
253
|
def url(self) -> str:
|