new PrefixWrapper constructor, for wrapping existing handlers and altering the prefix

5 files changed, 238 insertions, 4 deletions

diff --git a/docs/lib/passlib.utils.handlers.rst b/docs/lib/passlib.utils.handlers.rst
index 509e234..bacedfa 100644
--- a/docs/lib/passlib.utils.handlers.rst
+++ b/docs/lib/passlib.utils.handlers.rst
@@ -142,6 +142,10 @@ The StaticHandler class
 
 .. _testing-hash-handlers:
 
+Other Constructors
+==================
+.. autoclass:: PrefixWrapper
+
 Testing Hash Handlers
 =====================
 Within it's unittests, Passlib provides the :class:`~passlib.tests.utils.HandlerCase` class,
diff --git a/passlib/registry.py b/passlib/registry.py
index 70abf19..1f04ee3 100644
--- a/passlib/registry.py
+++ b/passlib/registry.py
@@ -294,7 +294,7 @@ def has_crypt_handler(name, loaded_only=False):
     global _handlers, _handler_locations
     return (name in _handlers) or (not loaded_only and name in _handler_locations)
 
-def _unload_handler_name(name):
+def _unload_handler_name(name, locations=True):
     """unloads a handler from the registry.
 
     .. warning::
@@ -306,14 +306,16 @@ def _unload_handler_name(name):
     if path to lazy load handler is found, its' removed.
 
     missing names are a noop.
+
+    :arg name: name of handler to unload
+    :param locations: if False, won't purge registered handler locations (default True)
     """
     global _handlers, _handler_locations
 
     if name in _handlers:
         del _handlers[name]
 
-    #NOTE: this messes w/ internals of registry, shouldn't be used publically.
-    if name in _handler_locations:
+    if locations and name in _handler_locations:
         del _handler_locations[name]
 
 #=========================================================
diff --git a/passlib/tests/test_utils_handlers.py b/passlib/tests/test_utils_handlers.py
index 4958716..ecc458e 100644
--- a/passlib/tests/test_utils_handlers.py
+++ b/passlib/tests/test_utils_handlers.py
@@ -10,8 +10,12 @@ from logging import getLogger
 import warnings
 #site
 #pkg
+from passlib.hash import ldap_md5
+from passlib.registry import _unload_handler_name as unload_handler_name, \
+    register_crypt_handler, get_crypt_handler
 from passlib.utils import rng, getrandstr, handlers as uh
-from passlib.tests.utils import HandlerCase, TestCase, catch_warnings
+from passlib.tests.utils import HandlerCase, TestCase, catch_warnings, \
+    dummy_handler_in_registry
 #module
 log = getLogger(__name__)
 
@@ -233,6 +237,84 @@ class SkeletonTest(TestCase):
     #=========================================================
 
 #=========================================================
+#PrefixWrapper
+#=========================================================
+class PrefixWrapperTest(TestCase):
+    "test PrefixWrapper class"
+
+    def test_00_lazy_loading(self):
+        "test PrefixWrapper lazy loading of handler"
+        d1 = uh.PrefixWrapper("d1", "ldap_md5", "{XXX}", "{MD5}", lazy=True)
+
+        #check base state
+        self.assertEqual(d1._wrapped_name, "ldap_md5")
+        self.assertIs(d1._wrapped_handler, None)
+
+        #check loading works
+        self.assertIs(d1.wrapped, ldap_md5)
+        self.assertIs(d1._wrapped_handler, ldap_md5)
+
+        #replace w/ wrong handler, make sure doesn't reload w/ dummy
+        with dummy_handler_in_registry("ldap_md5") as dummy:
+            self.assertIs(d1.wrapped, ldap_md5)
+
+    def test_01_active_loading(self):
+        "test PrefixWrapper active loading of handler"
+        d1 = uh.PrefixWrapper("d1", "ldap_md5", "{XXX}", "{MD5}")
+
+        #check base state
+        self.assertEqual(d1._wrapped_name, "ldap_md5")
+        self.assertIs(d1._wrapped_handler, ldap_md5)
+        self.assertIs(d1.wrapped, ldap_md5)
+
+        #replace w/ wrong handler, make sure doesn't reload w/ dummy
+        with dummy_handler_in_registry("ldap_md5") as dummy:
+            self.assertIs(d1.wrapped, ldap_md5)
+
+    def test_02_explicit(self):
+        "test PrefixWrapper with explicitly specified handler"
+
+        d1 = uh.PrefixWrapper("d1", ldap_md5, "{XXX}", "{MD5}")
+
+        #check base state
+        self.assertEqual(d1._wrapped_name, None)
+        self.assertIs(d1._wrapped_handler, ldap_md5)
+        self.assertIs(d1.wrapped, ldap_md5)
+
+        #replace w/ wrong handler, make sure doesn't reload w/ dummy
+        with dummy_handler_in_registry("ldap_md5") as dummy:
+            self.assertIs(d1.wrapped, ldap_md5)
+
+    def test_10_wrapped_attributes(self):
+        d1 = uh.PrefixWrapper("d1", "ldap_md5", "{XXX}", "{MD5}")
+        self.assertEqual(d1.name, "d1")
+        self.assertIs(d1.setting_kwds, ldap_md5.setting_kwds)
+
+    def test_11_wrapped_methods(self):
+        d1 = uh.PrefixWrapper("d1", "ldap_md5", "{XXX}", "{MD5}")
+        dph = "{XXX}X03MO1qnZdYdgyfeuILPmQ=="
+        lph = "{MD5}X03MO1qnZdYdgyfeuILPmQ=="
+
+        #genconfig
+        self.assertIs(d1.genconfig(), None)
+
+        #genhash
+        self.assertEqual(d1.genhash("password", None), dph)
+        self.assertEqual(d1.genhash("password", dph), dph)
+        self.assertRaises(ValueError, d1.genhash, "password", lph)
+
+        #encrypt
+        self.assertEqual(d1.encrypt("password"), dph)
+
+        #identify
+        self.assertTrue(d1.identify(dph))
+        self.assertFalse(d1.identify(lph))
+
+        #verify
+        self.assertRaises(ValueError, d1.verify, "password", lph)
+        self.assertTrue(d1.verify("password", dph))
+
+#=========================================================
 #sample algorithms - these serve as known quantities
 # to test the unittests themselves, as well as other
 # parts of passlib. they shouldn't be used as actual password schemes.
diff --git a/passlib/tests/utils.py b/passlib/tests/utils.py
index a1e35d4..4835f24 100644
--- a/passlib/tests/utils.py
+++ b/passlib/tests/utils.py
@@ -698,6 +698,27 @@ def create_backend_case(base_test, name):
     return dummy
 
 #=========================================================
+#misc helpers
+#=========================================================
+class dummy_handler_in_registry(object):
+    "context manager that inserts dummy handler in registry"
+    def __init__(self, name):
+        self.name = name
+        self.dummy = type('dummy_' + name, (uh.GenericHandler,), dict(
+            name=name,
+            setting_kwds=(),
+        ))
+
+    def __enter__(self):
+        registry._unload_handler_name(self.name, locations=False)
+        registry.register_crypt_handler(self.dummy)
+        assert registry.get_crypt_handler(self.name) is self.dummy
+        return self.dummy
+
+    def __exit__(self, *exc_info):
+        registry._unload_handler_name(self.name, locations=False)
+
+#=========================================================
 #helper for creating temp files - all cleaned up when prog exits
 #=========================================================
 tmp_files = []
diff --git a/passlib/utils/handlers.py b/passlib/utils/handlers.py
index fb11c34..ea346d7 100644
--- a/passlib/utils/handlers.py
+++ b/passlib/utils/handlers.py
@@ -13,6 +13,7 @@ import os
 from warnings import warn
 #site
 #libs
+from passlib.registry import get_crypt_handler
 from passlib.utils import classproperty, h64, getrandstr, getrandbytes, rng, is_crypt_handler, ALL_BYTE_VALUES
 #pkg
 #local
@@ -27,6 +28,7 @@ __all__ = [
             'HasRawSalt',
         'HasRounds',
         'HasManyBackends',
+    'PrefixWrapper',
 ]
 
 #=========================================================
@@ -1460,5 +1462,128 @@ class HasManyBackends(GenericHandler):
         return self.calc_checksum(secret)
 
 #=========================================================
+#wrappers
+#=========================================================
+class PrefixWrapper(object):
+    """wraps another handler, adding a constant prefix.
+
+    instances of this class wrap another password hash handler,
+    altering the constant prefix that's prepended to the wrapped
+    handlers' hashes.
+
+    this is used mainly by the :samp:`ldap_{digest}_crypt` methods,
+    such as :class:`ldap_md5_crypt`, which wraps :class:`md5_crypt` and adds a ``{CRYPT}`` prefix.
+
+    usage::
+
+        myhandler = PrefixWrapper("myhandler", "md5_crypt", prefix="$mh$", orig_prefix="$1$")
+
+    :param name: name to assign to handler
+    :param wrapped: handler object or name of registered handler
+    :param prefix: identifying prefix to prepend to all hashes
+    :param orig_prefix: prefix to strip (defaults to '').
+    :param lazy: if True and wrapped handler is specified by name, don't look it up until needed.
+    """
+
+    def __init__(self, name, wrapped, prefix='', orig_prefix='', lazy=False, doc=None):
+        self.name = name
+        self.prefix = prefix
+        self.orig_prefix = orig_prefix
+        if doc:
+            self.__doc__ = doc
+        if hasattr(wrapped, "name"):
+            self._check_handler(wrapped)
+            self._wrapped_handler = wrapped
+        else:
+            self._wrapped_name = wrapped
+            if not lazy:
+                self._get_wrapped()
+
+    _wrapped_name = None
+    _wrapped_handler = None
+
+    def _check_handler(self, handler):
+        if 'ident' in handler.setting_kwds:
+            #TODO: look into way to fix the issues.
+            warn("PrefixWrapper may not work correctly for handlers which have multiple identifiers: %r" % (handler.name,))
+
+    def _get_wrapped(self):
+        handler = self._wrapped_handler
+        if handler is None:
+            handler = get_crypt_handler(self._wrapped_name)
+            self._check_handler(handler)
+            self._wrapped_handler = handler
+        return handler
+
+    wrapped = property(_get_wrapped)
+
+    ##@property
+    ##def ident(self):
+    ##    return self._prefix
+
+    #attrs that should be proxied
+    _proxy_attrs = (
+                    "setting_kwds", "context_kwds",
+                    "default_rounds", "min_rounds", "max_rounds", "rounds_cost",
+                    "backends", "has_backend", "get_backend", "set_backend",
+                    )
+
+    def __repr__(self):
+        args = [ repr(self._wrapped_name or self._wrapped_handler) ]
+        if self.prefix:
+            args.append("prefix=%r" % self.prefix)
+        if self.orig_prefix:
+            args.append("orig_prefix=%r", self.orig_prefix)
+        args = ", ".join(args)
+        return 'PrefixWrapper(%r, %s)' % (self.name, args)
+
+    def __getattr__(self, attr):
+        "proxy most attributes from wrapped class (eg rounds, salt size, etc)"
+        if attr in self._proxy_attrs:
+            return getattr(self.wrapped, attr)
+        raise AttributeError("missing attribute: %r" % (attr,))
+
+    def _unwrap_hash(self, hash):
+        "given hash belonging to wrapper, return orig version"
+        prefix = self.prefix
+        if not hash.startswith(prefix):
+            raise ValueError("not a valid %s hash" % (self.name,))
+        return self.orig_prefix + hash[len(prefix):]
+
+    def _wrap_hash(self, hash):
+        "given orig hash; return one belonging to wrapper"
+        prefix = self.orig_prefix
+        if not hash.startswith(prefix):
+            raise ValueError("not a valid %s hash" % (self.wrapped.name,))
+        return self.prefix + hash[len(prefix):]
+
+    def identify(self, hash):
+        if not hash or not hash.startswith(self.prefix):
+            return False
+        hash = self._unwrap_hash(hash)
+        return self.wrapped.identify(hash)
+
+    def genconfig(self, **kwds):
+        config = self.wrapped.genconfig(**kwds)
+        if config:
+            return self._wrap_hash(config)
+        else:
+            return config
+
+    def genhash(self, secret, config, **kwds):
+        if config:
+            config = self._unwrap_hash(config)
+        return self._wrap_hash(self.wrapped.genhash(secret, config, **kwds))
+
+    def encrypt(self, secret, **kwds):
+        return self._wrap_hash(self.wrapped.encrypt(secret, **kwds))
+
+    def verify(self, secret, hash, **kwds):
+        if not hash:
+            raise ValueError("no %s hash specified" % (self.name,))
+        hash = self._unwrap_hash(hash)
+        return self.wrapped.verify(secret, hash, **kwds)
+
+#=========================================================
 # eof
 #=========================================================