bpo-40698: Improve distutils upload hash digests (GH-20260)

- Fix upload test on systems that blocks MD5
- Add SHA2-256 and Blake2b-256 digests based on new Warehous and twine
  specs.

Signed-off-by: Christian Heimes <christian@python.org>

2 files changed, 41 insertions, 5 deletions

diff --git a/command/upload.py b/command/upload.py
index d822ba01..95e9fda1 100644
--- a/command/upload.py
+++ b/command/upload.py
@@ -16,6 +16,16 @@ from distutils.core import PyPIRCCommand
 from distutils.spawn import spawn
 from distutils import log
 
+
+# PyPI Warehouse supports MD5, SHA256, and Blake2 (blake2-256)
+# https://bugs.python.org/issue40698
+_FILE_CONTENT_DIGESTS = {
+    "md5_digest": getattr(hashlib, "md5", None),
+    "sha256_digest": getattr(hashlib, "sha256", None),
+    "blake2_256_digest": getattr(hashlib, "blake2b", None),
+}
+
+
 class upload(PyPIRCCommand):
 
     description = "upload binary package to PyPI"
@@ -87,6 +97,7 @@ class upload(PyPIRCCommand):
             content = f.read()
         finally:
             f.close()
+
         meta = self.distribution.metadata
         data = {
             # action
@@ -101,7 +112,6 @@ class upload(PyPIRCCommand):
             'content': (os.path.basename(filename),content),
             'filetype': command,
             'pyversion': pyversion,
-            'md5_digest': hashlib.md5(content).hexdigest(),
 
             # additional meta-data
             'metadata_version': '1.0',
@@ -123,6 +133,16 @@ class upload(PyPIRCCommand):
 
         data['comment'] = ''
 
+        # file content digests
+        for digest_name, digest_cons in _FILE_CONTENT_DIGESTS.items():
+            if digest_cons is None:
+                continue
+            try:
+                data[digest_name] = digest_cons(content).hexdigest()
+            except ValueError:
+                # hash digest not available or blocked by security policy
+                pass
+
         if self.sign:
             with open(filename + ".asc", "rb") as f:
                 data['gpg_signature'] = (os.path.basename(filename) + ".asc",
diff --git a/tests/test_upload.py b/tests/test_upload.py
index c17d8e7d..bca5516d 100644
--- a/tests/test_upload.py
+++ b/tests/test_upload.py
@@ -130,14 +130,30 @@ class uploadTestCase(BasePyPIRCCommandTestCase):
 
         # what did we send ?
         headers = dict(self.last_open.req.headers)
-        self.assertEqual(headers['Content-length'], '2162')
+        self.assertGreaterEqual(int(headers['Content-length']), 2162)
         content_type = headers['Content-type']
         self.assertTrue(content_type.startswith('multipart/form-data'))
         self.assertEqual(self.last_open.req.get_method(), 'POST')
         expected_url = 'https://upload.pypi.org/legacy/'
         self.assertEqual(self.last_open.req.get_full_url(), expected_url)
-        self.assertTrue(b'xxx' in self.last_open.req.data)
-        self.assertIn(b'protocol_version', self.last_open.req.data)
+        data = self.last_open.req.data
+        self.assertIn(b'xxx',data)
+        self.assertIn(b'protocol_version', data)
+        self.assertIn(b'sha256_digest', data)
+        self.assertIn(
+            b'cd2eb0837c9b4c962c22d2ff8b5441b7b45805887f051d39bf133b583baf'
+            b'6860',
+            data
+        )
+        if b'md5_digest' in data:
+            self.assertIn(b'f561aaf6ef0bf14d4208bb46a4ccb3ad', data)
+        if b'blake2_256_digest' in data:
+            self.assertIn(
+                b'b6f289a27d4fe90da63c503bfe0a9b761a8f76bb86148565065f040be'
+                b'6d1c3044cf7ded78ef800509bccb4b648e507d88dc6383d67642aadcc'
+                b'ce443f1534330a',
+                data
+            )
 
         # The PyPI response body was echoed
         results = self.get_logs(INFO)
@@ -166,7 +182,7 @@ class uploadTestCase(BasePyPIRCCommandTestCase):
         cmd.run()
 
         headers = dict(self.last_open.req.headers)
-        self.assertEqual(headers['Content-length'], '2172')
+        self.assertGreaterEqual(int(headers['Content-length']), 2172)
         self.assertIn(b'long description\r', self.last_open.req.data)
 
     def test_upload_fails(self):