importer: share blacklist logic between master/parent

mitogen/core.py

@@ -167,6 +167,11 @@ def restart(func, *args):
                 raise
 
 
+def is_blacklisted_import(importer, fullname):
+    return ((not any(fullname.startswith(s) for s in importer.whitelist)) or
+                 (any(fullname.startswith(s) for s in importer.blacklist)))
+
+
 def set_cloexec(fd):
     flags = fcntl.fcntl(fd, fcntl.F_GETFD)
     fcntl.fcntl(fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC)
@@ -461,8 +466,7 @@ class Importer(object):
             del _tls.running
 
     def _refuse_imports(self, fullname):
-        if ((not any(fullname.startswith(s) for s in self.whitelist)) or
-                (any(fullname.startswith(s) for s in self.blacklist))):
+        if is_blacklisted_import(self, fullname):
             raise ImportError('Refused')
 
         f = sys._getframe(2)

mitogen/master.py

@@ -442,7 +442,7 @@ class ModuleResponder(object):
         self._finder = ModuleFinder()
         self._cache = {}  # fullname -> pickled
         self.blacklist = []
-        self.whitelist = []
+        self.whitelist = ['']
         router.add_handler(self._on_get_module, mitogen.core.GET_MODULE)
 
     def __repr__(self):
@@ -451,6 +451,8 @@ class ModuleResponder(object):
     MAIN_RE = re.compile(r'^if\s+__name__\s*==\s*.__main__.\s*:', re.M)
 
     def whitelist_prefix(self, fullname):
+        if self.whitelist == ['']:
+            self.whitelist = ['mitogen']
         self.whitelist.append(fullname)
 
     def blacklist_prefix(self, fullname):
@@ -466,7 +468,7 @@ class ModuleResponder(object):
         return src
 
     def _build_tuple(self, fullname):
-        if fullname in self._blacklist:
+        if mitogen.core.is_blacklisted_import(self, fullname):
             raise ImportError('blacklisted')
 
         if fullname in self._cache: