diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 38618ba2..dc8e8cf2 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -5,6 +5,7 @@ Brewtils Changelog ------ TBD +- Apply MD5 Check Sum of chunked files to ensure files are loaded into memory properly - Updated Plugin `max_concurrent` to support -1 to utilize the default formula that `concurrent.futures.ThreadPoolExecutor` supports `min(32, os.cpu_count() + 4)` - Updated SystemClient to utilize the local Garden name for default Namespace if none can be determined diff --git a/Makefile b/Makefile index fa242155..3cafce4a 100644 --- a/Makefile +++ b/Makefile @@ -81,7 +81,7 @@ clean: clean-build clean-docs clean-python clean-test ## remove everything but s # Formatting format: ## Run black formatter in-line - black --target-version py27 $(MODULE_NAME) $(TEST_DIR) + black $(MODULE_NAME) $(TEST_DIR) # Linting diff --git a/brewtils/models.py b/brewtils/models.py index 5c5aa020..06f1ffc9 100644 --- a/brewtils/models.py +++ b/brewtils/models.py @@ -520,6 +520,7 @@ def __init__( owner=None, job=None, request=None, + md5_sum=None, ): self.id = id self.owner_id = owner_id @@ -532,6 +533,7 @@ def __init__( self.file_size = file_size self.chunks = chunks self.chunk_size = chunk_size + self.md5_sum = md5_sum def __str__(self): return self.file_name @@ -593,6 +595,7 @@ def __init__( chunks_ok=None, operation_complete=None, message=None, + md5_sum=None, ): # Top-level file info self.file_id = file_id @@ -603,6 +606,7 @@ def __init__( self.chunks = chunks self.owner_id = owner_id self.owner_type = owner_type + self.md5_sum = md5_sum # Chunk info self.chunk_id = chunk_id self.offset = offset diff --git a/brewtils/rest/easy_client.py b/brewtils/rest/easy_client.py index 919f6673..eece1d78 100644 --- a/brewtils/rest/easy_client.py +++ b/brewtils/rest/easy_client.py @@ -4,6 +4,7 @@ from io import BytesIO from pathlib import Path from typing import Any, Callable, List, NoReturn, Optional, Type, Union +from hashlib import md5 import six import wrapt @@ -999,6 +1000,8 @@ def upload_chunked_file( fd = file_to_upload require_close = False + default_file_params["md5_sum"] = md5(fd.getbuffer()).hexdigest() + try: default_file_params["file_name"] = desired_filename or fd.name except AttributeError: @@ -1066,6 +1069,15 @@ def download_chunked_file(self, file_id): file_obj.seek(0) + if ( + "md5_sum" in meta + and meta["md5_sum"] != md5(file_obj.getbuffer()).hexdigest() + ): + raise ValidationError( + "Requested file %s MD5 SUM %s does match actual MD5 SUM %s" + % (file_id, meta["md5_sum"], md5(file_obj.getbuffer()).hexdigest()) + ) + return file_obj def delete_chunked_file(self, file_id): diff --git a/brewtils/schemas.py b/brewtils/schemas.py index 3d2eb385..ccaeee87 100644 --- a/brewtils/schemas.py +++ b/brewtils/schemas.py @@ -261,6 +261,7 @@ class FileSchema(BaseSchema): file_size = fields.Int(allow_none=False) chunks = fields.Dict(allow_none=True) chunk_size = fields.Int(allow_none=False) + md5_sum = fields.Str(allow_none=True) class FileChunkSchema(BaseSchema): @@ -281,6 +282,7 @@ class FileStatusSchema(BaseSchema): chunks = fields.Dict(allow_none=True) owner_id = fields.Str(allow_none=True) owner_type = fields.Str(allow_none=True) + md5_sum = fields.Str(allow_none=True) # Chunk info chunk_id = fields.Str(allow_none=True) offset = fields.Int(allow_none=True) diff --git a/test/rest/easy_client_test.py b/test/rest/easy_client_test.py index b58e990a..0f42ac33 100644 --- a/test/rest/easy_client_test.py +++ b/test/rest/easy_client_test.py @@ -39,6 +39,7 @@ def target_file(): fp.tell = Mock(return_value=0) fp.seek = Mock(return_value=1024) fp.read = Mock(side_effect=iter([b"content", None])) + fp.getbuffer = Mock(return_value=b"content") return fp @@ -596,6 +597,12 @@ def test_upload_chunked_file( client._check_chunked_file_validity = Mock(return_value=(True, {})) resolvable = client.upload_chunked_file(target_file, "desired_name") + _, called_kwargs = rest_client.post_chunked_file.call_args[0] + + assert called_kwargs["md5_sum"] == "9a0364b9e99bb480dd25e1f0284c8555" + assert called_kwargs["file_name"] == "desired_name" + assert called_kwargs["file_size"] == 1024 + assert called_kwargs["chunk_size"] == 261120 assert resolvable == bg_resolvable_chunk def test_upload_file_fail(self, client, rest_client, server_error, target_file):