pangea-sdk 3.7.0__py3-none-any.whl → 3.7.1__py3-none-any.whl

pangea/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "3.7.0"
+__version__ = "3.7.1"
 
 from pangea.asyncio.request import PangeaRequestAsync
 from pangea.config import PangeaConfig

pangea/asyncio/request.py

@@ -7,12 +7,11 @@ import time
 from typing import Dict, List, Optional, Tuple, Type, Union
 
 import aiohttp
-import pangea.exceptions as pe
 from aiohttp import FormData
 
-# from requests.adapters import HTTPAdapter, Retry
-from pangea.request import PangeaRequestBase
-from pangea.response import AcceptedResult, PangeaResponse, PangeaResponseResult, ResponseStatus, TransferMethod
+import pangea.exceptions as pe
+from pangea.request import MultipartResponse, PangeaRequestBase
+from pangea.response import AttachedFile, PangeaResponse, PangeaResponseResult, ResponseStatus, TransferMethod
 from pangea.utils import default_encoder
 
 
@@ -30,7 +29,7 @@ class PangeaRequestAsync(PangeaRequestBase):
         endpoint: str,
         result_class: Type[PangeaResponseResult],
         data: Union[str, Dict] = {},
-        files: Optional[List[Tuple]] = None,
+        files: List[Tuple] = [],
         poll_result: bool = True,
         url: Optional[str] = None,
     ) -> PangeaResponse:
@@ -56,7 +55,7 @@ class PangeaRequestAsync(PangeaRequestBase):
         )
         transfer_method = data.get("transfer_method", None)  # type: ignore[union-attr]
 
-        if files is not None and type(data) is dict and (transfer_method == TransferMethod.POST_URL.value):
+        if files and type(data) is dict and (transfer_method == TransferMethod.POST_URL.value):
             requests_response = await self._full_post_presigned_url(
                 endpoint, result_class=result_class, data=data, files=files
             )
@@ -66,10 +65,26 @@ class PangeaRequestAsync(PangeaRequestBase):
             )
 
         await self._check_http_errors(requests_response)
-        json_resp = await requests_response.json()
-        self.logger.debug(json.dumps({"service": self.service, "action": "post", "url": url, "response": json_resp}))
 
-        pangea_response = PangeaResponse(requests_response, result_class=result_class, json=json_resp)  # type: ignore[var-annotated]
+        if "multipart/form-data" in requests_response.headers.get("content-type", ""):
+            multipart_response = await self._process_multipart_response(requests_response)
+            pangea_response: PangeaResponse = PangeaResponse(
+                requests_response,
+                result_class=result_class,
+                json=multipart_response.pangea_json,
+                attached_files=multipart_response.attached_files,
+            )
+        else:
+            try:
+                json_resp = await requests_response.json()
+                self.logger.debug(
+                    json.dumps({"service": self.service, "action": "post", "url": url, "response": json_resp})
+                )
+
+                pangea_response = PangeaResponse(requests_response, result_class=result_class, json=json_resp)
+            except aiohttp.ContentTypeError as e:
+                raise pe.PangeaException(f"Failed to decode json response. {e}. Body: {await requests_response.text()}")
+
         if poll_result:
             pangea_response = await self._handle_queued_result(pangea_response)
 
@@ -77,7 +92,7 @@ class PangeaRequestAsync(PangeaRequestBase):
 
     async def get(
         self, path: str, result_class: Type[PangeaResponseResult], check_response: bool = True
-    ) -> PangeaResponse:
+    ) -> PangeaResponse[Type[PangeaResponseResult]]:
         """Makes the GET call to a Pangea Service endpoint.
 
         Args:
@@ -115,7 +130,7 @@ class PangeaRequestAsync(PangeaRequestBase):
             raise pe.ServiceTemporarilyUnavailable(await resp.json())
 
     async def poll_result_by_id(
-        self, request_id: str, result_class: Union[Type[PangeaResponseResult], dict], check_response: bool = True
+        self, request_id: str, result_class: Union[Type[PangeaResponseResult], Type[dict]], check_response: bool = True
     ):
         path = self._get_poll_path(request_id)
         self.logger.debug(json.dumps({"service": self.service, "action": "poll_result_once", "url": path}))
@@ -157,12 +172,84 @@ class PangeaRequestAsync(PangeaRequestBase):
         if resp.status < 200 or resp.status >= 300:
             raise pe.PresignedUploadError(f"presigned PUT failure: {resp.status}", await resp.text())
 
+    async def download_file(self, url: str, filename: Optional[str] = None) -> AttachedFile:
+        self.logger.debug(
+            json.dumps(
+                {
+                    "service": self.service,
+                    "action": "download_file",
+                    "url": url,
+                    "status": "start",
+                }
+            )
+        )
+        async with self.session.get(url, headers={}) as response:
+            if response.status == 200:
+                if filename is None:
+                    content_disposition = response.headers.get("Content-Disposition", "")
+                    filename = self._get_filename_from_content_disposition(content_disposition)
+                    if filename is None:
+                        filename = self._get_filename_from_url(url)
+                        if filename is None:
+                            filename = "default_filename"
+
+                content_type = response.headers.get("Content-Type", "")
+                self.logger.debug(
+                    json.dumps(
+                        {
+                            "service": self.service,
+                            "action": "download_file",
+                            "url": url,
+                            "filename": filename,
+                            "status": "success",
+                        }
+                    )
+                )
+
+                return AttachedFile(filename=filename, file=await response.read(), content_type=content_type)
+            else:
+                raise pe.DownloadFileError(f"Failed to download file. Status: {response.status}", await response.text())
+
+    async def _get_pangea_json(self, reader: aiohttp.MultipartReader) -> Optional[Dict]:
+        # Iterate through parts
+        async for part in reader:
+            return await part.json()
+        return None
+
+    async def _get_attached_files(self, reader: aiohttp.MultipartReader) -> List[AttachedFile]:
+        files = []
+        i = 0
+
+        async for part in reader:
+            content_type = part.headers.get("Content-Type", "")
+            content_disposition = part.headers.get("Content-Disposition", "")
+            name = self._get_filename_from_content_disposition(content_disposition)
+            if name is None:
+                name = f"default_file_name_{i}"
+                i += 1
+            files.append(AttachedFile(name, await part.read(), content_type))
+
+        return files
+
+    async def _process_multipart_response(self, resp: aiohttp.ClientResponse) -> MultipartResponse:
+        # Parse the multipart response
+        multipart_reader = aiohttp.MultipartReader.from_response(resp)
+
+        pangea_json = await self._get_pangea_json(multipart_reader)  # type: ignore[arg-type]
+        self.logger.debug(
+            json.dumps({"service": self.service, "action": "multipart response", "response": pangea_json})
+        )
+
+        multipart_reader = multipart_reader.__aiter__()
+        attached_files = await self._get_attached_files(multipart_reader)  # type: ignore[arg-type]
+        return MultipartResponse(pangea_json, attached_files)  # type: ignore[arg-type]
+
     async def _http_post(
         self,
         url: str,
         headers: Dict = {},
         data: Union[str, Dict] = {},
-        files: Optional[List[Tuple]] = None,
+        files: List[Tuple] = [],
         presigned_url_post: bool = False,
     ) -> aiohttp.ClientResponse:
         if files:
@@ -193,26 +280,29 @@ class PangeaRequestAsync(PangeaRequestBase):
         self.logger.debug(
             json.dumps({"service": self.service, "action": "http_put", "url": url}, default=default_encoder)
         )
-        form = FormData()
-        name, value = files[0]
-        form.add_field(name, value[1], filename=value[0], content_type=value[2])
-        return await self.session.put(url, headers=headers, data=form)
+        _, value = files[0]
+        return await self.session.put(url, headers=headers, data=value[1])
 
     async def _full_post_presigned_url(
         self,
         endpoint: str,
         result_class: Type[PangeaResponseResult],
         data: Union[str, Dict] = {},
-        files: Optional[List[Tuple]] = None,
+        files: List[Tuple] = [],
     ):
-        if len(files) == 0:  # type: ignore[arg-type]
+        if len(files) == 0:
             raise AttributeError("files attribute should have at least 1 file")
 
         response = await self.request_presigned_url(endpoint=endpoint, result_class=result_class, data=data)
-        data_to_presigned = response.accepted_result.post_form_data  # type: ignore[union-attr]
-        presigned_url = response.accepted_result.post_url  # type: ignore[union-attr]
+        if response.accepted_result is None:
+            raise pe.PangeaException("No accepted_result field when requesting presigned url")
+        if response.accepted_result.post_url is None:
+            raise pe.PresignedURLException("No presigned url", response)
+
+        data_to_presigned = response.accepted_result.post_form_data
+        presigned_url = response.accepted_result.post_url
 
-        await self.post_presigned_url(url=presigned_url, data=data_to_presigned, files=files)  # type: ignore[arg-type]
+        await self.post_presigned_url(url=presigned_url, data=data_to_presigned, files=files)
         return response.raw_response
 
     async def request_presigned_url(
@@ -232,14 +322,16 @@ class PangeaRequestAsync(PangeaRequestBase):
             raise e
 
         # Receive 202
-        return await self._poll_presigned_url(accepted_exception.response)  # type: ignore[return-value]
+        return await self._poll_presigned_url(accepted_exception.response)
 
-    async def _poll_presigned_url(self, response: PangeaResponse) -> AcceptedResult:
+    async def _poll_presigned_url(
+        self, response: PangeaResponse[Type[PangeaResponseResult]]
+    ) -> PangeaResponse[Type[PangeaResponseResult]]:
         if response.http_status != 202:
             raise AttributeError("Response should be 202")
 
         if response.accepted_result is not None and response.accepted_result.has_upload_url:
-            return response  # type: ignore[return-value]
+            return response
 
         self.logger.debug(json.dumps({"service": self.service, "action": "poll_presigned_url", "step": "start"}))
         retry_count = 1
@@ -276,7 +368,7 @@ class PangeaRequestAsync(PangeaRequestBase):
         self.logger.debug(json.dumps({"service": self.service, "action": "poll_presigned_url", "step": "exit"}))
 
         if loop_resp.accepted_result is not None and not loop_resp.accepted_result.has_upload_url:
-            return loop_resp  # type: ignore[return-value]
+            return loop_resp
         else:
             raise loop_exc
 

pangea/asyncio/services/audit.py

@@ -1,18 +1,24 @@
 # Copyright 2022 Pangea Cyber Corporation
 # Author: Pangea Cyber Corporation
 import datetime
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Sequence, Union
 
 import pangea.exceptions as pexc
+from pangea.asyncio.services.base import ServiceBaseAsync
 from pangea.response import PangeaResponse
 from pangea.services.audit.audit import AuditBase
 from pangea.services.audit.exceptions import AuditException
 from pangea.services.audit.models import (
+    DownloadFormat,
+    DownloadRequest,
+    DownloadResult,
     Event,
     LogBulkResult,
     LogResult,
+    PublishedRoot,
     RootRequest,
     RootResult,
+    RootSource,
     SearchOrder,
     SearchOrderBy,
     SearchOutput,
@@ -22,8 +28,6 @@ from pangea.services.audit.models import (
 )
 from pangea.services.audit.util import format_datetime
 
-from .base import ServiceBaseAsync
-
 
 class AuditAsync(ServiceBaseAsync, AuditBase):
     """Audit service client.
@@ -175,8 +179,10 @@ class AuditAsync(ServiceBaseAsync, AuditBase):
         """
 
         input = self._get_log_request(event, sign_local=sign_local, verify=verify, verbose=verbose)
-        response = await self.request.post("v1/log", LogResult, data=input.dict(exclude_none=True))
-        if response.success:
+        response: PangeaResponse[LogResult] = await 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
 
@@ -209,8 +215,10 @@ class AuditAsync(ServiceBaseAsync, AuditBase):
         """
 
         input = self._get_log_request(events, sign_local=sign_local, verify=False, verbose=verbose)
-        response = await self.request.post("v2/log", LogBulkResult, data=input.dict(exclude_none=True))
-        if response.success:
+        response: PangeaResponse[LogBulkResult] = await self.request.post(
+            "v2/log", LogBulkResult, data=input.dict(exclude_none=True)
+        )
+        if response.success and response.result is not None:
             for result in response.result.results:
                 self._process_log_result(result, verify=True)
         return response
@@ -245,12 +253,12 @@ class AuditAsync(ServiceBaseAsync, AuditBase):
 
         input = self._get_log_request(events, sign_local=sign_local, verify=False, verbose=verbose)
         try:
-            response = await self.request.post(
+            response: PangeaResponse[LogBulkResult] = await 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:
             for result in response.result.results:
                 self._process_log_result(result, verify=True)
         return response
@@ -264,7 +272,7 @@ class AuditAsync(ServiceBaseAsync, 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,
@@ -293,7 +301,7 @@ class AuditAsync(ServiceBaseAsync, 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
@@ -326,7 +334,12 @@ class AuditAsync(ServiceBaseAsync, AuditBase):
             verbose=verbose,
         )
 
-        response = await self.request.post("v1/search", SearchOutput, data=input.dict(exclude_none=True))
+        response: PangeaResponse[SearchOutput] = await self.request.post(
+            "v1/search", SearchOutput, data=input.dict(exclude_none=True)
+        )
+        if verify_consistency:
+            await self.update_published_roots(response.result)  # type: ignore[arg-type]
+
         return self.handle_search_response(response, verify_consistency, verify_events)
 
     async def results(
@@ -334,6 +347,7 @@ class AuditAsync(ServiceBaseAsync, 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]:
@@ -348,6 +362,7 @@ class AuditAsync(ServiceBaseAsync, 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:
@@ -365,7 +380,8 @@ class AuditAsync(ServiceBaseAsync, AuditBase):
             result_res: PangeaResponse[SearchResultsOutput] = audit.results(
                 id=search_res.result.id,
                 limit=10,
-                offset=0)
+                offset=0,
+                assert_search_restriction={'source': ["monitor"]})
         """
 
         if limit <= 0:  # type: ignore[operator]
@@ -378,8 +394,12 @@ class AuditAsync(ServiceBaseAsync, AuditBase):
             id=id,
             limit=limit,
             offset=offset,
+            assert_search_restriction=assert_search_restriction,
         )
         response = await self.request.post("v1/results", SearchResultOutput, data=input.dict(exclude_none=True))
+        if verify_consistency and response.result is not None:
+            await self.update_published_roots(response.result)
+
         return self.handle_results_response(response, verify_consistency, verify_events)
 
     async def root(self, tree_size: Optional[int] = None) -> PangeaResponse[RootResult]:
@@ -405,3 +425,67 @@ class AuditAsync(ServiceBaseAsync, AuditBase):
         """
         input = RootRequest(tree_size=tree_size)
         return await self.request.post("v1/root", RootResult, data=input.dict(exclude_none=True))
+
+    async def download_results(
+        self, result_id: str, format: Optional[DownloadFormat] = 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.
+
+        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 = await audit.download_results(
+                result_id="pas_[...]",
+                format=DownloadFormat.JSON,
+            )
+        """
+
+        input = DownloadRequest(result_id=result_id, format=format)
+        return await self.request.post("v1/download_results", DownloadResult, data=input.dict(exclude_none=True))
+
+    async 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 AuditAsync.search() or AuditAsync.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 = await 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/asyncio/services/authn.py

@@ -4,10 +4,9 @@
 from typing import Dict, List, Optional, Union
 
 import pangea.services.authn.models as m
+from pangea.asyncio.services.base import ServiceBaseAsync
 from pangea.response import PangeaResponse
 
-from .base import ServiceBaseAsync
-
 SERVICE_NAME = "authn"
 
 
@@ -114,6 +113,9 @@ class AuthNAsync(ServiceBaseAsync):
             Examples:
                 response = authn.session.list()
             """
+            if isinstance(filter, dict):
+                filter = m.SessionListFilter(**filter)
+
             input = m.SessionListRequest(filter=filter, last=last, order=order, order_by=order_by, size=size)
             return await self.request.post("v2/session/list", m.SessionListResults, data=input.dict(exclude_none=True))
 
@@ -272,6 +274,9 @@ class AuthNAsync(ServiceBaseAsync):
                         token="ptu_wuk7tvtpswyjtlsx52b7yyi2l7zotv4a",
                     )
                 """
+                if isinstance(filter, dict):
+                    filter = m.SessionListFilter(**filter)
+
                 input = m.ClientSessionListRequest(
                     token=token, filter=filter, last=last, order=order, order_by=order_by, size=size
                 )
@@ -593,6 +598,9 @@ class AuthNAsync(ServiceBaseAsync):
             Examples:
                 response = authn.user.list()
             """
+            if isinstance(filter, dict):
+                filter = m.UserListFilter(**filter)
+
             input = m.UserListRequest(
                 filter=filter,
                 last=last,
@@ -642,6 +650,9 @@ class AuthNAsync(ServiceBaseAsync):
                 Examples:
                     response = authn.user.invites.list()
                 """
+                if isinstance(filter, dict):
+                    filter = m.UserInviteListFilter(**filter)
+
                 input = m.UserInviteListRequest(filter=filter, last=last, order=order, order_by=order_by, size=size)
                 return await self.request.post(
                     "v2/user/invite/list", m.UserInviteListResult, data=input.dict(exclude_none=True)
@@ -1064,6 +1075,8 @@ class AuthNAsync(ServiceBaseAsync):
             Examples:
                 response = authn.agreements.list()
             """
+            if isinstance(filter, dict):
+                filter = m.AgreementListFilter(**filter)
 
             input = m.AgreementListRequest(filter=filter, last=last, order=order, order_by=order_by, size=size)
             return await self.request.post(

pangea/asyncio/services/base.py

@@ -1,18 +1,18 @@
 # Copyright 2022 Pangea Cyber Corporation
 # Author: Pangea Cyber Corporation
 
-from typing import Optional, Type, Union
+from typing import Dict, Optional, Type, Union
 
 from pangea.asyncio.request import PangeaRequestAsync
 from pangea.exceptions import AcceptedRequestException
-from pangea.response import PangeaResponse, PangeaResponseResult
-from pangea.services.base import ServiceBase
+from pangea.response import AttachedFile, PangeaResponse, PangeaResponseResult
+from pangea.services.base import PangeaRequest, ServiceBase
 
 
 class ServiceBaseAsync(ServiceBase):
     @property
-    def request(self):
-        if not self._request:
+    def request(self) -> PangeaRequestAsync:  # type: ignore[override]
+        if self._request is None or isinstance(self._request, PangeaRequest):
             self._request = PangeaRequestAsync(
                 config=self.config,
                 token=self.token,
@@ -28,7 +28,7 @@ class ServiceBaseAsync(ServiceBase):
         exception: Optional[AcceptedRequestException] = None,
         response: Optional[PangeaResponse] = None,
         request_id: Optional[str] = None,
-        result_class: Union[Type[PangeaResponseResult], dict] = dict,  # type: ignore[assignment]
+        result_class: Union[Type[PangeaResponseResult], Type[Dict]] = dict,
     ) -> PangeaResponse:
         """
         Poll result
@@ -58,6 +58,9 @@ class ServiceBaseAsync(ServiceBase):
         else:
             raise AttributeError("Need to set exception, response or request_id")
 
+    async def download_file(self, url: str, filename: Optional[str] = None) -> AttachedFile:  # type: ignore[override]
+        return await self.request.download_file(url=url, filename=filename)
+
     async def close(self):
         await self.request.session.close()
         # Loop over all attributes to check if they are derived from ServiceBaseAsync and close them

pangea/asyncio/services/file_scan.py

@@ -2,16 +2,15 @@
 # Author: Pangea Cyber Corporation
 import io
 import logging
-from typing import Dict, Optional
+from typing import Dict, List, Optional, Tuple
 
 import pangea.services.file_scan as m
 from pangea.asyncio.request import PangeaRequestAsync
+from pangea.asyncio.services.base import ServiceBaseAsync
 from pangea.request import PangeaConfig
 from pangea.response import PangeaResponse, TransferMethod
 from pangea.utils import FileUploadParams, get_file_upload_params
 
-from .base import ServiceBaseAsync
-
 
 class FileScanAsync(ServiceBaseAsync):
     """FileScan service client.
@@ -96,7 +95,7 @@ class FileScanAsync(ServiceBaseAsync):
                 size = params.size
             else:
                 crc, sha, size = None, None, None
-            files = [("upload", ("filename", file, "application/octet-stream"))]
+            files: List[Tuple] = [("upload", ("filename", file, "application/octet-stream"))]
         else:
             raise ValueError("Need to set file_path or file arguments")
 

pangea/asyncio/services/intel.py

@@ -4,11 +4,10 @@ import hashlib
 from typing import List, Optional
 
 import pangea.services.intel as m
+from pangea.asyncio.services.base import ServiceBaseAsync
 from pangea.response import PangeaResponse
 from pangea.utils import hash_256_filepath
 
-from .base import ServiceBaseAsync
-
 
 class FileIntelAsync(ServiceBaseAsync):
     """File Intel service client

pangea/asyncio/services/redact.py

@@ -4,10 +4,9 @@
 from typing import Dict, List, Optional, Union
 
 import pangea.services.redact as m
+from pangea.asyncio.services.base import ServiceBaseAsync
 from pangea.response import PangeaResponse
 
-from .base import ServiceBaseAsync
-
 
 class RedactAsync(ServiceBaseAsync):
     """Redact service client.

pangea/asyncio/services/vault.py

@@ -3,6 +3,7 @@
 import datetime
 from typing import Dict, Optional, Union
 
+from pangea.asyncio.services.base import ServiceBaseAsync
 from pangea.response import PangeaResponse
 from pangea.services.vault.models.asymmetric import (
     AsymmetricGenerateRequest,
@@ -47,8 +48,8 @@ from pangea.services.vault.models.common import (
     StateChangeRequest,
     StateChangeResult,
     SymmetricAlgorithm,
-    TDict,
     Tags,
+    TDict,
     UpdateRequest,
     UpdateResult,
 )
@@ -69,8 +70,6 @@ from pangea.services.vault.models.symmetric import (
     SymmetricStoreResult,
 )
 
-from .base import ServiceBaseAsync
-
 
 class VaultAsync(ServiceBaseAsync):
     """Vault service client.
@@ -92,7 +91,7 @@ class VaultAsync(ServiceBaseAsync):
         vault_config = PangeaConfig(domain="pangea.cloud")
 
         # Setup Pangea Vault service
-        vault = Vault(token=PANGEA_VAULT_TOKEN, config=audit_config)
+        vault = Vault(token=PANGEA_VAULT_TOKEN, config=vault_config)
     """
 
     service_name = "vault"
@@ -859,8 +858,8 @@ class VaultAsync(ServiceBaseAsync):
 
                 Default is `deactivated`.
             public_key (EncodedPublicKey, optional): The public key (in PEM format)
-            private_key: (EncodedPrivateKey, optional): The private key (in PEM format)
-            key: (EncodedSymmetricKey, optional): The key material (in base64)
+            private_key (EncodedPrivateKey, optional): The private key (in PEM format)
+            key (EncodedSymmetricKey, optional): The key material (in base64)
 
         Raises:
             PangeaAPIException: If an API Error happens
@@ -1222,7 +1221,7 @@ class VaultAsync(ServiceBaseAsync):
             )
         """
 
-        input = EncryptStructuredRequest(
+        input: EncryptStructuredRequest[TDict] = EncryptStructuredRequest(
             id=id, structured_data=structured_data, filter=filter, version=version, additional_data=additional_data
         )
         return await self.request.post(
@@ -1271,7 +1270,7 @@ class VaultAsync(ServiceBaseAsync):
             )
         """
 
-        input = EncryptStructuredRequest(
+        input: EncryptStructuredRequest[TDict] = EncryptStructuredRequest(
             id=id, structured_data=structured_data, filter=filter, version=version, additional_data=additional_data
         )
         return await self.request.post(