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

pangea/services/base.py

@@ -1,23 +1,40 @@
 # Copyright 2022 Pangea Cyber Corporation
 # Author: Pangea Cyber Corporation
+from __future__ import annotations
 
 import copy
 import logging
 from typing import Dict, Optional, Type, Union
 
+from typing_extensions import TypeVar
+
 from pangea.asyncio.request import PangeaRequestAsync
 from pangea.config import PangeaConfig
 from pangea.exceptions import AcceptedRequestException
 from pangea.request import PangeaRequest
 from pangea.response import AttachedFile, PangeaResponse, PangeaResponseResult
 
+TResult = TypeVar("TResult", bound=PangeaResponseResult, default=PangeaResponseResult)
+
 
 class ServiceBase(object):
     service_name: str = "base"
 
     def __init__(
-        self, token, config: Optional[PangeaConfig] = None, logger_name: str = "pangea", config_id: Optional[str] = None
-    ):
+        self,
+        token: str,
+        config: PangeaConfig | None = None,
+        logger_name: str = "pangea",
+        config_id: str | None = None,
+    ) -> None:
+        """
+        Initializes a new Pangea service client.
+
+        Args:
+            token: Pangea API token.
+            config: Configuration.
+            logger_name: Logger name.
+        """
         if not token:
             raise Exception("No token provided")
 
@@ -30,11 +47,11 @@ class ServiceBase(object):
         self.request.set_extra_headers(extra_headers)
 
     @property
-    def token(self):
+    def token(self) -> str:
         return self._token
 
     @token.setter
-    def token(self, value):
+    def token(self, value: str) -> None:
         self._token = value
 
     @property
@@ -55,15 +72,16 @@ class ServiceBase(object):
         exception: Optional[AcceptedRequestException] = None,
         response: Optional[PangeaResponse] = None,
         request_id: Optional[str] = None,
-        result_class: Union[Type[PangeaResponseResult], Type[Dict]] = dict,
-    ) -> PangeaResponse:
+        result_class: Type[TResult] = PangeaResponseResult,  # type: ignore[assignment]
+    ) -> PangeaResponse[TResult]:
         """
         Poll result
 
         Returns request's result that has been accepted by the server
 
         Args:
-            exception (AcceptedRequestException): Exception raise by SDK on the call that is been processed.
+            exception: Exception that was previously raised by the SDK on a call
+              that is being processed.
 
         Returns:
             PangeaResponse
@@ -83,5 +101,18 @@ class ServiceBase(object):
         else:
             raise AttributeError("Need to set exception, response or request_id")
 
-    def download_file(self, url: str, filename: Optional[str] = None) -> AttachedFile:
+    def download_file(self, url: str, filename: str | None = None) -> AttachedFile:
+        """
+        Download file
+
+        Download a file from the specified URL and save it with the given
+        filename.
+
+        Args:
+            url: URL of the file to download
+            filename: Name to save the downloaded file as. If not provided, the
+              filename will be determined from the Content-Disposition header or
+              the URL.
+        """
+
         return self.request.download_file(url=url, filename=filename)

pangea/services/embargo.py

@@ -103,7 +103,7 @@ class Embargo(ServiceBase):
             response = embargo.ip_check("190.6.64.94")
         """
         input = IPCheckRequest(ip=ip)
-        return self.request.post("v1/ip/check", EmbargoResult, data=input.dict())
+        return self.request.post("v1/ip/check", EmbargoResult, data=input.model_dump())
 
     def iso_check(self, iso_code: str) -> PangeaResponse[EmbargoResult]:
         """
@@ -130,4 +130,4 @@ class Embargo(ServiceBase):
             response = embargo.iso_check("CU")
         """
         input = ISOCheckRequest(iso_code=iso_code)
-        return self.request.post("v1/iso/check", result_class=EmbargoResult, data=input.dict())
+        return self.request.post("v1/iso/check", result_class=EmbargoResult, data=input.model_dump())

pangea/services/file_scan.py

@@ -11,22 +11,25 @@ from pangea.utils import FileUploadParams, get_file_upload_params
 
 
 class FileScanRequest(APIRequestModel):
-    """
-    File Scan request data
-
-    provider (str, optional): Provider of the information. Default provider defined by the configuration.
-    verbose (bool, optional): Echo back the parameters of the API in the response
-    raw (bool, optional): Return additional details from the provider.
-    """
+    """File Scan request data."""
 
     verbose: Optional[bool] = None
+    """Echo back the parameters of the API in the response."""
+
     raw: Optional[bool] = None
+    """Return additional details from the provider."""
+
     provider: Optional[str] = None
+    """Provider of the information. Default provider defined by the configuration."""
+
     size: Optional[int] = None
     crc32c: Optional[str] = None
     sha256: Optional[str] = None
     source_url: Optional[str] = None
+    """A URL where the file to be scanned can be downloaded."""
+
     transfer_method: TransferMethod = TransferMethod.POST_URL
+    """The transfer method used to upload the file data."""
 
 
 class FileScanData(PangeaResponseResult):
@@ -71,7 +74,6 @@ class FileScan(ServiceBase):
     """
 
     service_name = "file-scan"
-    version = "v1"
 
     def file_scan(
         self,
@@ -92,12 +94,14 @@ class FileScan(ServiceBase):
         OperationId: file_scan_post_v1_scan
 
         Args:
-            file (io.BufferedReader, optional): file to be scanned (should be opened with read permissions and in binary format)
             file_path (str, optional): filepath to be opened and scanned
+            file (io.BufferedReader, optional): file to be scanned (should be opened with read permissions and in binary format)
             verbose (bool, optional): Echo the API parameters in the response
             raw (bool, optional): Include raw data from this provider
             provider (str, optional): Scan file using this provider
             sync_call (bool, optional): True to wait until server returns a result, False to return immediately and retrieve result asynchronously
+            transfer_method (TransferMethod, optional): Transfer method used to upload the file data.
+            source_url (str, optional): A URL where the Pangea APIs can fetch the contents of the input file.
 
         Raises:
             PangeaAPIException: If an API Error happens
@@ -118,6 +122,15 @@ class FileScan(ServiceBase):
                     print(f"\\t{err.detail} \\n")
         """
 
+        if transfer_method == TransferMethod.SOURCE_URL and source_url is None:
+            raise ValueError("`source_url` argument is required when using `TransferMethod.SOURCE_URL`.")
+
+        if source_url is not None and transfer_method != TransferMethod.SOURCE_URL:
+            raise ValueError(
+                "`transfer_method` should be `TransferMethod.SOURCE_URL` when using the `source_url` argument."
+            )
+
+        files: Optional[List[Tuple]] = None
         if file or file_path:
             if file_path:
                 file = open(file_path, "rb")
@@ -128,9 +141,9 @@ class FileScan(ServiceBase):
                 size = params.size
             else:
                 crc, sha, size = None, None, None
-            files: List[Tuple] = [("upload", ("filename", file, "application/octet-stream"))]
-        else:
-            raise ValueError("Need to set file_path or file arguments")
+            files = [("upload", ("filename", file, "application/octet-stream"))]
+        elif source_url is None:
+            raise ValueError("Need to set one of `file_path`, `file`, or `source_url` arguments.")
 
         input = FileScanRequest(
             verbose=verbose,
@@ -142,8 +155,12 @@ class FileScan(ServiceBase):
             transfer_method=transfer_method,
             source_url=source_url,
         )
-        data = input.dict(exclude_none=True)
-        return self.request.post("v1/scan", FileScanResult, data=data, files=files, poll_result=sync_call)
+        data = input.model_dump(exclude_none=True)
+        try:
+            return self.request.post("v1/scan", FileScanResult, data=data, files=files, poll_result=sync_call)
+        finally:
+            if file_path and file:
+                file.close()
 
     def request_upload_url(
         self,
@@ -164,7 +181,7 @@ class FileScan(ServiceBase):
             input.sha256 = params.sha256_hex
             input.size = params.size
 
-        data = input.dict(exclude_none=True)
+        data = input.model_dump(exclude_none=True)
         return self.request.request_presigned_url("v1/scan", FileScanResult, data=data)