pangea-sdk 3.7.0__py3-none-any.whl → 3.8.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 +1 -1
- pangea/asyncio/request.py +123 -32
- pangea/asyncio/services/__init__.py +1 -0
- pangea/asyncio/services/audit.py +236 -21
- pangea/asyncio/services/authn.py +79 -50
- pangea/asyncio/services/authz.py +259 -0
- pangea/asyncio/services/base.py +9 -6
- pangea/asyncio/services/file_scan.py +3 -4
- pangea/asyncio/services/intel.py +5 -6
- pangea/asyncio/services/redact.py +21 -3
- pangea/asyncio/services/vault.py +28 -12
- pangea/config.py +10 -18
- pangea/dump_audit.py +1 -0
- pangea/exceptions.py +8 -0
- pangea/request.py +164 -74
- pangea/response.py +63 -17
- pangea/services/__init__.py +1 -0
- pangea/services/audit/audit.py +241 -55
- pangea/services/audit/exceptions.py +1 -2
- pangea/services/audit/models.py +83 -21
- pangea/services/audit/signing.py +1 -0
- pangea/services/audit/util.py +1 -0
- pangea/services/authn/authn.py +38 -4
- pangea/services/authn/models.py +9 -9
- pangea/services/authz.py +377 -0
- pangea/services/base.py +34 -14
- pangea/services/embargo.py +1 -2
- pangea/services/file_scan.py +3 -4
- pangea/services/intel.py +3 -4
- pangea/services/redact.py +21 -3
- pangea/services/vault/vault.py +29 -12
- pangea/utils.py +2 -3
- {pangea_sdk-3.7.0.dist-info → pangea_sdk-3.8.0.dist-info}/METADATA +33 -6
- pangea_sdk-3.8.0.dist-info/RECORD +46 -0
- pangea_sdk-3.7.0.dist-info/RECORD +0 -44
- {pangea_sdk-3.7.0.dist-info → pangea_sdk-3.8.0.dist-info}/WHEEL +0 -0
pangea/asyncio/services/vault.py
CHANGED
@@ -1,8 +1,12 @@
|
|
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
|
from typing import Dict, Optional, Union
|
5
7
|
|
8
|
+
from pangea.asyncio.services.base import ServiceBaseAsync
|
9
|
+
from pangea.config import PangeaConfig
|
6
10
|
from pangea.response import PangeaResponse
|
7
11
|
from pangea.services.vault.models.asymmetric import (
|
8
12
|
AsymmetricGenerateRequest,
|
@@ -47,8 +51,8 @@ from pangea.services.vault.models.common import (
|
|
47
51
|
StateChangeRequest,
|
48
52
|
StateChangeResult,
|
49
53
|
SymmetricAlgorithm,
|
50
|
-
TDict,
|
51
54
|
Tags,
|
55
|
+
TDict,
|
52
56
|
UpdateRequest,
|
53
57
|
UpdateResult,
|
54
58
|
)
|
@@ -69,8 +73,6 @@ from pangea.services.vault.models.symmetric import (
|
|
69
73
|
SymmetricStoreResult,
|
70
74
|
)
|
71
75
|
|
72
|
-
from .base import ServiceBaseAsync
|
73
|
-
|
74
76
|
|
75
77
|
class VaultAsync(ServiceBaseAsync):
|
76
78
|
"""Vault service client.
|
@@ -92,17 +94,31 @@ class VaultAsync(ServiceBaseAsync):
|
|
92
94
|
vault_config = PangeaConfig(domain="pangea.cloud")
|
93
95
|
|
94
96
|
# Setup Pangea Vault service
|
95
|
-
vault = Vault(token=PANGEA_VAULT_TOKEN, config=
|
97
|
+
vault = Vault(token=PANGEA_VAULT_TOKEN, config=vault_config)
|
96
98
|
"""
|
97
99
|
|
98
100
|
service_name = "vault"
|
99
101
|
|
100
102
|
def __init__(
|
101
103
|
self,
|
102
|
-
token,
|
103
|
-
config=None,
|
104
|
-
logger_name="pangea",
|
105
|
-
):
|
104
|
+
token: str,
|
105
|
+
config: PangeaConfig | None = None,
|
106
|
+
logger_name: str = "pangea",
|
107
|
+
) -> None:
|
108
|
+
"""
|
109
|
+
Vault client
|
110
|
+
|
111
|
+
Initializes a new Vault client.
|
112
|
+
|
113
|
+
Args:
|
114
|
+
token: Pangea API token.
|
115
|
+
config: Configuration.
|
116
|
+
logger_name: Logger name.
|
117
|
+
|
118
|
+
Examples:
|
119
|
+
config = PangeaConfig(domain="pangea_domain")
|
120
|
+
vault = VaultAsync(token="pangea_token", config=config)
|
121
|
+
"""
|
106
122
|
super().__init__(token, config, logger_name)
|
107
123
|
|
108
124
|
# Delete endpoint
|
@@ -859,8 +875,8 @@ class VaultAsync(ServiceBaseAsync):
|
|
859
875
|
|
860
876
|
Default is `deactivated`.
|
861
877
|
public_key (EncodedPublicKey, optional): The public key (in PEM format)
|
862
|
-
private_key
|
863
|
-
key
|
878
|
+
private_key (EncodedPrivateKey, optional): The private key (in PEM format)
|
879
|
+
key (EncodedSymmetricKey, optional): The key material (in base64)
|
864
880
|
|
865
881
|
Raises:
|
866
882
|
PangeaAPIException: If an API Error happens
|
@@ -1222,7 +1238,7 @@ class VaultAsync(ServiceBaseAsync):
|
|
1222
1238
|
)
|
1223
1239
|
"""
|
1224
1240
|
|
1225
|
-
input = EncryptStructuredRequest(
|
1241
|
+
input: EncryptStructuredRequest[TDict] = EncryptStructuredRequest(
|
1226
1242
|
id=id, structured_data=structured_data, filter=filter, version=version, additional_data=additional_data
|
1227
1243
|
)
|
1228
1244
|
return await self.request.post(
|
@@ -1271,7 +1287,7 @@ class VaultAsync(ServiceBaseAsync):
|
|
1271
1287
|
)
|
1272
1288
|
"""
|
1273
1289
|
|
1274
|
-
input = EncryptStructuredRequest(
|
1290
|
+
input: EncryptStructuredRequest[TDict] = EncryptStructuredRequest(
|
1275
1291
|
id=id, structured_data=structured_data, filter=filter, version=version, additional_data=additional_data
|
1276
1292
|
)
|
1277
1293
|
return await self.request.post(
|
pangea/config.py
CHANGED
@@ -9,65 +9,57 @@ from typing import Optional
|
|
9
9
|
class PangeaConfig:
|
10
10
|
"""Holds run time configuration information used by SDK components."""
|
11
11
|
|
12
|
+
domain: str = "aws.us.pangea.cloud"
|
12
13
|
"""
|
13
14
|
Used to set Pangea domain (and port if needed), it should not include service subdomain
|
14
15
|
just for particular use cases when environment = "local", domain could be set to an url including:
|
15
16
|
scheme (http:// or https://), subdomain, domain and port.
|
16
|
-
|
17
17
|
"""
|
18
|
-
domain: str = "aws.us.pangea.cloud"
|
19
18
|
|
19
|
+
environment: str = "production"
|
20
20
|
"""
|
21
21
|
Used to generate service url.
|
22
22
|
It should be only 'production' or 'local' in cases of particular services that can run locally as Redact.
|
23
|
-
|
24
23
|
"""
|
25
|
-
environment: str = "production"
|
26
24
|
|
25
|
+
config_id: Optional[str] = None
|
27
26
|
"""
|
28
27
|
Only used for services that support multiconfig (e.g.: Audit service)
|
29
28
|
|
30
29
|
@deprecated("config_id will be deprecated from PangeaConfig. Set it on service initialization instead")
|
31
30
|
"""
|
32
|
-
config_id: Optional[str] = None
|
33
31
|
|
32
|
+
insecure: bool = False
|
34
33
|
"""
|
35
34
|
Set to true to use plain http
|
36
|
-
|
37
35
|
"""
|
38
|
-
insecure: bool = False
|
39
36
|
|
37
|
+
request_retries: int = 3
|
40
38
|
"""
|
41
39
|
Number of retries on the initial request
|
42
|
-
|
43
40
|
"""
|
44
|
-
request_retries: int = 3
|
45
41
|
|
42
|
+
request_backoff: float = 0.5
|
46
43
|
"""
|
47
44
|
Backoff strategy passed to 'requests'
|
48
|
-
|
49
45
|
"""
|
50
|
-
request_backoff: float = 0.5
|
51
46
|
|
47
|
+
request_timeout: int = 5
|
52
48
|
"""
|
53
49
|
Timeout used on initial request attempts
|
54
|
-
|
55
50
|
"""
|
56
|
-
request_timeout: int = 5
|
57
51
|
|
52
|
+
poll_result_timeout: int = 30
|
58
53
|
"""
|
59
54
|
Timeout used to poll results after 202 (in secs)
|
60
|
-
|
61
55
|
"""
|
62
|
-
poll_result_timeout: int = 30
|
63
56
|
|
57
|
+
queued_retry_enabled: bool = True
|
64
58
|
"""
|
65
59
|
Enable queued request retry support
|
66
60
|
"""
|
67
|
-
queued_retry_enabled: bool = True
|
68
61
|
|
62
|
+
custom_user_agent: Optional[str] = None
|
69
63
|
"""
|
70
64
|
Extra user agent to be added to request user agent
|
71
|
-
|
72
65
|
"""
|
73
|
-
custom_user_agent: Optional[str] = None
|
pangea/dump_audit.py
CHANGED
@@ -10,6 +10,7 @@ from datetime import datetime
|
|
10
10
|
from typing import Tuple
|
11
11
|
|
12
12
|
import dateutil.parser
|
13
|
+
|
13
14
|
from pangea.response import PangeaResponse
|
14
15
|
from pangea.services import Audit
|
15
16
|
from pangea.services.audit.models import SearchEvent, SearchOrder, SearchOrderBy, SearchOutput, SearchResultOutput
|
pangea/exceptions.py
CHANGED
@@ -30,6 +30,14 @@ class PresignedUploadError(PangeaException):
|
|
30
30
|
self.body = body
|
31
31
|
|
32
32
|
|
33
|
+
class DownloadFileError(PangeaException):
|
34
|
+
body: str
|
35
|
+
|
36
|
+
def __init__(self, message: str, body: str):
|
37
|
+
super().__init__(message)
|
38
|
+
self.body = body
|
39
|
+
|
40
|
+
|
33
41
|
class PangeaAPIException(PangeaException):
|
34
42
|
"""Exceptions raised during API calls"""
|
35
43
|
|
pangea/request.py
CHANGED
@@ -5,16 +5,30 @@ import copy
|
|
5
5
|
import json
|
6
6
|
import logging
|
7
7
|
import time
|
8
|
-
from typing import Dict, List, Optional, Tuple, Type, Union
|
8
|
+
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type, Union
|
9
|
+
|
10
|
+
import requests
|
11
|
+
from requests.adapters import HTTPAdapter, Retry
|
12
|
+
from requests_toolbelt import MultipartDecoder # type: ignore
|
13
|
+
from typing_extensions import TypeVar
|
9
14
|
|
10
|
-
import aiohttp
|
11
15
|
import pangea
|
12
16
|
import pangea.exceptions as pe
|
13
|
-
import requests
|
14
17
|
from pangea.config import PangeaConfig
|
15
|
-
from pangea.response import PangeaResponse, PangeaResponseResult, ResponseStatus, TransferMethod
|
18
|
+
from pangea.response import AttachedFile, PangeaResponse, PangeaResponseResult, ResponseStatus, TransferMethod
|
16
19
|
from pangea.utils import default_encoder
|
17
|
-
|
20
|
+
|
21
|
+
if TYPE_CHECKING:
|
22
|
+
import aiohttp
|
23
|
+
|
24
|
+
|
25
|
+
class MultipartResponse(object):
|
26
|
+
pangea_json: Dict[str, str]
|
27
|
+
attached_files: List = []
|
28
|
+
|
29
|
+
def __init__(self, pangea_json: Dict[str, str], attached_files: List = []):
|
30
|
+
self.pangea_json = pangea_json
|
31
|
+
self.attached_files = attached_files
|
18
32
|
|
19
33
|
|
20
34
|
class PangeaRequestBase(object):
|
@@ -30,9 +44,10 @@ class PangeaRequestBase(object):
|
|
30
44
|
self._queued_retry_enabled = config.queued_retry_enabled
|
31
45
|
|
32
46
|
# Custom headers
|
33
|
-
self._extra_headers = {}
|
47
|
+
self._extra_headers: Dict = {}
|
34
48
|
self._user_agent = ""
|
35
|
-
|
49
|
+
|
50
|
+
self.set_custom_user_agent(config.custom_user_agent)
|
36
51
|
self._session: Optional[Union[requests.Session, aiohttp.ClientSession]] = None
|
37
52
|
|
38
53
|
self.logger = logger
|
@@ -57,7 +72,7 @@ class PangeaRequestBase(object):
|
|
57
72
|
if isinstance(headers, dict):
|
58
73
|
self._extra_headers = headers
|
59
74
|
|
60
|
-
def set_custom_user_agent(self, user_agent: str):
|
75
|
+
def set_custom_user_agent(self, user_agent: Optional[str]):
|
61
76
|
self.config.custom_user_agent = user_agent
|
62
77
|
self._user_agent = f"pangea-python/{pangea.__version__}"
|
63
78
|
if self.config.custom_user_agent:
|
@@ -110,6 +125,16 @@ class PangeaRequestBase(object):
|
|
110
125
|
self._extra_headers.update(headers)
|
111
126
|
return self._extra_headers
|
112
127
|
|
128
|
+
def _get_filename_from_content_disposition(self, content_disposition: str) -> Optional[str]:
|
129
|
+
filename_parts = content_disposition.split("name=")
|
130
|
+
if len(filename_parts) > 1:
|
131
|
+
return filename_parts[1].split(";")[0].strip('"')
|
132
|
+
else:
|
133
|
+
return None
|
134
|
+
|
135
|
+
def _get_filename_from_url(self, url: str) -> Optional[str]:
|
136
|
+
return url.split("/")[-1].split("?")[0]
|
137
|
+
|
113
138
|
def _check_response(self, response: PangeaResponse) -> PangeaResponse:
|
114
139
|
status = response.status
|
115
140
|
summary = response.summary
|
@@ -149,7 +174,7 @@ class PangeaRequestBase(object):
|
|
149
174
|
elif status == ResponseStatus.TREE_NOT_FOUND.value:
|
150
175
|
raise pe.TreeNotFoundException(summary, response)
|
151
176
|
elif status == ResponseStatus.IP_NOT_FOUND.value:
|
152
|
-
raise pe.IPNotFoundException(summary)
|
177
|
+
raise pe.IPNotFoundException(summary, response)
|
153
178
|
elif status == ResponseStatus.BAD_OFFSET.value:
|
154
179
|
raise pe.BadOffsetException(summary, response)
|
155
180
|
elif status == ResponseStatus.FORBIDDEN_VAULT_OPERATION.value:
|
@@ -157,7 +182,7 @@ class PangeaRequestBase(object):
|
|
157
182
|
elif status == ResponseStatus.VAULT_ITEM_NOT_FOUND.value:
|
158
183
|
raise pe.VaultItemNotFound(summary, response)
|
159
184
|
elif status == ResponseStatus.NOT_FOUND.value:
|
160
|
-
raise pe.NotFound(response.raw_response.url if response.raw_response is not None else "", response) # type: ignore[arg-type]
|
185
|
+
raise pe.NotFound(str(response.raw_response.url) if response.raw_response is not None else "", response) # type: ignore[arg-type]
|
161
186
|
elif status == ResponseStatus.INTERNAL_SERVER_ERROR.value:
|
162
187
|
raise pe.InternalServerError(response)
|
163
188
|
elif status == ResponseStatus.ACCEPTED.value:
|
@@ -165,6 +190,9 @@ class PangeaRequestBase(object):
|
|
165
190
|
raise pe.PangeaAPIException(f"{summary} ", response)
|
166
191
|
|
167
192
|
|
193
|
+
TResult = TypeVar("TResult", bound=PangeaResponseResult, default=PangeaResponseResult)
|
194
|
+
|
195
|
+
|
168
196
|
class PangeaRequest(PangeaRequestBase):
|
169
197
|
"""An object that makes direct calls to Pangea Service APIs.
|
170
198
|
|
@@ -180,12 +208,12 @@ class PangeaRequest(PangeaRequestBase):
|
|
180
208
|
def post(
|
181
209
|
self,
|
182
210
|
endpoint: str,
|
183
|
-
result_class: Type[
|
211
|
+
result_class: Type[TResult],
|
184
212
|
data: Union[str, Dict] = {},
|
185
213
|
files: Optional[List[Tuple]] = None,
|
186
214
|
poll_result: bool = True,
|
187
215
|
url: Optional[str] = None,
|
188
|
-
) -> PangeaResponse:
|
216
|
+
) -> PangeaResponse[TResult]:
|
189
217
|
"""Makes the POST call to a Pangea Service endpoint.
|
190
218
|
|
191
219
|
Args:
|
@@ -200,13 +228,13 @@ class PangeaRequest(PangeaRequestBase):
|
|
200
228
|
url = self._url(endpoint)
|
201
229
|
|
202
230
|
# Set config ID if available
|
203
|
-
if self.config_id and data.get("config_id", None) is None:
|
204
|
-
data["config_id"] = self.config_id
|
231
|
+
if self.config_id and isinstance(data, dict) and data.get("config_id", None) is None:
|
232
|
+
data["config_id"] = self.config_id
|
205
233
|
|
206
234
|
self.logger.debug(
|
207
235
|
json.dumps({"service": self.service, "action": "post", "url": url, "data": data}, default=default_encoder)
|
208
236
|
)
|
209
|
-
transfer_method = data.get("transfer_method", None)
|
237
|
+
transfer_method = data.get("transfer_method", None) if isinstance(data, dict) else None
|
210
238
|
|
211
239
|
if files is not None and type(data) is dict and (transfer_method == TransferMethod.POST_URL.value):
|
212
240
|
requests_response = self._full_post_presigned_url(
|
@@ -218,15 +246,68 @@ class PangeaRequest(PangeaRequestBase):
|
|
218
246
|
)
|
219
247
|
|
220
248
|
self._check_http_errors(requests_response)
|
221
|
-
json_resp = requests_response.json()
|
222
|
-
self.logger.debug(json.dumps({"service": self.service, "action": "post", "url": url, "response": json_resp}))
|
223
249
|
|
224
|
-
|
250
|
+
if "multipart/form-data" in requests_response.headers.get("content-type", ""):
|
251
|
+
multipart_response = self._process_multipart_response(requests_response)
|
252
|
+
pangea_response: PangeaResponse = PangeaResponse(
|
253
|
+
requests_response,
|
254
|
+
result_class=result_class,
|
255
|
+
json=multipart_response.pangea_json,
|
256
|
+
attached_files=multipart_response.attached_files,
|
257
|
+
)
|
258
|
+
else:
|
259
|
+
try:
|
260
|
+
json_resp = requests_response.json()
|
261
|
+
self.logger.debug(
|
262
|
+
json.dumps({"service": self.service, "action": "post", "url": url, "response": json_resp})
|
263
|
+
)
|
264
|
+
|
265
|
+
pangea_response = PangeaResponse(requests_response, result_class=result_class, json=json_resp)
|
266
|
+
except requests.exceptions.JSONDecodeError as e:
|
267
|
+
raise pe.PangeaException(f"Failed to decode json response. {e}. Body: {requests_response.text}")
|
268
|
+
|
225
269
|
if poll_result:
|
226
270
|
pangea_response = self._handle_queued_result(pangea_response)
|
227
271
|
|
228
272
|
return self._check_response(pangea_response)
|
229
273
|
|
274
|
+
def _get_pangea_json(self, decoder: MultipartDecoder) -> Optional[Dict]:
|
275
|
+
# Iterate through parts
|
276
|
+
for i, part in enumerate(decoder.parts):
|
277
|
+
if i == 0:
|
278
|
+
json_str = part.content.decode("utf-8")
|
279
|
+
return json.loads(json_str)
|
280
|
+
|
281
|
+
return None
|
282
|
+
|
283
|
+
def _get_attached_files(self, decoder: MultipartDecoder) -> List[AttachedFile]:
|
284
|
+
files = []
|
285
|
+
|
286
|
+
for i, part in enumerate(decoder.parts):
|
287
|
+
content_type = part.headers.get(b"Content-Type", b"").decode("utf-8")
|
288
|
+
# if "application/octet-stream" in content_type:
|
289
|
+
if i > 0:
|
290
|
+
content_disposition = part.headers.get(b"Content-Disposition", b"").decode("utf-8")
|
291
|
+
name = self._get_filename_from_content_disposition(content_disposition)
|
292
|
+
if name is None:
|
293
|
+
name = f"default_file_name_{i}"
|
294
|
+
|
295
|
+
files.append(AttachedFile(name, part.content, content_type))
|
296
|
+
|
297
|
+
return files
|
298
|
+
|
299
|
+
def _process_multipart_response(self, resp: requests.Response) -> MultipartResponse:
|
300
|
+
# Parse the multipart response
|
301
|
+
decoder = MultipartDecoder.from_response(resp)
|
302
|
+
|
303
|
+
pangea_json = self._get_pangea_json(decoder)
|
304
|
+
self.logger.debug(
|
305
|
+
json.dumps({"service": self.service, "action": "multipart response", "response": pangea_json})
|
306
|
+
)
|
307
|
+
|
308
|
+
attached_files = self._get_attached_files(decoder)
|
309
|
+
return MultipartResponse(pangea_json, attached_files) # type: ignore
|
310
|
+
|
230
311
|
def _check_http_errors(self, resp: requests.Response):
|
231
312
|
if resp.status_code == 503:
|
232
313
|
raise pe.ServiceTemporarilyUnavailable(resp.json())
|
@@ -268,48 +349,8 @@ class PangeaRequest(PangeaRequestBase):
|
|
268
349
|
|
269
350
|
return data, files
|
270
351
|
|
271
|
-
def
|
272
|
-
self
|
273
|
-
endpoint: str,
|
274
|
-
result_class: Type[PangeaResponseResult],
|
275
|
-
data: Union[str, Dict] = {},
|
276
|
-
files: Optional[List[Tuple]] = None,
|
277
|
-
):
|
278
|
-
if len(files) == 0: # type: ignore[arg-type]
|
279
|
-
raise AttributeError("files attribute should have at least 1 file")
|
280
|
-
|
281
|
-
# Send request
|
282
|
-
try:
|
283
|
-
# This should return 202 (AcceptedRequestException)
|
284
|
-
resp = self.post(endpoint=endpoint, result_class=result_class, data=data, poll_result=False)
|
285
|
-
raise pe.PresignedURLException("Should return 202", resp)
|
286
|
-
|
287
|
-
except pe.AcceptedRequestException as e:
|
288
|
-
accepted_exception = e
|
289
|
-
except Exception as e:
|
290
|
-
raise e
|
291
|
-
|
292
|
-
# Receive 202 with accepted_status
|
293
|
-
result = self._poll_presigned_url(accepted_exception) # type: ignore[arg-type]
|
294
|
-
data_to_presigned = result.accepted_status.upload_details # type: ignore[attr-defined]
|
295
|
-
presigned_url = result.accepted_status.upload_url # type: ignore[attr-defined]
|
296
|
-
|
297
|
-
# Send multipart request with file and upload_details as body
|
298
|
-
resp = self._http_post(url=presigned_url, data=data_to_presigned, files=files, multipart_post=False) # type: ignore[assignment]
|
299
|
-
self.logger.debug(
|
300
|
-
json.dumps(
|
301
|
-
{"service": self.service, "action": "post presigned", "url": presigned_url, "response": resp.text}, # type: ignore[attr-defined]
|
302
|
-
default=default_encoder,
|
303
|
-
)
|
304
|
-
)
|
305
|
-
|
306
|
-
if resp.status_code < 200 or resp.status_code >= 300: # type: ignore[attr-defined]
|
307
|
-
raise pe.PresignedUploadError(f"presigned POST failure: {resp.status_code}", resp.text) # type: ignore[attr-defined]
|
308
|
-
|
309
|
-
return accepted_exception.response.raw_response
|
310
|
-
|
311
|
-
def _handle_queued_result(self, response: PangeaResponse) -> PangeaResponse:
|
312
|
-
if self._queued_retry_enabled and response.raw_response.status_code == 202: # type: ignore[union-attr]
|
352
|
+
def _handle_queued_result(self, response: PangeaResponse[TResult]) -> PangeaResponse[TResult]:
|
353
|
+
if self._queued_retry_enabled and response.http_status == 202:
|
313
354
|
self.logger.debug(
|
314
355
|
json.dumps(
|
315
356
|
{"service": self.service, "action": "poll_result", "response": response.json},
|
@@ -320,7 +361,7 @@ class PangeaRequest(PangeaRequestBase):
|
|
320
361
|
|
321
362
|
return response
|
322
363
|
|
323
|
-
def get(self, path: str, result_class: Type[
|
364
|
+
def get(self, path: str, result_class: Type[TResult], check_response: bool = True) -> PangeaResponse[TResult]:
|
324
365
|
"""Makes the GET call to a Pangea Service endpoint.
|
325
366
|
|
326
367
|
Args:
|
@@ -336,7 +377,9 @@ class PangeaRequest(PangeaRequestBase):
|
|
336
377
|
self.logger.debug(json.dumps({"service": self.service, "action": "get", "url": url}))
|
337
378
|
requests_response = self.session.get(url, headers=self._headers())
|
338
379
|
self._check_http_errors(requests_response)
|
339
|
-
pangea_response
|
380
|
+
pangea_response: PangeaResponse = PangeaResponse(
|
381
|
+
requests_response, result_class=result_class, json=requests_response.json()
|
382
|
+
)
|
340
383
|
|
341
384
|
self.logger.debug(
|
342
385
|
json.dumps(
|
@@ -350,20 +393,61 @@ class PangeaRequest(PangeaRequestBase):
|
|
350
393
|
|
351
394
|
return self._check_response(pangea_response)
|
352
395
|
|
396
|
+
def download_file(self, url: str, filename: Optional[str] = None) -> AttachedFile:
|
397
|
+
self.logger.debug(
|
398
|
+
json.dumps(
|
399
|
+
{
|
400
|
+
"service": self.service,
|
401
|
+
"action": "download_file",
|
402
|
+
"url": url,
|
403
|
+
"filename": filename,
|
404
|
+
"status": "start",
|
405
|
+
}
|
406
|
+
)
|
407
|
+
)
|
408
|
+
response = self.session.get(url, headers={})
|
409
|
+
if response.status_code == 200:
|
410
|
+
if filename is None:
|
411
|
+
content_disposition = response.headers.get(b"Content-Disposition", b"").decode("utf-8")
|
412
|
+
filename = self._get_filename_from_content_disposition(content_disposition)
|
413
|
+
if filename is None:
|
414
|
+
filename = self._get_filename_from_url(url)
|
415
|
+
if filename is None:
|
416
|
+
filename = "default_filename"
|
417
|
+
|
418
|
+
content_type = response.headers.get(b"Content-Type", b"").decode("utf-8")
|
419
|
+
|
420
|
+
self.logger.debug(
|
421
|
+
json.dumps(
|
422
|
+
{
|
423
|
+
"service": self.service,
|
424
|
+
"action": "download_file",
|
425
|
+
"url": url,
|
426
|
+
"filename": filename,
|
427
|
+
"status": "success",
|
428
|
+
}
|
429
|
+
)
|
430
|
+
)
|
431
|
+
return AttachedFile(filename=filename, file=response.content, content_type=content_type)
|
432
|
+
else:
|
433
|
+
raise pe.DownloadFileError(f"Failed to download file. Status: {response.status_code}", response.text)
|
434
|
+
|
353
435
|
def poll_result_by_id(
|
354
|
-
self, request_id: str, result_class:
|
355
|
-
):
|
436
|
+
self, request_id: str, result_class: Type[TResult], check_response: bool = True
|
437
|
+
) -> PangeaResponse[TResult]:
|
356
438
|
path = self._get_poll_path(request_id)
|
357
439
|
self.logger.debug(json.dumps({"service": self.service, "action": "poll_result_once", "url": path}))
|
358
|
-
return self.get(path, result_class, check_response=check_response)
|
440
|
+
return self.get(path, result_class, check_response=check_response)
|
359
441
|
|
360
|
-
def poll_result_once(
|
442
|
+
def poll_result_once(
|
443
|
+
self, response: PangeaResponse[TResult], check_response: bool = True
|
444
|
+
) -> PangeaResponse[TResult]:
|
361
445
|
request_id = response.request_id
|
362
446
|
if not request_id:
|
363
447
|
raise pe.PangeaException("Poll result error: response did not include a 'request_id'")
|
364
448
|
|
365
449
|
if response.status != ResponseStatus.ACCEPTED.value:
|
366
|
-
raise pe.PangeaException("Response already
|
450
|
+
raise pe.PangeaException("Response already processed")
|
367
451
|
|
368
452
|
return self.poll_result_by_id(request_id, response.result_class, check_response=check_response)
|
369
453
|
|
@@ -422,7 +506,8 @@ class PangeaRequest(PangeaRequestBase):
|
|
422
506
|
self.logger.debug(
|
423
507
|
json.dumps({"service": self.service, "action": "http_put", "url": url}, default=default_encoder)
|
424
508
|
)
|
425
|
-
|
509
|
+
_, value = files[0]
|
510
|
+
return self.session.put(url, headers=headers, data=value[1])
|
426
511
|
|
427
512
|
def _full_post_presigned_url(
|
428
513
|
self,
|
@@ -431,17 +516,22 @@ class PangeaRequest(PangeaRequestBase):
|
|
431
516
|
data: Union[str, Dict] = {},
|
432
517
|
files: Optional[List[Tuple]] = None,
|
433
518
|
):
|
434
|
-
if len(files) == 0:
|
519
|
+
if files is None or len(files) == 0:
|
435
520
|
raise AttributeError("files attribute should have at least 1 file")
|
436
521
|
|
437
522
|
response = self.request_presigned_url(endpoint=endpoint, result_class=result_class, data=data)
|
438
|
-
|
439
|
-
|
523
|
+
if response.accepted_result is None:
|
524
|
+
raise pe.PangeaException("No accepted_result field when requesting presigned url")
|
525
|
+
if response.accepted_result.post_url is None:
|
526
|
+
raise pe.PresignedURLException("No presigned url", response)
|
527
|
+
|
528
|
+
data_to_presigned = response.accepted_result.post_form_data
|
529
|
+
presigned_url = response.accepted_result.post_url
|
440
530
|
|
441
|
-
self.post_presigned_url(url=presigned_url, data=data_to_presigned, files=files)
|
531
|
+
self.post_presigned_url(url=presigned_url, data=data_to_presigned, files=files)
|
442
532
|
return response.raw_response
|
443
533
|
|
444
|
-
def _poll_result_retry(self, response: PangeaResponse) -> PangeaResponse:
|
534
|
+
def _poll_result_retry(self, response: PangeaResponse[TResult]) -> PangeaResponse[TResult]:
|
445
535
|
retry_count = 1
|
446
536
|
start = time.time()
|
447
537
|
|
@@ -453,7 +543,7 @@ class PangeaRequest(PangeaRequestBase):
|
|
453
543
|
self.logger.debug(json.dumps({"service": self.service, "action": "poll_result_retry", "step": "exit"}))
|
454
544
|
return self._check_response(response)
|
455
545
|
|
456
|
-
def _poll_presigned_url(self, response: PangeaResponse) -> PangeaResponse:
|
546
|
+
def _poll_presigned_url(self, response: PangeaResponse[TResult]) -> PangeaResponse[TResult]:
|
457
547
|
if response.http_status != 202:
|
458
548
|
raise AttributeError("Response should be 202")
|
459
549
|
|