Hash: Generic Crypto.Hash.new(algo, [data]) function

This allows us to instantiate a new hash given only an existing hash
object.

2 files changed, 154 insertions, 0 deletions

diff --git a/lib/Crypto/Hash/__init__.py b/lib/Crypto/Hash/__init__.py
index b9f9525..1050c78 100644
--- a/lib/Crypto/Hash/__init__.py
+++ b/lib/Crypto/Hash/__init__.py
@@ -53,4 +53,123 @@ __all__ = ['HMAC', 'MD2', 'MD4', 'MD5', 'RIPEMD160', 'SHA1',
            'SHA224', 'SHA256', 'SHA384', 'SHA512']
 __revision__ = "$Id$"
 
+import sys
+if sys.version_info[0] == 2 and sys.version_info[1] == 1:
+    from Crypto.Util.py21compat import *
+from Crypto.Util.py3compat import *
 
+def new(algo, *args):
+    """Initialize a new hash object.
+
+    The first argument to this function may be an algorithm name or another
+    hash object.
+
+    This function has significant overhead.  It's recommended that you instead
+    import and use the individual hash modules directly.
+    """
+
+    # Try just invoking algo.new()
+    # We do this first so that this is the fastest.
+    try:
+        new_func = algo.new
+    except AttributeError:
+        pass
+    else:
+        return new_func(*args)
+
+    # Try getting the algorithm name.
+    if isinstance(algo, str):
+        name = algo
+    else:
+        try:
+            name = algo.name
+        except AttributeError:
+            raise ValueError("unsupported hash type %r" % (algo,))
+
+    # Got the name.  Let's see if we have a PyCrypto implementation.
+    try:
+        new_func = _new_funcs[name]
+    except KeyError:
+        # No PyCrypto implementation.  Try hashlib.
+        try:
+            import hashlib
+        except ImportError:
+            # There is no hashlib.
+            raise ValueError("unsupported hash type %s" % (name,))
+        return hashlib.new(name, *args)
+    else:
+        # We have a PyCrypto implementation.  Instantiate it.
+        return new_func(*args)
+
+# This dict originally gets the following _*_new methods, but its members get
+# replaced with the real new() methods of the various hash modules as they are
+# used.  We do it without locks to improve performance, which is safe in
+# CPython because dict access is atomic in CPython.  This might break PyPI.
+_new_funcs = {}
+
+def _md2_new(*args):
+    from Crypto.Hash import MD2
+    _new_funcs['MD2'] = _new_funcs['md2'] = MD2.new
+    return MD2.new(*args)
+_new_funcs['MD2'] = _new_funcs['md2'] = _md2_new
+del _md2_new
+
+def _md4_new(*args):
+    from Crypto.Hash import MD4
+    _new_funcs['MD4'] = _new_funcs['md4'] = MD4.new
+    return MD4.new(*args)
+_new_funcs['MD4'] = _new_funcs['md4'] = _md4_new
+del _md4_new
+
+def _md5_new(*args):
+    from Crypto.Hash import MD5
+    _new_funcs['MD5'] = _new_funcs['md5'] = MD5.new
+    return MD5.new(*args)
+_new_funcs['MD5'] = _new_funcs['md5'] = _md5_new
+del _md5_new
+
+def _ripemd160_new(*args):
+    from Crypto.Hash import RIPEMD160
+    _new_funcs['RIPEMD160'] = _new_funcs['ripemd160'] = \
+        _new_funcs['RIPEMD'] = _new_funcs['ripemd'] = RIPEMD160.new
+    return RIPEMD160.new(*args)
+_new_funcs['RIPEMD160'] = _new_funcs['ripemd160'] = \
+    _new_funcs['RIPEMD'] = _new_funcs['ripemd'] = _ripemd160_new
+del _ripemd160_new
+
+def _sha1_new(*args):
+    from Crypto.Hash import SHA1
+    _new_funcs['SHA1'] = _new_funcs['sha1'] = \
+        _new_funcs['SHA'] = _new_funcs['sha'] = SHA1.new
+    return SHA1.new(*args)
+_new_funcs['SHA1'] = _new_funcs['sha1'] = \
+    _new_funcs['SHA'] = _new_funcs['sha'] = _sha1_new
+del _sha1_new
+
+def _sha224_new(*args):
+    from Crypto.Hash import SHA224
+    _new_funcs['SHA224'] = _new_funcs['sha224'] = SHA224.new
+    return SHA224.new(*args)
+_new_funcs['SHA224'] = _new_funcs['sha224'] = _sha224_new
+del _sha224_new
+
+def _sha256_new(*args):
+    from Crypto.Hash import SHA256
+    _new_funcs['SHA256'] = _new_funcs['sha256'] = SHA256.new
+    return SHA256.new(*args)
+_new_funcs['SHA256'] = _new_funcs['sha256'] = _sha256_new
+del _sha256_new
+
+def _sha384_new(*args):
+    from Crypto.Hash import SHA384
+    _new_funcs['SHA384'] = _new_funcs['sha384'] = SHA384.new
+    return SHA384.new(*args)
+_new_funcs['SHA384'] = _new_funcs['sha384'] = _sha384_new
+del _sha384_new
+
+def _sha512_new(*args):
+    from Crypto.Hash import SHA512
+    _new_funcs['SHA512'] = _new_funcs['sha512'] = SHA512.new
+    return SHA512.new(*args)
+_new_funcs['SHA512'] = _new_funcs['sha512'] = _sha512_new
+del _sha512_new
diff --git a/lib/Crypto/SelfTest/Hash/common.py b/lib/Crypto/SelfTest/Hash/common.py
index fa9fd22..487e2cb 100644
--- a/lib/Crypto/SelfTest/Hash/common.py
+++ b/lib/Crypto/SelfTest/Hash/common.py
@@ -29,6 +29,7 @@ __revision__ = "$Id$"
 import sys
 import unittest
 import binascii
+import Crypto.Hash
 from Crypto.Util.py3compat import *
 
 # For compatibility with Python 2.1 and Python 2.2
@@ -100,6 +101,19 @@ class HashSelfTest(unittest.TestCase):
         out5 = binascii.b2a_hex(h2.digest())
         self.assertEqual(self.expected, out5)
 
+        # Verify that Crypto.Hash.new(h) produces a fresh hash object
+        h3 = Crypto.Hash.new(h)
+        h3.update(self.input)
+        out6 = binascii.b2a_hex(h3.digest())
+        self.assertEqual(self.expected, out6)
+
+        if hasattr(h, 'name'):
+            # Verify that Crypto.Hash.new(h.name) produces a fresh hash object
+            h4 = Crypto.Hash.new(h.name)
+            h4.update(self.input)
+            out7 = binascii.b2a_hex(h4.digest())
+            self.assertEqual(self.expected, out7)
+
 class HashTestOID(unittest.TestCase):
     def __init__(self, hashmod, oid):
         unittest.TestCase.__init__(self)
@@ -111,6 +125,25 @@ class HashTestOID(unittest.TestCase):
         h = self.hashmod.new()
         self.assertEqual(PKCS1_v1_5._HASH_OIDS[h.name], self.oid)
 
+class GenericHashConstructorTest(unittest.TestCase):
+    def __init__(self, hashmod):
+        unittest.TestCase.__init__(self)
+        self.hashmod = hashmod
+
+    def runTest(self):
+        obj1 = self.hashmod.new("foo")
+        obj2 = self.hashmod.new()
+        obj3 = Crypto.Hash.new(obj1.name, "foo")
+        obj4 = Crypto.Hash.new(obj1.name)
+        obj5 = Crypto.Hash.new(obj1, "foo")
+        obj6 = Crypto.Hash.new(obj1)
+        self.assert_(isinstance(self.hashmod, obj1))
+        self.assert_(isinstance(self.hashmod, obj2))
+        self.assert_(isinstance(self.hashmod, obj3))
+        self.assert_(isinstance(self.hashmod, obj4))
+        self.assert_(isinstance(self.hashmod, obj5))
+        self.assert_(isinstance(self.hashmod, obj6))
+
 class MACSelfTest(unittest.TestCase):
 
     def __init__(self, hashmod, description, expected_dict, input, key, hashmods):
@@ -175,6 +208,8 @@ def make_hash_tests(module, module_name, test_data, digest_size, oid=None):
     tests.append(HashDigestSizeSelfTest(module, name, digest_size))
     if oid is not None:
         tests.append(HashTestOID(module, b(oid)))
+    if getattr(module, 'name', None) is not None:
+        tests.append(GenericHashConstructorTest(module))
     return tests
 
 def make_mac_tests(module, module_name, test_data, hashmods):