Repository: planetlabs/planet-client-python
Pull request #474 (2022-05-03 14:05:56)

Checksum match orders #432


Patch:
diff --git a/planet/cli/orders.py b/planet/cli/orders.py
index 764530b..65353f4 100644
--- a/planet/cli/orders.py
+++ b/planet/cli/orders.py
@@ -179,6 +179,10 @@ async def wait(ctx, order_id, delay, max_attempts, state):
 @translate_exceptions
 @coro
 @click.argument('order_id', type=click.UUID)
+@click.option('--checksum',
+              default=None,
+              type=click.Choice(['MD5', 'SHA256'], case_sensitive=False),
+              help=('Verify that checksums match.'))
 @click.option('--directory',
               default='.',
               help=('Base directory for file download.'),
@@ -190,14 +194,18 @@ async def wait(ctx, order_id, delay, max_attempts, state):
               is_flag=True,
               default=False,
               help=('Overwrite files if they already exist.'))
-async def download(ctx, order_id, overwrite, directory):
-    """Download order by order ID."""
+async def download(ctx, order_id, overwrite, directory, checksum):
+    """Download order by order ID.
+
+If --checksum is provided, the associated checksums given in the manifest
+are compared against the downloaded files to verify that they match."""
     quiet = ctx.obj['QUIET']
     async with orders_client(ctx) as cl:
         await cl.download_order(str(order_id),
                                 directory=directory,
                                 overwrite=overwrite,
-                                progress_bar=not quiet)
+                                progress_bar=not quiet,
+                                checksum=checksum)
 
 
 def read_file_geojson(ctx, param, value):
diff --git a/planet/clients/orders.py b/planet/clients/orders.py
index 311711f..fee7019 100644
--- a/planet/clients/orders.py
+++ b/planet/clients/orders.py
@@ -18,7 +18,10 @@ import os
 import time
 import typing
 import uuid
+import json
+import hashlib
 
+from pathlib import Path
 from .. import exceptions
 from ..constants import PLANET_BASE_URL
 from ..http import Session
@@ -271,18 +274,68 @@ class OrdersClient():
                              progress_bar=progress_bar)
         return dl_path
 
+    @staticmethod
+    def _validate_checksum(manifest_data: dict, filenames: list,
+                           checksum: str):
+        """Calculate checksum and validate that it passes.
+
+        Parameters:
+            manifest_data: The contents of the opened manifest.json file.
+            filenames: List of downloaded assets.
+            checksum: The type of checksum hash- 'MD5' or 'SHA256'.
+
+        Returns:
+            Nothing if checksums were succesful.
+
+        Raises:
+            planet.exceptions.ClientError: If the checksum fails."""
+
+        # Save hashkey & respective filename from manifest json into dictionary
+        # This is the origin hashkey
+        if checksum == 'MD5':
+            hash_type = hashlib.md5
+            checksum = checksum.lower()
+        elif checksum == 'SHA256':
+            hash_type = hashlib.sha256
+            checksum = checksum.lower()
+        file_key_pairs = {}
+        for json_entry in manifest_data['files']:
+            file_name = Path(json_entry['path']).name
+            origin_hash = json_entry['digests'][checksum]
+            file_key_pairs[file_name] = origin_hash
+        # For each downloaded file, retrieve origin hashkey from dict
+        filenames_loop = [
+            x for x in filenames if not x.endswith('manifest.json')
+        ]
+        for filename in filenames_loop:
+            downloaded_file_name = Path(filename).name
+            origin_hash = file_key_pairs[downloaded_file_name]
+            # For each file (not including manifest json),
+            # calculate hash on contents. This is the returned hash.
+            with open(filename, 'rb') as file_to_check:
+                json_data = file_to_check.read()
+                returned_hash = hash_type(json_data).hexdigest()
+                # Compare original hashkey in dict with calculated
+                if origin_hash != returned_hash:
+                    raise exceptions.ClientError(
+                        'Checksum failed. File ({filename}) not correctly \
+                        downloaded.')
+        return None
+
     async def download_order(self,
                              order_id: str,
                              directory: str = None,
                              overwrite: bool = False,
-                             progress_bar: bool = False) -> typing.List[str]:
+                             progress_bar: bool = False,
+                             checksum: str = None) -> typing.List[str]:
         """Download all assets in an order.
 
         Parameters:
-            order_id: The ID of the order
-            directory: Root directory for file download.
+            order_id: The ID of the order.
+            directory: Base directory for file download.
             overwrite: Overwrite files if they already exist.
             progress_bar: Show progress bar during download.
+            checksum: The type of checksum hash- MD5 or SHA256.
 
         Returns:
             Paths to downloaded files.
@@ -291,9 +344,11 @@ class OrdersClient():
             planet.exceptions.APIError: On API error.
             planet.exceptions.ClientError: If the order is not in a final
                 state.
+            planet.exceptions.ClientError: If more than one manifest file
+                per order.
+            planet.exceptions.ClientError: If no manifest file found in order.
         """
         order = await self.get_order(order_id)
-
         order_state = order['state']
         if not OrderStates.is_final(order_state):
             raise exceptions.ClientError(
@@ -313,6 +368,24 @@ class OrdersClient():
                                       progress_bar=progress_bar)
             for location in locations
         ]
+        if checksum:
+            # Checksum Implementation
+            manifest_files = [
+                x for x in filenames if Path(x).name == 'manifest.json'
+            ]
+            manifest_count = len(manifest_files)
+            if manifest_count > 1:
+                raise exceptions.ClientError(
+                    'Only 1 manifest.json expected per order.\
+                                             Recieved: {manifest_count}')
+            elif manifest_count is None:
+                raise exceptions.ClientError('No manifest file found.')
+            manifest_json = manifest_files[0]
+            with open(manifest_json, 'rb') as manifest:
+                manifest_data = json.load(manifest)
+                self._validate_checksum(manifest_data=manifest_data,
+                                        filenames=filenames,
+                                        checksum=checksum)
         return filenames
 
     @staticmethod


Test patch:
diff --git a/tests/integration/test_orders_api.py b/tests/integration/test_orders_api.py
index c542ece..6643f85 100644
--- a/tests/integration/test_orders_api.py
+++ b/tests/integration/test_orders_api.py
@@ -19,6 +19,7 @@ import math
 import os
 from pathlib import Path
 from unittest.mock import call, create_autospec
+import hashlib
 
 import httpx
 import pytest
@@ -536,6 +537,171 @@ async def test_download_asset_md(tmpdir, session):
     assert Path(filename).name == 'metadata.json'
 
 
+@respx.mock
+@pytest.mark.asyncio
+@pytest.mark.parametrize("checksum", [("MD5"), ("SHA256")])
+async def test_checksum_success(tmpdir,
+                                order_description,
+                                oid,
+                                session,
+                                checksum):
+    # Mock an HTTP response for download
+    order_description['state'] = 'success'
+    dl_url1 = TEST_DOWNLOAD_URL + '/asset1'
+    dl_url2 = TEST_DOWNLOAD_URL + '/asset2'
+    dl_url3 = TEST_DOWNLOAD_URL + '/manifest'
+    order_description['_links']['results'] = [{
+        'location': dl_url1
+    }, {
+        'location': dl_url2
+    }, {
+        'location': dl_url3
+    }]
+    get_url = f'{TEST_ORDERS_URL}/{oid}'
+    get_order_response = httpx.Response(HTTPStatus.OK, json=order_description)
+    respx.get(get_url).return_value = get_order_response
+
+    # The first asset in the order:
+    # A 1-byte binary file containing the value "1".
+    asset1_content = b"1"
+    asset1_response = httpx.Response(HTTPStatus.OK,
+                                     content=asset1_content,
+                                     headers={
+                                         'Content-Type':
+                                         'image/tiff',
+                                         'Content-Disposition':
+                                         'attachment; filename="asset1.tif"'
+                                     })
+    respx.get(dl_url1).return_value = asset1_response
+
+    asset2_response = httpx.Response(HTTPStatus.OK,
+                                     json={'foo': 'bar'},
+                                     headers={
+                                         'Content-Type':
+                                         'application/json',
+                                         'Content-Disposition':
+                                         'attachment; filename="asset2.json"'
+                                     })
+    respx.get(dl_url2).return_value = asset2_response
+    # Create a mock manifest which has the correct hashkeys for each asset
+    manifest = httpx.Response(
+        HTTPStatus.OK,
+        json={
+            "name":
+            "",
+            "files":
+            [{
+                "path": "asset1.tif",
+                "digests": {
+                    "md5": hashlib.md5(asset1_response.content).hexdigest(),
+                    "sha256":
+                    hashlib.sha256(asset1_response.content).hexdigest()
+                },
+            },
+             {
+                 "path": "asset2.json",
+                 "digests": {
+                     "md5": hashlib.md5(asset2_response.content).hexdigest(),
+                     "sha256":
+                     hashlib.sha256(asset2_response.content).hexdigest()
+                 }
+             }]
+        },
+        headers={
+            'Content-Type': 'application/json',
+            'Content-Disposition': 'attachment; filename="manifest.json"'
+        })
+    respx.get(dl_url3).return_value = manifest
+    # Test Checksum
+    cl = OrdersClient(session, base_url=TEST_URL)
+    await cl.download_order(oid, directory=str(tmpdir), checksum=checksum)
+
+
+@respx.mock
+@pytest.mark.asyncio
+@pytest.mark.parametrize("checksum", [("MD5"), ("SHA256")])
+async def test_checksum_failure(tmpdir,
+                                order_description,
+                                oid,
+                                session,
+                                checksum):
+    # Note: the hashkeys in the mock manifest below were changed
+    # from the correct keys to temporary keys
+    # This should cause the checksum to fail.
+    # Mock an HTTP response for download
+    order_description['state'] = 'success'
+    dl_url1 = TEST_DOWNLOAD_URL + '/asset1'
+    dl_url2 = TEST_DOWNLOAD_URL + '/asset2'
+    dl_url3 = TEST_DOWNLOAD_URL + '/manifest'
+    order_description['_links']['results'] = [{
+        'location': dl_url1
+    }, {
+        'location': dl_url2
+    }, {
+        'location': dl_url3
+    }]
+    get_url = f'{TEST_ORDERS_URL}/{oid}'
+    get_order_response = httpx.Response(HTTPStatus.OK, json=order_description)
+    respx.get(get_url).return_value = get_order_response
+
+    # The first asset in the order:
+    # A 1-byte binary file containing the value "1".
+    asset1_content = b"1"
+    asset1_response = httpx.Response(HTTPStatus.OK,
+                                     content=asset1_content,
+                                     headers={
+                                         'Content-Type':
+                                         'image/tiff',
+                                         'Content-Disposition':
+                                         'attachment; filename="asset1.tif"'
+                                     })
+    respx.get(dl_url1).return_value = asset1_response
+
+    # The second asset in the order is a json file
+    asset2_response = httpx.Response(HTTPStatus.OK,
+                                     json={'foo': 'bar'},
+                                     headers={
+                                         'Content-Type':
+                                         'application/json',
+                                         'Content-Disposition':
+                                         'attachment; filename="asset2.json"'
+                                     })
+    respx.get(dl_url2).return_value = asset2_response
+
+    manifest = httpx.Response(HTTPStatus.OK,
+                              json={
+                                  "name":
+                                  "",
+                                  "files": [{
+                                      "path": "asset1.tif",
+                                      "digests": {
+                                          "md5": "manifest_key1_md5",
+                                          "sha256": "manifest_key1_sha256"
+                                      },
+                                  },
+                                            {
+                                                "path": "asset2.json",
+                                                "digests": {
+                                                    "md5":
+                                                    "manifest_key2_md5",
+                                                    "sha256":
+                                                    "manifest_key2_sha256"
+                                                }
+                                            }]
+                              },
+                              headers={
+                                  'Content-Type':
+                                  'application/json',
+                                  'Content-Disposition':
+                                  'attachment; filename="manifest.json"'
+                              })
+    respx.get(dl_url3).return_value = manifest
+    # Test Checksum
+    cl = OrdersClient(session, base_url=TEST_URL)
+    with pytest.raises(exceptions.ClientError):
+        await cl.download_order(oid, directory=str(tmpdir), checksum=checksum)
+
+
 @respx.mock
 @pytest.mark.asyncio
 async def test_download_asset_img(tmpdir, open_test_img, session):