kdf enhancements

* renamed _resolve_prf() to get_prf(), documented interface
* added pbkdf1 support
* added 'encoding' kwd to pbkdf1 & pbkdf2, for easier handling of unicode inputs
*

1 files changed, 185 insertions, 53 deletions

diff --git a/passlib/utils/pbkdf2.py b/passlib/utils/pbkdf2.py
index 6e0222a..adb071f 100644
--- a/passlib/utils/pbkdf2.py
+++ b/passlib/utils/pbkdf2.py
@@ -1,5 +1,7 @@
 """passlib.pbkdf2 - PBKDF2 support
 
+this module is getting increasingly poorly named.
+maybe rename to "kdf" since it's getting more key derivation functions added.
 """
 #=================================================================================
 #imports
@@ -23,11 +25,13 @@ from passlib.utils import xor_bytes
 #local
 __all__ = [
     "hmac_sha1",
+    "get_prf",
+    "pbkdf1",
     "pbkdf2",
 ]
 
 #=================================================================================
-#hmac sha1 support
+#quick hmac_sha1 implementation used various places
 #=================================================================================
 def hmac_sha1(key, msg):
     "perform raw hmac-sha1 of a message"
@@ -45,85 +49,211 @@ if _EVP:
             hmac_sha1 = _EVP.hmac
 
 #=================================================================================
-#backend
+#general prf lookup
 #=================================================================================
-MAX_BLOCKS = 0xffffffff #2**32-1
-MAX_HMAC_SHA1_KEYLEN = MAX_BLOCKS*20
+def _get_hmac_prf(digest):
+    "helper to return HMAC prf for specific digest"
+    #check if m2crypto is present and supports requested digest
+    if _EVP:
+        try:
+            result = _EVP.hmac('x', 'y', digest)
+        except ValueError:
+            pass
+        else:
+            #it does. so use M2Crypto's hmac & digest code
+            hmac_const = _EVP.hmac
+            def prf(key, msg):
+                "prf(key,msg)->digest; generated by passlib.utils.pbkdf2.get_prf()"
+                return hmac_const(key, msg, digest)
+            prf.__name__ = "hmac_" + digest
+            digest_size = len(result)
+            return prf, digest_size
 
-def _resolve_prf(prf):
-    "resolve prf string or callable -> func & digest_size"
-    if isinstance(prf, str):
-        if prf.startswith("hmac-"):
-            digest = prf[5:]
-
-            #check if m2crypto is present and supports requested digest
-            if _EVP:
-                try:
-                    result = _EVP.hmac('x', 'y', digest)
-                except ValueError:
-                    pass
-                else:
-                    #it does. so use M2Crypto's hmac & digest code
-                    hmac_const = _EVP.hmac
-                    def encode_block(key, msg):
-                        return hmac_const(key, msg, digest)
-                    digest_size = len(result)
-                    return encode_block, digest_size
-
-            #fall back to stdlib implementation
-            digest_const = getattr(hashlib, digest, None)
-            if not digest_const:
-                raise ValueError("unknown hash algorithm: %r" % (digest,))
-            digest_size = digest_const().digest_size
-            hmac_const = hmac.new
-            def encode_block(key, msg):
-                return hmac_const(key, msg, digest_const).digest()
-            return encode_block, digest_size
+    #fall back to stdlib implementation
+    digest_const = getattr(hashlib, digest, None)
+    if not digest_const:
+        raise ValueError("unknown hash algorithm: %r" % (digest,))
+    digest_size = digest_const().digest_size
+    hmac_const = hmac.new
+    def prf(key, msg):
+        "prf(key,msg)->digest; generated by passlib.utils.pbkdf2.get_prf()"
+        return hmac_const(key, msg, digest_const).digest()
+    prf.__name__ = "hmac_" + digest
+    return prf, digest_size
 
-        else:
-            raise ValueError("unknown prf algorithm: %r" % (prf,))
+#cache mapping prf name/func -> (func, digest_size)
+_prf_cache = {}
 
-    elif callable(prf):
-        #assume it's a callable, use it directly
-        digest_size = len(prf('',''))
-        return prf, digest_size
+def _clear_prf_cache():
+    "helper for unit tests"
+    _prf_cache.clear()
 
+def get_prf(name):
+    """lookup pseudo-random family (prf) by name.
+    
+    :arg name:
+        this must be the name of a recognized prf.
+        currently this only recognizes names with the format
+        :samp:`hmac-{digest}`, where :samp:`{digest}`
+        is the name of a hash function such as
+        ``md5``, ``sha256``, etc.
+        
+        this can also be a callable with the signature
+        ``prf(secret, message) -> digest``,
+        in which case it will be returned unchanged.
+        
+    :raises ValueError: if the name is not known
+    :raises TypeError: if the name is not a callable or string
+    
+    :returns:
+        a tuple of :samp:`({func}, {digest_size})`.
+        
+        * :samp:`{func}` is a function implementing
+          the specified prf, and has the signature
+          ``func(secret, message) -> digest``.
+          
+        * :samp:`{digest_size}` is an integer indicating
+          the number of bytes the function returns.
+          
+    usage example::
+    
+        >>> from passlib.utils.pbkdf2 import get_prf
+        >>> hmac_sha256, dsize = get_prf("hmac-sha256")
+        >>> hmac_sha256
+        <function hmac_sha256 at 0x1e37c80>
+        >>> dsize
+        32
+        >>> digest = hmac_sha256('password', 'message')
+
+    this function will attempt to return the fastest implementation
+    it can find; if M2Crypto is present, and supports the specified prf, 
+    :func:`M2Crypto.EVP.hmac` will be used behind the scenes.
+    """
+    global _prf_cache
+    if name in _prf_cache:
+        return _prf_cache[name]
+    if isinstance(name, str):
+        if name.startswith("hmac-") or name.startswith("hmac_"):
+            retval = _get_hmac_prf(name[5:])
+        else:
+            raise ValueError("unknown prf algorithm: %r" % (name,))
+    elif callable(name):
+        #assume it's a callable, use it directly
+        digest_size = len(name('x','y'))
+        retval = (name, digest_size)
     else:
         raise TypeError("prf must be string or callable")
+    _prf_cache[name] = retval
+    return retval
 
-def pbkdf2(secret, salt, rounds, keylen, prf="hmac-sha1"):
-    """pkcs#5 password-based key derivation v2.0
+#=================================================================================
+#pbkdf1 support
+#=================================================================================
+def pbkdf1(secret, salt, rounds, keylen, hash="sha1", encoding="utf8"):
+    """pkcs#5 password-based key derivation v1.5
 
     :arg secret: passphrase to use to generate key
     :arg salt: salt string to use when generating key
     :param rounds: number of rounds to use to generate key
     :arg keylen: number of bytes to generate
-    :param prf:
-        psuedo-random function to use for key strengthening.
+    :param hash:
+        hash function to use.
         if specified, it must be one of the following:
 
-        * a callable with the prototype ``prf(secret, plaintext) -> ciphertext``
-        * a string of the form :samp:`hmac-{digest}`, where :samp:`{digest}`
-          is the name of a hash function such as ``md5``, ``sha256``, etc.
+        * a callable with the prototype ``hash(message) -> raw digest``
+        * a string matching one of the hashes recognized by hashlib
+
+    :param encoding:
+        encoding to use when converting unicode secret and salt to bytes.
+        defaults to ``utf8``.
+        
+    :returns:
+        raw bytes of generated key
+    
+    This algorithm is deprecated, new code should use PBKDF2.
+    Among other reasons, ``keylen`` cannot be larger
+    than the digest size of the specified hash.
+    
+    """
+    #prepare secret
+    if isinstance(secret, unicode):
+        secret = secret.encode(encoding)
+    elif not isinstance(secret, str):
+        raise TypeError("secret must be str or unicode")
+
+    #prepare salt
+    if isinstance(salt, unicode):
+        salt = salt.encode(encoding)
+    elif not isinstance(salt, str):
+        raise TypeError("salt must be str or unicode")
 
-        this defaults to ``hmac-sha1``, the only prf listed in the PBKDF2 specification.
+    #preprare rounds
+    if not isinstance(rounds, (int, long)):
+        raise TypeError("rounds must be an integer")
+    if rounds < 1:
+        raise ValueError("rounds must be at least 1")
+
+    #prep keylen
+    if keylen < 0:
+        raise ValueError("keylen must be at least 0")
 
-    If M2Crypto is present, and supports the specified prf, 
-    :func:`M2Crypto.EVP.hmac` will be used to accelerate this function.
+    #resolve hash
+    if isinstance(hash, str):
+        #check for builtin hash
+        hf = getattr(hashlib, hash, None)
+        if hf is None:
+            #check for ssl hash
+            #NOTE: if hash unknown, will throw ValueError, which we'd just
+            # reraise anyways; so instead of checking, we just let it get
+            # thrown during first use, below
+            def hf(msg):
+                return hashlib.new(hash, msg)
 
+    #run pbkdf1
+    block = hf(secret + salt).digest()
+    if keylen > len(block):
+        raise ValueError, "keylength too large for digest: %r > %r" % (keylen, len(block))
+    r = 1
+    while r < rounds:
+        block = hf(block).digest()
+        r += 1
+    return block[:keylen]
+    
+#=================================================================================
+#pbkdf2
+#=================================================================================
+MAX_BLOCKS = 0xffffffff #2**32-1
+MAX_HMAC_SHA1_KEYLEN = MAX_BLOCKS*20
+
+def pbkdf2(secret, salt, rounds, keylen, prf="hmac-sha1", encoding="utf8"):
+    """pkcs#5 password-based key derivation v2.0
+
+    :arg secret: passphrase to use to generate key
+    :arg salt: salt string to use when generating key
+    :param rounds: number of rounds to use to generate key
+    :arg keylen: number of bytes to generate
+    :param prf:
+        psuedo-random family to use for key strengthening.
+        this can be any string or callable accepted by :func:`get_prf`.
+        this defaults to ``hmac-sha1`` (the only prf explicitly listed in
+        the PBKDF2 specification)
+    :param encoding:
+        encoding to use when converting unicode secret and salt to bytes.
+        defaults to ``utf8``.
+        
     :returns:
         raw bytes of generated key
     """
 
     #prepare secret
     if isinstance(secret, unicode):
-        secret = secret.encode("utf-8")
+        secret = secret.encode(encoding)
     elif not isinstance(secret, str):
         raise TypeError("secret must be str or unicode")
 
     #prepare salt
     if isinstance(salt, unicode):
-        salt = salt.encode("utf-8")
+        salt = salt.encode(encoding)
     elif not isinstance(salt, str):
         raise TypeError("salt must be str or unicode")
 
@@ -146,7 +276,7 @@ def pbkdf2(secret, salt, rounds, keylen, prf="hmac-sha1"):
             return _EVP.pbkdf2(secret, salt, rounds, keylen)
 
     #resolve prf
-    encode_block, digest_size = _resolve_prf(prf)
+    encode_block, digest_size = get_prf(prf)
 
     #figure out how many blocks we'll need
     bcount = (keylen+digest_size-1)//digest_size
@@ -160,9 +290,11 @@ def pbkdf2(secret, salt, rounds, keylen, prf="hmac-sha1"):
         block = tmp = encode_block(secret, salt + pack(">L", i))
         #NOTE: could potentially unroll this loop somewhat for speed,
         # or find some faster way to accumulate & xor tmp values together
-        for j in xrange(rounds-1):
+        j = 1
+        while j < rounds:
             tmp = encode_block(secret, tmp)
             block = xor_bytes(block, tmp)
+            j += 1
         write(block)
 
     #and done