Merge pull request #368 from pypa/hashing-and-testing

Refactor hash generation for testability

tests/test_package.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from __future__ import unicode_literals
+import platform
 from twine import package
 
 import pretend
@@ -163,3 +164,47 @@ def test_metadata_dictionary(gpg_signature):
 
     # GPG signature
     assert result.get('gpg_signature') == gpg_signature
+
+
+TWINE_1_5_0_WHEEL_HEXDIGEST = package.Hexdigest(
+    '1919f967e990bee7413e2a4bc35fd5d1',
+    'd86b0f33f0c7df49e888b11c43b417da5520cbdbce9f20618b1494b600061e67',
+    'b657a4148d05bd0098c1d6d8cc4e14e766dbe93c3a5ab6723b969da27a87bac0',
+)
+
+if platform.python_implementation().lower() == 'pypy':
+    # pyblake2 refuses to install on PyPy
+    TWINE_1_5_0_WHEEL_HEXDIGEST = TWINE_1_5_0_WHEEL_HEXDIGEST._replace(
+        blake2=None,
+    )
+
+
+def test_hash_manager():
+    """Verify our HashManager works."""
+    filename = 'tests/fixtures/twine-1.5.0-py2.py3-none-any.whl'
+    hasher = package.HashManager(filename)
+    hasher.hash()
+    assert hasher.hexdigest() == TWINE_1_5_0_WHEEL_HEXDIGEST
+
+
+def test_fips_hash_manager(monkeypatch):
+    """Verify the behaviour if hashlib is using FIPS mode."""
+    replaced_md5 = pretend.raiser(ValueError('fipsmode'))
+    monkeypatch.setattr(package.hashlib, 'md5', replaced_md5)
+
+    filename = 'tests/fixtures/twine-1.5.0-py2.py3-none-any.whl'
+    hasher = package.HashManager(filename)
+    hasher.hash()
+    hashes = TWINE_1_5_0_WHEEL_HEXDIGEST._replace(md5=None)
+    assert hasher.hexdigest() == hashes
+
+
+def test_no_blake2_hash_manager(monkeypatch):
+    """Verify the behaviour with missing blake2."""
+    monkeypatch.setattr(package, 'blake2b', None)
+
+    filename = 'tests/fixtures/twine-1.5.0-py2.py3-none-any.whl'
+    hasher = package.HashManager(filename)
+    hasher.hash()
+    hashes = TWINE_1_5_0_WHEEL_HEXDIGEST._replace(blake2=None)
+    assert hasher.hexdigest() == hashes

tox.ini

@@ -6,6 +6,7 @@ deps =
     coverage
     pretend
     pytest
+    py27,py34,py35: pyblake2
 commands =
     coverage run --source twine -m pytest {posargs:tests}
     coverage report -m

twine/package.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from __future__ import absolute_import, unicode_literals, print_function
+import collections
 import hashlib
 import io
 import os
@@ -61,33 +62,13 @@ def __init__(self, filename, comment, metadata, python_version, filetype):
         self.signed_basefilename = self.basefilename + '.asc'
         self.gpg_signature = None
 
-        blake2_256_hash = None
-        if blake2b is not None:
-            blake2_256_hash = blake2b(digest_size=256 // 8)
-        # NOTE(sigmavirus24): We may or may not be able to use blake2 so let's
-        # either use the methods or lambdas to do nothing.
-        blake_update = getattr(blake2_256_hash, 'update', lambda *args: None)
-        blake_hexdigest = getattr(blake2_256_hash, 'hexdigest', lambda: None)
-
-        # NOTE: MD5 is not available on FIPS-enabled systems
-        md5_hash = None
-        try:
-            md5_hash = hashlib.md5()
-        except ValueError:
-            md5_hash = None
-        md5_update = getattr(md5_hash, 'update', lambda *args: None)
-        md5_hexdigest = getattr(md5_hash, 'hexdigest', lambda: None)
-
-        sha2_hash = hashlib.sha256()
-        with open(filename, "rb") as fp:
-            for content in iter(lambda: fp.read(io.DEFAULT_BUFFER_SIZE), b''):
-                md5_update(content)
-                sha2_hash.update(content)
-                blake_update(content)
+        hasher = HashManager(filename)
+        hasher.hash()
+        hexdigest = hasher.hexdigest()
 
-        self.md5_digest = md5_hexdigest()
-        self.sha2_digest = sha2_hash.hexdigest()
-        self.blake2_256_digest = blake_hexdigest()
+        self.md5_digest = hexdigest.md5
+        self.sha2_digest = hexdigest.sha2
+        self.blake2_256_digest = hexdigest.blake2
 
     @classmethod
     def from_filename(cls, filename, comment):
@@ -184,3 +165,70 @@ def sign(self, sign_with, identity):
         subprocess.check_call(gpg_args)
 
         self.add_gpg_signature(self.signed_filename, self.signed_basefilename)
+
+
+Hexdigest = collections.namedtuple('Hexdigest', ['md5', 'sha2', 'blake2'])
+
+
+class HashManager(object):
+    """Manage our hashing objects for simplicity.
+
+    This will also allow us to better test this logic.
+    """
+
+    def __init__(self, filename):
+        """Initialize our manager and hasher objects."""
+        self.filename = filename
+        try:
+            self._md5_hasher = hashlib.md5()
+        except ValueError:
+            # FIPs mode disables MD5
+            self._md5_hasher = None
+
+        self._sha2_hasher = hashlib.sha256()
+        self._blake_hasher = None
+        if blake2b is not None:
+            self._blake_hasher = blake2b(digest_size=256 // 8)
+
+    def _md5_update(self, content):
+        if self._md5_hasher is not None:
+            self._md5_hasher.update(content)
+
+    def _md5_hexdigest(self):
+        if self._md5_hasher is not None:
+            return self._md5_hasher.hexdigest()
+        return None
+
+    def _sha2_update(self, content):
+        if self._sha2_hasher is not None:
+            self._sha2_hasher.update(content)
+
+    def _sha2_hexdigest(self):
+        if self._sha2_hasher is not None:
+            return self._sha2_hasher.hexdigest()
+        return None
+
+    def _blake_update(self, content):
+        if self._blake_hasher is not None:
+            self._blake_hasher.update(content)
+
+    def _blake_hexdigest(self):
+        if self._blake_hasher is not None:
+            return self._blake_hasher.hexdigest()
+        return None
+
+    def hash(self):
+        """Hash the file contents."""
+        with open(self.filename, "rb") as fp:
+            for content in iter(lambda: fp.read(io.DEFAULT_BUFFER_SIZE), b''):
+                self._md5_update(content)
+                self._sha2_update(content)
+                self._blake_update(content)
+
+    def hexdigest(self):
+        """Return the hexdigest for the file."""
+        return Hexdigest(
+            self._md5_hexdigest(),
+            self._sha2_hexdigest(),
+            self._blake_hexdigest(),
+        )