1 files changed, 83 insertions, 14 deletions

diff --git a/Lib/hashlib.py b/Lib/hashlib.py
index 21454c7..316cece 100644
--- a/Lib/hashlib.py
+++ b/Lib/hashlib.py
@@ -1,4 +1,4 @@
-#  Copyright (C) 2005-2010   Gregory P. Smith (greg@krypto.org)
+#.  Copyright (C) 2005-2010   Gregory P. Smith (greg@krypto.org)
 #  Licensed to PSF under a Contributor Agreement.
 #
 
@@ -60,34 +60,38 @@ algorithms_guaranteed = set(__always_supported)
 algorithms_available = set(__always_supported)
 
 __all__ = __always_supported + ('new', 'algorithms_guaranteed',
-                                'algorithms_available')
+                                'algorithms_available', 'pbkdf2_hmac')
 
 
+__builtin_constructor_cache = {}
+
 def __get_builtin_constructor(name):
+    cache = __builtin_constructor_cache
+    constructor = cache.get(name)
+    if constructor is not None:
+        return constructor
     try:
         if name in ('SHA1', 'sha1'):
             import _sha1
-            return _sha1.sha1
+            cache['SHA1'] = cache['sha1'] = _sha1.sha1
         elif name in ('MD5', 'md5'):
             import _md5
-            return _md5.md5
+            cache['MD5'] = cache['md5'] = _md5.md5
         elif name in ('SHA256', 'sha256', 'SHA224', 'sha224'):
             import _sha256
-            bs = name[3:]
-            if bs == '256':
-                return _sha256.sha256
-            elif bs == '224':
-                return _sha256.sha224
+            cache['SHA224'] = cache['sha224'] = _sha256.sha224
+            cache['SHA256'] = cache['sha256'] = _sha256.sha256
         elif name in ('SHA512', 'sha512', 'SHA384', 'sha384'):
             import _sha512
-            bs = name[3:]
-            if bs == '512':
-                return _sha512.sha512
-            elif bs == '384':
-                return _sha512.sha384
+            cache['SHA384'] = cache['sha384'] = _sha512.sha384
+            cache['SHA512'] = cache['sha512'] = _sha512.sha512
     except ImportError:
         pass  # no extension module, this hash is unsupported.
 
+    constructor = cache.get(name)
+    if constructor is not None:
+        return constructor
+
     raise ValueError('unsupported hash type ' + name)
 
 
@@ -134,6 +138,71 @@ except ImportError:
     new = __py_new
     __get_hash = __get_builtin_constructor
 
+try:
+    # OpenSSL's PKCS5_PBKDF2_HMAC requires OpenSSL 1.0+ with HMAC and SHA
+    from _hashlib import pbkdf2_hmac
+except ImportError:
+    _trans_5C = bytes((x ^ 0x5C) for x in range(256))
+    _trans_36 = bytes((x ^ 0x36) for x in range(256))
+
+    def pbkdf2_hmac(hash_name, password, salt, iterations, dklen=None):
+        """Password based key derivation function 2 (PKCS #5 v2.0)
+
+        This Python implementations based on the hmac module about as fast
+        as OpenSSL's PKCS5_PBKDF2_HMAC for short passwords and much faster
+        for long passwords.
+        """
+        if not isinstance(hash_name, str):
+            raise TypeError(hash_name)
+
+        if not isinstance(password, (bytes, bytearray)):
+            password = bytes(memoryview(password))
+        if not isinstance(salt, (bytes, bytearray)):
+            salt = bytes(memoryview(salt))
+
+        # Fast inline HMAC implementation
+        inner = new(hash_name)
+        outer = new(hash_name)
+        blocksize = getattr(inner, 'block_size', 64)
+        if len(password) > blocksize:
+            password = new(hash_name, password).digest()
+        password = password + b'\x00' * (blocksize - len(password))
+        inner.update(password.translate(_trans_36))
+        outer.update(password.translate(_trans_5C))
+
+        def prf(msg, inner=inner, outer=outer):
+            # PBKDF2_HMAC uses the password as key. We can re-use the same
+            # digest objects and just update copies to skip initialization.
+            icpy = inner.copy()
+            ocpy = outer.copy()
+            icpy.update(msg)
+            ocpy.update(icpy.digest())
+            return ocpy.digest()
+
+        if iterations < 1:
+            raise ValueError(iterations)
+        if dklen is None:
+            dklen = outer.digest_size
+        if dklen < 1:
+            raise ValueError(dklen)
+
+        dkey = b''
+        loop = 1
+        from_bytes = int.from_bytes
+        while len(dkey) < dklen:
+            prev = prf(salt + loop.to_bytes(4, 'big'))
+            # endianess doesn't matter here as long to / from use the same
+            rkey = int.from_bytes(prev, 'big')
+            for i in range(iterations - 1):
+                prev = prf(prev)
+                # rkey = rkey ^ prev
+                rkey ^= from_bytes(prev, 'big')
+            loop += 1
+            dkey += rkey.to_bytes(inner.digest_size, 'big')
+
+        return dkey[:dklen]
+
+
 for __func_name in __always_supported:
     # try them all, some may not work due to the OpenSSL
     # version not supporting that algorithm.