Merge pull request #371 from humanprotocol/fix-store-artifact

Fix store artifact to public or private bucket

hmt_escrow/storage.py

@@ -29,6 +29,13 @@
 ESCROW_AWS_ACCESS_KEY_ID = os.getenv("ESCROW_AWS_ACCESS_KEY_ID", "minio")
 ESCROW_AWS_SECRET_ACCESS_KEY = os.getenv("ESCROW_AWS_SECRET_ACCESS_KEY", "minio123")
 
+ESCROW_RESULTS_AWS_S3_ACCESS_KEY_ID = os.getenv(
+    "ESCROW_RESULTS_AWS_S3_ACCESS_KEY_ID", ""
+)
+ESCROW_RESULTS_AWS_S3_SECRET_ACCESS_KEY = os.getenv(
+    "ESCROW_RESULTS_AWS_S3_SECRET_ACCESS_KEY", ""
+)
+
 ESCROW_AWS_REGION = os.getenv("ESCROW_AWS_REGION", "us-west-2")
 
 ESCROW_ENDPOINT_URL = os.getenv("ESCROW_ENDPOINT_URL", "http://minio:9000")
@@ -47,22 +54,29 @@ class StorageFileNotFoundError(StorageClientError):
     pass
 
 
-def _connect_s3():
+def _connect_s3(use_public_bucket=False):
     try:
-        return boto3.client(
-            "s3",
-            aws_access_key_id=ESCROW_AWS_ACCESS_KEY_ID,
-            aws_secret_access_key=ESCROW_AWS_SECRET_ACCESS_KEY,
-            endpoint_url=ESCROW_ENDPOINT_URL,
-            region_name=ESCROW_AWS_REGION,
-        )
+        if use_public_bucket:
+            return boto3.client(
+                "s3",
+                aws_access_key_id=ESCROW_RESULTS_AWS_S3_ACCESS_KEY_ID,
+                aws_secret_access_key=ESCROW_RESULTS_AWS_S3_SECRET_ACCESS_KEY,
+            )
+        else:
+            return boto3.client(
+                "s3",
+                aws_access_key_id=ESCROW_AWS_ACCESS_KEY_ID,
+                aws_secret_access_key=ESCROW_AWS_SECRET_ACCESS_KEY,
+                endpoint_url=ESCROW_ENDPOINT_URL,
+                region_name=ESCROW_AWS_REGION,
+            )
     except Exception as e:
         LOG.error(f"Connection with S3 failed because of: {e}")
         raise e
 
 
 def get_bucket(public: bool = False) -> str:
-    """Retrieves correct bucket according to ACL visibility.
+    """Retrieves correct bucket (private/public).
 
     Args:
           public(bool): whether the public bucket should be retrieved or internal one.
@@ -172,8 +186,8 @@ def download(key: str, private_key: bytes, public: bool = False) -> Dict:
 def upload(
     msg: Dict,
     public_key: bytes,
-    encrypt_data: Optional[bool] = True,
-    use_public_bucket: Optional[bool] = False,
+    encrypt_data=True,
+    use_public_bucket=False,
 ) -> Tuple[str, str]:
     """Upload and encrypt a string for later retrieval.
     This can be manifest files, results, or anything that's been already
@@ -203,21 +217,19 @@ def upload(
     hash_ = hashlib.sha1(content).hexdigest()
     key = f"s3{hash_}"
 
-    # If encryption is on, even if usage of public bucket is true, encrypted data is always private
-    is_public = use_public_bucket is True and encrypt_data is False
-    bucket_name = get_bucket(public=is_public)
+    # Get private or public bucket name
+    bucket_name = get_bucket(public=use_public_bucket)
 
-    bucket_kwargs: Dict[str, Union[str, bytes]] = {"Key": key, "Bucket": bucket_name}
+    # If encryption is on, use crypto.encrypt function, else use utf-8 encoded artifact
+    body = crypto.encrypt(public_key, artifact) if encrypt_data is True else content
+    bucket_kwargs: Dict[str, Union[str, bytes]] = {
+        "Body": body,
+        "Bucket": bucket_name,
+        "Key": key,
+    }
 
-    if encrypt_data is True:
-        # If encryption is enabled, the bucket is private and data is encrypted
-        bucket_kwargs.update({"Body": crypto.encrypt(public_key, artifact)})
-    else:
-        # If encryption is off, file will be publicly readable
-        bucket_kwargs.update({"ACL": "public-read", "Body": content})
-
-    BOTO3_CLIENT = _connect_s3()
-    BOTO3_CLIENT.put_object(**bucket_kwargs)
+    boto3_client = _connect_s3(use_public_bucket)
+    boto3_client.put_object(**bucket_kwargs)
 
     LOG.debug(f"Uploaded to S3, key: {key}")
     return hash_, key

hmt_escrow/utils.py

@@ -83,7 +83,8 @@ def parse_transfer_transaction(
         return hmt_transferred, tx_balance
 
     transfer_event = hmtoken_contract.events.Transfer().processReceipt(tx_receipt)
-    print(f"transfer_event {transfer_event}")
+    logger.info(f"Transfer_event {transfer_event}.")
+
     hmt_transferred = bool(transfer_event)
 
     if hmt_transferred:

package.json

@@ -1,6 +1,6 @@
 {
   "name": "hmt-escrow",
-  "version": "0.14.3",
+  "version": "0.14.4",
   "description": "Launch escrow contracts to the HUMAN network",
   "main": "truffle.js",
   "directories": {

setup.py

@@ -2,7 +2,7 @@
 
 setuptools.setup(
     name="hmt-escrow",
-    version="0.14.3",
+    version="0.14.4",
     author="HUMAN Protocol",
     description="A python library to launch escrow contracts to the HUMAN network.",
     url="https://github.com/humanprotocol/hmt-escrow",

test/hmt_escrow/storage/test_bucket.py

@@ -1,11 +1,16 @@
 import unittest
 from unittest.mock import patch
 
-from hmt_escrow.storage import get_bucket, get_public_bucket_url, get_key_from_url
+from hmt_escrow.storage import _connect_s3, get_bucket, get_public_bucket_url, get_key_from_url
 
 ESCROW_TEST_BUCKETNAME = "test-escrow-results"
 ESCROW_TEST_PUBLIC_BUCKETNAME = "test-escrow-public-results"
 
+ESCROW_RESULTS_AWS_S3_ACCESS_KEY_ID = "ESCROW_RESULTS_AWS_S3_ACCESS_KEY_ID"
+ESCROW_RESULTS_AWS_S3_SECRET_ACCESS_KEY = "ESCROW_RESULTS_AWS_S3_SECRET_ACCESS_KEY"
+ESCROW_AWS_ACCESS_KEY_ID = "ESCROW_AWS_ACCESS_KEY_ID"
+ESCROW_AWS_SECRET_ACCESS_KEY = "ESCROW_AWS_SECRET_ACCESS_KEY"
+
 ESCROW_TEST_PUBLIC_RESULTS_URL = 'http://my-public-bucket.s3.amazon.com'
 
 
@@ -44,6 +49,39 @@ def test_key_retrieval_from_file_url(self):
         key = get_key_from_url(key)
         self.assertEqual(key, expected_key)
 
+    @patch("hmt_escrow.storage.ESCROW_AWS_ACCESS_KEY_ID", ESCROW_AWS_ACCESS_KEY_ID)
+    @patch("hmt_escrow.storage.ESCROW_AWS_SECRET_ACCESS_KEY", ESCROW_AWS_SECRET_ACCESS_KEY)
+    @patch("hmt_escrow.storage.boto3")
+    def test_connect_private_bucket_when_param_passed(self, boto3):
+        """ Tests connection to private bucket with False param. """
+
+        _connect_s3(False)
+        self.assertIn("endpoint_url", boto3.client.call_args.kwargs.keys())
+        self.assertEqual(boto3.client.call_args.kwargs["aws_access_key_id"], ESCROW_AWS_ACCESS_KEY_ID)
+        self.assertEqual(boto3.client.call_args.kwargs["aws_secret_access_key"], ESCROW_AWS_SECRET_ACCESS_KEY)
+
+    @patch("hmt_escrow.storage.ESCROW_AWS_ACCESS_KEY_ID", ESCROW_AWS_ACCESS_KEY_ID)
+    @patch("hmt_escrow.storage.ESCROW_AWS_SECRET_ACCESS_KEY", ESCROW_AWS_SECRET_ACCESS_KEY)
+    @patch("hmt_escrow.storage.boto3")
+    def test_connect_private_bucket_without_param(self, boto3):
+        """ Tests connection to private bucket without param. """
+
+        _connect_s3(False)
+        self.assertIn("endpoint_url", boto3.client.call_args.kwargs.keys())
+        self.assertEqual(boto3.client.call_args.kwargs["aws_access_key_id"], ESCROW_AWS_ACCESS_KEY_ID)
+        self.assertEqual(boto3.client.call_args.kwargs["aws_secret_access_key"], ESCROW_AWS_SECRET_ACCESS_KEY)
+
+    @patch("hmt_escrow.storage.ESCROW_RESULTS_AWS_S3_ACCESS_KEY_ID", ESCROW_RESULTS_AWS_S3_ACCESS_KEY_ID)
+    @patch("hmt_escrow.storage.ESCROW_RESULTS_AWS_S3_SECRET_ACCESS_KEY", ESCROW_RESULTS_AWS_S3_SECRET_ACCESS_KEY)
+    @patch("hmt_escrow.storage.boto3")
+    def test_connect_public_bucket(self, boto3):
+        """ Tests connection to public bucket. """
+
+        _connect_s3(True)
+        self.assertNotIn("endpoint_url", boto3.client.call_args.kwargs.keys())
+        self.assertEqual(boto3.client.call_args.kwargs["aws_access_key_id"], ESCROW_RESULTS_AWS_S3_ACCESS_KEY_ID)
+        self.assertEqual(boto3.client.call_args.kwargs["aws_secret_access_key"], ESCROW_RESULTS_AWS_S3_SECRET_ACCESS_KEY)
+
 
 if __name__ == "__main__":
     unittest.main(exit=True)

test/hmt_escrow/storage/test_storage.py

@@ -28,52 +28,23 @@ def setUp(self) -> None:
         self.pub_key = b"2dbc2c2c86052702e7c219339514b2e8bd4687ba1236c478ad41b43330b08488c12c8c1797aa181f3a4596a1bd8a0c18344ea44d6655f61fa73e56e743f79e0d"
         self.priv_key = b"28e516f1e2f99e96a48a23cea1f94ee5f073403a1c68e818263f0eb898f1c8e5"
 
+    @patch("hmt_escrow.storage.ESCROW_PUBLIC_BUCKETNAME", ESCROW_TEST_PUBLIC_BUCKETNAME)
     @patch("hmt_escrow.storage.ESCROW_BUCKETNAME", ESCROW_TEST_BUCKETNAME)
-    def test_upload_to_private_bucket_when_encryption_on(self):
+    def test_upload_to_private_bucket(self):
         """
-        Tests uploading file to storage to private bucket when encryption is on, no matter the decision to use public
-        bucket.
+        Tests uploading file to storage to private bucket when encryption is on.
         """
 
         s3_client_mock = MagicMock()
         with patch('hmt_escrow.storage._connect_s3') as mock_s3:
             mock_s3.return_value = s3_client_mock
 
-            # Uploads file with encryption on and choice not to use public bucket
             upload(self.get_manifest(), self.pub_key, encrypt_data=True, use_public_bucket=False)
 
             mock_s3.assert_called()
             self.assertIn('Bucket', s3_client_mock.put_object.call_args.kwargs.keys())
 
-            # With encryption on, bucket MUST be the private one
-            self.assertEqual(s3_client_mock.put_object.call_args.kwargs['Bucket'], ESCROW_TEST_BUCKETNAME)
-
-            s3_client_mock.reset_mock()
-
-            # Uploads file with encryption on and choice to use public bucket
-            upload(self.get_manifest(), self.pub_key, encrypt_data=True, use_public_bucket=True)
-
-            mock_s3.assert_called()
-            self.assertIn('Bucket', s3_client_mock.put_object.call_args.kwargs.keys())
-
-            # With encryption on, bucket MUST be the private one
-            self.assertEqual(s3_client_mock.put_object.call_args.kwargs['Bucket'], ESCROW_TEST_BUCKETNAME)
-
-    @patch("hmt_escrow.storage.ESCROW_BUCKETNAME", ESCROW_TEST_BUCKETNAME)
-    def test_upload_to_private_bucket_when_encryption_off(self):
-        """ Tests uploading file to storage to private bucket when encryption is off. """
-
-        s3_client_mock = MagicMock()
-        with patch('hmt_escrow.storage._connect_s3') as mock_s3:
-            mock_s3.return_value = s3_client_mock
-
-            # Uploads file with encryption on
-            upload(self.get_manifest(), self.pub_key, encrypt_data=False, use_public_bucket=False)
-
-            mock_s3.assert_called()
-            self.assertIn('Bucket', s3_client_mock.put_object.call_args.kwargs.keys())
-
-            # With encryption on, bucket MUST be the private one
+            # With use_public_bucket False, bucket MUST be the private one
             self.assertEqual(s3_client_mock.put_object.call_args.kwargs['Bucket'], ESCROW_TEST_BUCKETNAME)
 
     @patch("hmt_escrow.storage.ESCROW_PUBLIC_BUCKETNAME", ESCROW_TEST_PUBLIC_BUCKETNAME)
@@ -84,19 +55,18 @@ def test_upload_to_public_bucket(self):
         with patch('hmt_escrow.storage._connect_s3') as mock_s3:
             mock_s3.return_value = s3_client_mock
 
-            # Uploads file with encryption off
-            upload(self.get_manifest(), self.pub_key, encrypt_data=False, use_public_bucket=True)
+            upload(self.get_manifest(), self.pub_key, encrypt_data=True, use_public_bucket=True)
 
             mock_s3.assert_called()
 
             self.assertIn('Bucket', s3_client_mock.put_object.call_args.kwargs.keys())
 
-            # With encryption on, bucket MUST be the private one
+            # With use_public_bucket True, bucket MUST be the public one
             self.assertIn(s3_client_mock.put_object.call_args.kwargs['Bucket'], ESCROW_TEST_PUBLIC_BUCKETNAME)
 
-    def test_upload_with_encryption_option(self):
+    def test_upload_with_enabled_encryption_option(self):
         """
-        Tests whether data must be persisted in storage encrypted or plain.
+        Tests data persisted in storage is encrypted.
         """
         s3_client_mock = MagicMock()
         with patch('hmt_escrow.storage._connect_s3') as mock_s3:
@@ -112,9 +82,14 @@ def test_upload_with_encryption_option(self):
             # Data to be uploaded must be encrypted
             uploaded_content = crypto.decrypt(self.priv_key, s3_client_mock.put_object.call_args.kwargs['Body'])
             self.assertEqual(json.dumps(data, sort_keys=True), uploaded_content)
-
-            s3_client_mock.reset_mock()
-
+
+    def test_upload_with_disabled_encryption_option(self):
+        """
+        Tests data persisted in storage is plain.
+        """
+        s3_client_mock = MagicMock()
+        with patch('hmt_escrow.storage._connect_s3') as mock_s3:
+            mock_s3.return_value = s3_client_mock
             # Encryption off.
             data = self.get_manifest()
             upload(data, self.pub_key, encrypt_data=False)

test/hmt_escrow/test_job.py

@@ -560,7 +560,7 @@ def test_get_hmt_balance(self):
             "0x56B532F1D090E4edb1c92F30d3087771AE6B6992",
             get_w3(),
         )
-        print(amount)
+
         self.assertGreater(amount, 10000)