Add SHA256 fallback if MD5 isn't available on the system

Author: nateprewitt

awscli/customizations/s3uploader.py

@@ -11,14 +11,14 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 
-import hashlib
 import logging
 import os
 import sys
 import threading
 
 import botocore
 import botocore.exceptions
+from botocore.compat import get_md5
 from s3transfer.manager import TransferManager
 from s3transfer.subscribers import BaseSubscriber
 
@@ -27,6 +27,16 @@
 LOG = logging.getLogger(__name__)
 
 
+def _get_checksum():
+    hashlib_params = {"usedforsecurity": False}
+    try:
+        checksum = get_md5(**hashlib_params)
+    except botocore.exceptions.MD5UnavailableError:
+        import hashlib
+        checksum = hashlib.sha256(**hashlib_params)
+    return checksum
+
+
 class NoSuchBucketError(Exception):
     def __init__(self, **kwargs):
         msg = self.fmt.format(**kwargs)
@@ -134,7 +144,7 @@ def upload(self, file_name, remote_path):
 
     def upload_with_dedup(self, file_name, extension=None):
         """
-        Makes and returns name of the S3 object based on the file's MD5 sum
+        Makes and returns name of the S3 object based on the file's checksum
 
         :param file_name: file to upload
         :param extension: String of file extension to append to the object
@@ -146,8 +156,8 @@ def upload_with_dedup(self, file_name, extension=None):
         # and re-upload only if necessary. So the template points to same file
         # in multiple places, this will upload only once
 
-        filemd5 = self.file_checksum(file_name)
-        remote_path = filemd5
+        file_checksum = self.file_checksum(file_name)
+        remote_path = file_checksum
         if extension:
             remote_path = remote_path + "." + extension
 
@@ -175,7 +185,7 @@ def make_url(self, obj_path):
 
     def file_checksum(self, file_name):
         with open(file_name, "rb") as file_handle:
-            md5 = hashlib.md5()
+            checksum = _get_checksum()
             # Read file in chunks of 4096 bytes
             block_size = 4096
 
@@ -185,13 +195,13 @@ def file_checksum(self, file_name):
 
             buf = file_handle.read(block_size)
             while len(buf) > 0:
-                md5.update(buf)
+                checksum.update(buf)
                 buf = file_handle.read(block_size)
 
             # Restore file cursor's position
             file_handle.seek(curpos)
 
-            return md5.hexdigest()
+            return checksum.hexdigest()
 
     def to_path_style_s3_url(self, key, version=None):
         """

tests/unit/customizations/test_s3uploader.py

@@ -348,6 +348,27 @@ def test_file_checksum(self):
         finally:
             shutil.rmtree(tempdir)
 
+    @mock.patch("awscli.customizations.s3uploader.get_md5")
+    def test_file_checksum_fips_fallback(self, get_md5_mock):
+        num_chars = 4096*5
+        data = ''.join(random.choice(string.ascii_uppercase)
+                       for _ in range(num_chars)).encode('utf-8')
+        checksum = hashlib.sha256(usedforsecurity=False)
+        checksum.update(data)
+        expected_checksum = checksum.hexdigest()
+
+        tempdir = tempfile.mkdtemp()
+        get_md5_mock.side_effect = botocore.exceptions.MD5UnavailableError()
+        try:
+            filename = os.path.join(tempdir, 'tempfile')
+            with open(filename, 'wb') as f:
+                f.write(data)
+
+            actual_checksum = self.s3uploader.file_checksum(filename)
+            self.assertEqual(expected_checksum, actual_checksum)
+        finally:
+            shutil.rmtree(tempdir)
+
     def test_make_url(self):
         path = "Hello/how/are/you"
         expected = f"s3://{self.bucket_name}/{path}"