[Python-checkins] r53623 - sandbox/trunk/import_in_py/controlled_importlib.py sandbox/trunk/import_in_py/test_controlled_importlib.py

brett.cannon python-checkins at python.org
Fri Feb 2 01:27:14 CET 2007


Author: brett.cannon
Date: Fri Feb  2 01:27:13 2007
New Revision: 53623

Added:
   sandbox/trunk/import_in_py/controlled_importlib.py   (contents, props changed)
   sandbox/trunk/import_in_py/test_controlled_importlib.py   (contents, props changed)
Log:
Begin implementation of a controlling import that usees whitelisting to
restrict what can and cannot be imported.


Added: sandbox/trunk/import_in_py/controlled_importlib.py
==============================================================================
--- (empty file)
+++ sandbox/trunk/import_in_py/controlled_importlib.py	Fri Feb  2 01:27:13 2007
@@ -0,0 +1,86 @@
+"""Create a controllable import.
+
+One should be able to do several things to control imports:
+
+    * Whitelisting of modules based on name and type
+        + built-ins
+        + frozen
+        + extensions
+    * Block all modules based on type
+        + .pyc
+    * Allow all modules of a certain type
+        + .py
+"""
+import importlib
+import sys
+
+class Whitelister(object):
+
+    """Whitelist the finding and/or loading of modules."""
+
+    def __init__(self, whitelist, *args, **kwargs):
+        """Store whitelist and continue on with instantiation."""
+        self._whitelist = frozenset(whitelist)
+        super(Whitelister, self).__init__(*args, **kwargs)
+
+    def check_name(self, name):
+        """Check a module name against the whitelist, returning True if it
+        passes or False otherwise."""
+        if name in self._whitelist:
+            return True
+        return False
+
+    def find_module(self, name, path=None):
+        """If the module name passes the whitelist, allow the finding of the
+        modules to continue, otherwise return None."""
+        if self.check_name(name):
+            return super(Whitelister, self).find_module(name, path)
+        return None
+
+    def load_module(self, name):
+        """If the module names passes the whitelist, allow the loading of the
+        module to continue, otherwise raise ImportError."""
+        if self.check_name(name):
+            return super(Whitelister, self).load_module(name)
+        raise ImportError("cannot import module")
+
+
+class WhitelistBuiltin(Whitelister, importlib.BuiltinImporter):
+
+    """Whitelisting importer/loader for built-in modules."""
+
+    pass
+    
+
+class WhitelistFrozen(Whitelister, importlib.FrozenImporter):
+
+    """Whitelisting importer/loader for frozen modules."""
+
+    pass
+
+
+class ControlledImport(importlib.Import):
+
+    """Represent a controllable version of import that allows for more
+    fine-grained control over what can and cannot be imported."""
+
+    def __init__(self, safe_builtins, safe_frozen, safe_extensions):
+        """Set up importation where built-in, frozen, and extension modules
+        must be whitelisted and .pyc files are completely blocked
+        while all. py files are allowed."""
+        # Clear out any traces of the previous import machinery.
+        sys.path_importer_cache.clear()
+        # Whitelist built-in and frozen modules on sys.meta_path.
+        # XXX
+        # Whitelist extension modules on sys.path.
+        # XXX
+        # Allow all .py files but not .pyc files on sys.path.
+        # XXX
+
+    def module_from_cache(self, name):
+        """Override so that any module name starting with a dot raises
+        ImportError."""
+        if name.startswith('.'):
+            raise ImportError("module names starting with '.' cannot be "
+                                "imported")
+        return importlib.Import.module_from_cache(self, name)

Added: sandbox/trunk/import_in_py/test_controlled_importlib.py
==============================================================================
--- (empty file)
+++ sandbox/trunk/import_in_py/test_controlled_importlib.py	Fri Feb  2 01:27:13 2007
@@ -0,0 +1,117 @@
+import controlled_importlib
+
+import StringIO
+import sys
+import unittest
+from test import test_support
+
+
+class DummyImporterLoader(object):
+
+    def find_module(self, name, path=None):
+        return True
+
+    def load_module(self, name):
+        return True
+
+class DummyWhitelist(controlled_importlib.Whitelister, DummyImporterLoader):
+
+    pass
+
+
+class WhitelistTests(unittest.TestCase):
+
+    """Test the general whitelisting mechanism."""
+
+    def test_whitelist_module(self):
+        # A direct module name should be allowed.
+        whitelist = 'ok_mod'
+        imp_load = DummyWhitelist([whitelist])
+        self.failUnless(imp_load.find_module(whitelist) is True)
+        self.failUnless(imp_load.find_module('fdssdf') is None)
+        self.failUnless(imp_load.load_module(whitelist) is True)
+        self.failUnlessRaises(ImportError, imp_load.load_module, 'asdfadsf')
+
+    def test_whitelist_list(self):
+        # When the whitelist is a list it should be able to properly whitelist
+        # no matter where a module is listed.
+        whitelist1 = 'A'
+        whitelist2 = 'B'
+        imp_load = DummyWhitelist([whitelist1, whitelist2])
+        for whitelist in (whitelist1, whitelist2):
+            self.failUnless(imp_load.find_module(whitelist) is True)
+            self.failUnless(imp_load.load_module(whitelist) is True)
+
+    def test_block_partial_name(self):
+        # A module that happens to be a prefix of a whitelisted module or has a
+        # whitelisted module as a prefix should still be blocked.
+        whitelist = 'mod'
+        imp_load = DummyWhitelist([whitelist])
+        # Module has a whitelisted module as a prefix.
+        self.failUnless(imp_load.find_module(whitelist+'2') is None)
+        self.failUnlessRaises(ImportError, imp_load.load_module, whitelist+'2')
+        # Module is a prefix of a whitelisted module.
+        self.failUnless(imp_load.find_module(whitelist[:-1]) is None)
+        self.failUnlessRaises(ImportError, imp_load.load_module,
+                whitelist[:-1])
+
+    def test_package(self):
+        # Whitelisting a package does not automatically allow the submodules.
+        whitelist = 'pkg'
+        imp_load = DummyWhitelist([whitelist])
+        mod = whitelist + '.' + 'mod'
+        self.failUnless(imp_load.find_module(mod) is None)
+        self.failUnlessRaises(ImportError, imp_load.load_module, mod)
+
+
+class WhitelistBuiltinTests(unittest.TestCase):
+
+    """Test the whitelisting support for built-in modules."""
+
+    def setUp(self):
+        self.whitelist = sys.builtin_module_names[0]
+        self.blacklist = sys.builtin_module_names[1]
+
+    def test_whitelist(self):
+        # Only modules on the whitelist should be allowed to be imported.
+        # Everything else should return None.
+        imp_load = controlled_importlib.WhitelistBuiltin([self.whitelist])
+        # Importer
+        self.failUnless(imp_load.find_module(self.whitelist) is not None)
+        self.failUnless(imp_load.find_module(self.blacklist) is None)
+        # Loader
+        self.failUnless(imp_load.load_module(self.whitelist))
+        self.failUnlessRaises(ImportError, imp_load.load_module,
+                                self.blacklist)
+
+
+class WhitelistFrozenTests(unittest.TestCase):
+
+    """Test whitelisting of frozen modules."""
+
+    def setUp(self):
+        sys.stdout = StringIO.StringIO()
+        self.whitelist = '__phello__'
+        self.blacklist = ('__hello__', '__phello__.spam')
+
+    def tearDown(self):
+        sys.stdout = sys.__stdout__
+
+    def test_whitelist(self):
+        imp_load = controlled_importlib.WhitelistFrozen([self.whitelist])
+        self.failUnless(imp_load.find_module(self.whitelist) is not None)
+        self.failUnless(imp_load.load_module(self.whitelist))
+        for blacklist in self.blacklist:
+            self.failUnless(imp_load.find_module(blacklist) is None)
+            self.failUnlessRaises(ImportError, imp_load.load_module, blacklist)
+
+
+
+def test_main():
+    test_support.run_unittest(WhitelistTests,
+                                WhitelistBuiltinTests,
+                                WhitelistFrozenTests)
+
+
+if __name__ == '__main__':
+    test_main()


More information about the Python-checkins mailing list