MD5 Check Sum generate and check

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
 

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

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

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):

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)

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):