pangea-sdk 3.8.0__py3-none-any.whl → 5.3.0__py3-none-any.whl

pangea/tools.py

@@ -95,7 +95,7 @@ def file_events(root_hashes: Dict[int, str], f: io.TextIOWrapper) -> Iterator[Ev
             else:
                 raise ValueError("invalid data")
         except (json.JSONDecodeError, ValueError, KeyError) as e:
-            exit_with_error(f"failed to parse line {idx}: {str(e)}")
+            exit_with_error(f"failed to parse line {idx}: {e!s}")
 
 
 def init_audit(token: str, domain: str) -> Audit:
@@ -108,15 +108,14 @@ def init_audit(token: str, domain: str) -> Audit:
 def make_aware_datetime(d: datetime) -> datetime:
     if d.tzinfo is None or d.tzinfo.utcoffset(d) is None:
         return d.replace(tzinfo=timezone.utc)
-    else:
-        return d
+    return d
 
 
 def filter_deep_none(data: Dict) -> Dict:
     return {k: v if not isinstance(v, Dict) else filter_deep_none(v) for k, v in data.items() if v is not None}
 
 
-def _load_env_var(env_var_name: str):
+def _load_env_var(env_var_name: str) -> str:
     value = os.getenv(env_var_name)
     if not value:
         raise PangeaException(f"{env_var_name} env var need to be set")
@@ -124,12 +123,12 @@ def _load_env_var(env_var_name: str):
     return value
 
 
-def get_test_domain(environment: TestEnvironment):
+def get_test_domain(environment: TestEnvironment) -> str:
     env_var_name = f"PANGEA_INTEGRATION_DOMAIN_{environment}"
     return _load_env_var(env_var_name)
 
 
-def get_test_token(environment: TestEnvironment):
+def get_test_token(environment: TestEnvironment) -> str:
     env_var_name = f"PANGEA_INTEGRATION_TOKEN_{environment}"
     return _load_env_var(env_var_name)
 
@@ -200,7 +199,7 @@ loggers: Dict[str, bool] = {}
 
 
 def logger_set_pangea_config(logger_name: str, level=logging.DEBUG):
-    if loggers.get(logger_name, None) is not None:
+    if loggers.get(logger_name) is not None:
         return
 
     loggers[logger_name] = True

pangea/utils.py

@@ -1,12 +1,13 @@
+from __future__ import annotations
+
 import base64
 import copy
 import datetime
 import io
 import json
-from collections import OrderedDict
-from hashlib import new, sha1, sha256, sha512
+from hashlib import md5, new, sha1, sha256, sha512
 
-from google_crc32c import Checksum as CRC32C  # type: ignore[import]
+from google_crc32c import Checksum as CRC32C
 from pydantic import BaseModel
 
 
@@ -34,20 +35,9 @@ def str2str_b64(data: str, encoding: str = "utf-8") -> str:
     return base64.b64encode(data.encode(encoding)).decode("ascii")
 
 
-def dict_order_keys(data: dict) -> OrderedDict:
-    if isinstance(data, dict):
-        return OrderedDict(sorted(data.items()))
-    else:
-        return data
-
-
-def dict_order_keys_recursive(data: dict) -> OrderedDict:
-    if isinstance(data, dict):
-        for k, v in data.items():
-            if type(v) is dict:
-                data[k] = dict_order_keys_recursive(v)
-
-    return data  # type: ignore[return-value]
+def str_b64_2bytes(data: str) -> bytes:
+    data += "=" * ((4 - len(data) % 4) % 4)  # add padding if needed
+    return base64.urlsafe_b64decode(data)
 
 
 def canonicalize_nested_json(data: dict) -> dict:
@@ -76,33 +66,97 @@ def canonicalize(data: dict) -> str:
         return str(data)
 
 
-def hash_sha256(data: str) -> str:
-    # Return sha256 hash in hex format
-    return sha256(data.encode("ascii")).hexdigest()
+def hash_sha256(input: str | io.BufferedReader) -> str:
+    # Return SHA256 hash in hex format
+    hash = sha256()
+    if isinstance(input, io.BufferedReader):
+        input.seek(0)  # restart reading
+        while True:
+            chunk = input.read(1024 * 1024)
+            if not chunk:
+                break
+            hash.update(chunk)
 
+        input.seek(0)  # restart reading
+    else:
+        hash.update(input.encode("utf-8"))
+
+    return hash.hexdigest()
 
-def hash_256_filepath(filepath: str) -> str:
-    data = open(filepath, "rb")
-    hash = sha256(data.read()).hexdigest()
-    data.close()
-    return hash
 
+def hash_sha1(input: str | io.BufferedReader) -> str:
+    # Return SHA1 hash in hex format
+    hash = sha1()
+    if isinstance(input, io.BufferedReader):
+        input.seek(0)  # restart reading
+        while True:
+            chunk = input.read(1024 * 1024)
+            if not chunk:
+                break
+            hash.update(chunk)
+
+        input.seek(0)  # restart reading
+    else:
+        hash.update(input.encode("utf-8"))
+
+    return hash.hexdigest()
 
-def hash_sha1(data: str) -> str:
-    # Return sha1 hash in hex format
-    return sha1(data.encode("ascii")).hexdigest()
 
+def hash_sha512(input: str | io.BufferedReader) -> str:
+    # Return SHA512 hash in hex format
+    hash = sha512()
+    if isinstance(input, io.BufferedReader):
+        input.seek(0)  # restart reading
+        while True:
+            chunk = input.read(1024 * 1024)
+            if not chunk:
+                break
+            hash.update(chunk)
+
+        input.seek(0)  # restart reading
+    else:
+        hash.update(input.encode("utf-8"))
 
-def hash_sha512(data: str) -> str:
-    # Return sha512 hash in hex format
-    return sha512(data.encode("ascii")).hexdigest()
+    return hash.hexdigest()
 
 
-def hash_ntlm(data: str):
-    # Calculate the NTLM hash
+def hash_ntlm(data: str) -> str:
+    # Return NTLM hash in hex format
     return new("md4", data.encode("utf-16le")).hexdigest()
 
 
+def hash_md5(input: str | io.BufferedReader) -> str:
+    # Return MD5 hash in hex format
+    hash = md5()
+    if isinstance(input, io.BufferedReader):
+        input.seek(0)  # restart reading
+
+        while True:
+            chunk = input.read(1024 * 1024)
+            if not chunk:
+                break
+            hash.update(chunk)
+
+        input.seek(0)  # restart reading
+    else:
+        hash.update(input.encode("utf-8"))
+
+    return hash.hexdigest()
+
+
+def get_crc32c(data: str) -> str:
+    crc = CRC32C()
+    crc.update(data)
+    return crc.hexdigest().decode("utf-8")
+
+
+def hash_256_filepath(filepath: str) -> str:
+    data = open(filepath, "rb")
+    hash = sha256(data.read()).hexdigest()
+    data.close()
+    return hash
+
+
 def get_prefix(hash: str, len: int = 5):
     return hash[0:len]
 
@@ -132,3 +186,10 @@ def get_file_upload_params(file: io.BufferedReader) -> FileUploadParams:
 
     file.seek(0)  # restart reading
     return FileUploadParams(crc_hex=crc.hexdigest().decode("utf-8"), sha256_hex=sha.hexdigest(), size=size)
+
+
+def get_file_size(file: io.BufferedReader) -> int:
+    file.seek(0, io.SEEK_END)
+    size = file.tell()
+    file.seek(0)  # restart reading
+    return size