✨ add support for HMAC signature usage for localresponses

Author: sebastianMindee

mindee/client.py

@@ -190,9 +190,9 @@ def load_prediction(
         :return:
         """
         try:
-            if local_response.json.get("job"):
-                return AsyncPredictResponse(product_class, local_response.json)
-            return PredictResponse(product_class, local_response.json)
+            if local_response.as_dict.get("job"):
+                return AsyncPredictResponse(product_class, local_response.as_dict)
+            return PredictResponse(product_class, local_response.as_dict)
         except KeyError as exc:
             raise MindeeError("No prediction found in local response.") from exc
 

mindee/input/local_response.py

@@ -1,29 +1,90 @@
+import hashlib
+import hmac
+import io
 import json
 from pathlib import Path
-from typing import Union, BinaryIO, Dict, Any
-import io
+from typing import Any, BinaryIO, Dict, Union
 
 from mindee.error import MindeeError
 
 
 class LocalResponse:
-    json: Dict[str, Any]
+    """Local response loaded from a file."""
+
+    _file: BinaryIO
+    """File object of the local response."""
 
     def __init__(self, input_file: Union[BinaryIO, str, Path, bytes]):
-        input_binary: BinaryIO
         if isinstance(input_file, BinaryIO):
-            input_binary = input_file
-            input_binary.seek(0)
-        elif isinstance(input_file, str) or isinstance(input_file, Path):
-            with open(input_file, 'rb') as f:
-                input_binary = io.BytesIO(f.read())
+            self._file = input_file
+            self._file.seek(0)
+        elif isinstance(input_file, (str, Path)):
+            with open(input_file, "r", encoding="utf-8") as file:
+                self._file = io.BytesIO(
+                    file.read().replace("\r", "").replace("\n", "").encode()
+                )
         elif isinstance(input_file, bytes):
-            input_binary = io.BytesIO(input_file)
+            self._file = io.BytesIO(input_file)
         else:
-            raise TypeError('Incompatible type for input.')
+            raise MindeeError("Incompatible type for input.")
+
+    @property
+    def as_dict(self) -> Dict[str, Any]:
+        """
+        Returns the dictionary representation of the file.
+
+        :return: A json-like dictionary.
+        """
         try:
-            input_binary.seek(0)
-            self.json = json.load(input_binary)
-            input_binary.close()
+            self._file.seek(0)
+            out_json = json.loads(self._file.read())
         except json.decoder.JSONDecodeError as exc:
-            raise MindeeError('File is not a valid dictionary.') from exc
+            raise MindeeError("File is not a valid dictionary.") from exc
+        return out_json
+
+    @staticmethod
+    def _process_secret_key(
+        secret_key: Union[str, bytes, bytearray]
+    ) -> Union[bytes, bytearray]:
+        """
+        Processes the secret key as a byte array.
+
+        :param secret_key: Secret key, either a string or a byte/byte array.
+        :return: a byte/byte array secret key.
+        """
+        if isinstance(secret_key, (bytes, bytearray)):
+            return secret_key
+        return secret_key.encode("utf-8")
+
+    def get_hmac_signature(self, secret_key: Union[str, bytes, bytearray]):
+        """
+        Returns the hmac signature of the local response, from the secret key provided.
+
+        :param secret_key: Secret key, either a string or a byte/byte array.
+        :return: The hmac signature of the local response.
+        """
+        algorithm = hashlib.sha256
+
+        try:
+            self._file.seek(0)
+            mac = hmac.new(
+                LocalResponse._process_secret_key(secret_key),
+                self._file.read(),
+                algorithm,
+            )
+        except (TypeError, ValueError) as exc:
+            raise MindeeError("Could not get HMAC signature from payload.") from exc
+
+        return mac.hexdigest()
+
+    def is_valid_hmac_signature(
+        self, secret_key: Union[str, bytes, bytearray], signature: str
+    ):
+        """
+        Checks if the hmac signature of the local response is valid.
+
+        :param secret_key: Secret key, given as a string.
+        :param signature:
+        :return: True if the HMAC signature is valid.
+        """
+        return signature == self.get_hmac_signature(secret_key)

tests/Input/test_local_response.py

@@ -0,0 +1,44 @@
+from pathlib import Path
+
+import pytest
+
+from mindee.input import LocalResponse
+from tests.api.test_async_response import ASYNC_DIR
+
+
+@pytest.fixture
+def dummy_secret_key():
+    return "ogNjY44MhvKPGTtVsI8zG82JqWQa68woYQH"
+
+@pytest.fixture
+def signature():
+    return "5ed1673e34421217a5dbfcad905ee62261a3dd66c442f3edd19302072bbf70d0"
+
+@pytest.fixture
+def file_path():
+    return Path(ASYNC_DIR / "get_completed_empty.json")
+
+def test_valid_file_local_response(dummy_secret_key, signature, file_path):
+    local_response = LocalResponse(file_path)
+    assert local_response._file is not None
+    assert not local_response.is_valid_hmac_signature(dummy_secret_key, "invalid signature")
+    assert signature == local_response.get_hmac_signature(dummy_secret_key)
+    assert local_response.is_valid_hmac_signature(dummy_secret_key, signature)
+
+def test_valid_path_local_response(dummy_secret_key, signature, file_path):
+    local_response = LocalResponse(file_path)
+    assert local_response._file is not None
+    assert not local_response.is_valid_hmac_signature(dummy_secret_key, "invalid signature")
+    assert signature == local_response.get_hmac_signature(dummy_secret_key)
+    assert local_response.is_valid_hmac_signature(dummy_secret_key, signature)
+
+def test_valid_bytes_local_response(dummy_secret_key, signature, file_path):
+    with open(file_path, 'r') as f:
+        str_response = f.read().replace('\r', '').replace('\n', '')
+    file_bytes = str_response.encode('utf-8')
+    local_response = LocalResponse(file_bytes)
+    assert local_response._file is not None
+    assert not local_response.is_valid_hmac_signature(dummy_secret_key, "invalid signature")
+    assert signature == local_response.get_hmac_signature(dummy_secret_key)
+    assert local_response.is_valid_hmac_signature(dummy_secret_key, signature)
+

tests/api/test_async_response.py

@@ -1,4 +1,5 @@
 import json
+from pathlib import Path
 
 import pytest
 import requests
@@ -10,7 +11,7 @@
 from mindee.parsing.common.async_predict_response import AsyncPredictResponse
 from mindee.product.invoice_splitter import InvoiceSplitterV1
 
-ASYNC_DIR = "./tests/data/async"
+ASYNC_DIR = Path("./tests/data/async")
 
 FILE_PATH_POST_SUCCESS = f"{ASYNC_DIR}/post_success.json"
 FILE_PATH_POST_FAIL = f"{ASYNC_DIR}/post_fail_forbidden.json"

tests/test_client.py

@@ -143,10 +143,11 @@ def test_local_response_from_async_json(dummy_client: Client):
 
 
 def test_local_response_from_invalid_file(dummy_client: Client):
+    local_response = LocalResponse(
+        PRODUCT_DATA_DIR / "invoices" / "response_v4" / "summary_full.rst"
+    )
     with pytest.raises(MindeeError):
-        LocalResponse(
-            PRODUCT_DATA_DIR / "invoices" / "response_v4" / "summary_full.rst"
-        )
+        print(local_response.as_dict)
 
 
 def test_local_response_from_invalid_dict(dummy_client: Client):