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

[brett.cannon]

Author: brett.cannon
Date: Mon Feb  5 22:42:38 2007
New Revision: 53636

Modified:
   sandbox/trunk/import_in_py/controlled_importlib.py
   sandbox/trunk/import_in_py/test_controlled_importlib.py
Log:
Add whitelisting support for extension modules.


Modified: sandbox/trunk/import_in_py/controlled_importlib.py
==============================================================================
--- sandbox/trunk/import_in_py/controlled_importlib.py	(original)
+++ sandbox/trunk/import_in_py/controlled_importlib.py	Mon Feb  5 22:42:38 2007
@@ -58,6 +58,25 @@
 
     pass
 
+class WhitelistExtHandler(importlib.ExtensionFileHandler):
+
+    """Add whitelisting to the extension module handler."""
+
+    def __init__(self, whitelist):
+        self._whitelist = frozenset(whitelist)
+        super(WhitelistExtHandler, self).__init__()
+
+    def cannot_handle(self, name):
+        if name in self._whitelist:
+            return False
+        return True
+
+    def handle_code(self, loader, mod_name, extension_path, package=None):
+        if mod_name in self._whitelist:
+            return super(WhitelistExtHandler, self).handle_code(loder, mod_name,
+                    extension_path, package)
+        raise ImportError("not allowed to load module")
+
 
 class ControlledImport(importlib.Import):
 

Modified: sandbox/trunk/import_in_py/test_controlled_importlib.py
==============================================================================
--- sandbox/trunk/import_in_py/test_controlled_importlib.py	(original)
+++ sandbox/trunk/import_in_py/test_controlled_importlib.py	Mon Feb  5 22:42:38 2007
@@ -1,7 +1,9 @@
 import controlled_importlib
+import importlib
 import mock_importlib
 
 from contextlib import contextmanager, nested
+import os
 import StringIO
 import sys
 import unittest
@@ -108,6 +110,55 @@
             self.failUnlessRaises(ImportError, imp_load.load_module, blacklist)
 
 
+class ExtensionHelper(unittest.TestCase):
+
+    whitelist = 'time'
+    blacklist = 'datetime'
+
+    def setUp(self):
+        self.imp_load = controlled_importlib.WhitelistExtHandler([self.whitelist])
+
+
+class WhitelistExtensionsTests(ExtensionHelper):
+
+    """Test the whitelisting of extension modules."""
+
+    def test_cannot_handle(self):
+        # Should return False for modules on the whitelist, True otherwise.
+        self.failUnless(not self.imp_load.cannot_handle(self.whitelist))
+        self.failUnless(self.imp_load.cannot_handle(self.blacklist))
+
+    def handle_code(self):
+        # Should raise ImportError if the module is not whitelisted, otherwise
+        # return the module.
+        module = self.imp_load.handle_code(None, self.whitelist,
+                self.imp_load.handles[0])
+        self.failUnlessEqual(module.__name__, self.whitelist)
+        self.failUnlessRaises(ImportError, self.imp_load.handle_code, None,
+                self.blacklist, self.imp_load.handles[0])
+
+
+class FileSystemImporterTests(ExtensionHelper):
+
+    """Make sure that the filesystem importer respects what the handler's
+    cannot_handle method says."""
+
+    def test_find_module(self):
+        # Make sure that when the handler is used in a filesystem importer that
+        # whitelisting is still respected for find_module.
+        for entry in sys.path:
+            blacklist_filename = self.blacklist + self.imp_load.handles[0]
+            blacklist_path = os.path.join(entry, blacklist_filename)
+            if os.path.exists(blacklist_path):
+                sys_entry = entry
+                break
+        else:
+            raise test_support.TestSkipped("sys.path empty")
+        importer = importlib.FileSystemImporter(sys_entry, self.imp_load)
+        self.failUnless(importer.find_module(self.whitelist))
+        self.failUnless(importer.find_module(self.blacklist) is None)
+
+
 @contextmanager
 def mutate_sys_modules(module, name):
     """Temporarily mutate sys.modules with a new module."""
@@ -242,6 +293,8 @@
     test_support.run_unittest(WhitelistTests,
                                 WhitelistBuiltinTests,
                                 WhitelistFrozenTests,
+                                WhitelistExtensionsTests,
+                                FileSystemImporterTests,
                                 ControlledImportMethodTests,
                                 ControlledImportUsageTests)