From b543b84e8045098dcd3feb1c0e5c02ce6211fd05 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Thu, 15 Feb 2018 20:11:38 +0545
Subject: [PATCH] importer: share blacklist logic between master/parent

---
 mitogen/core.py   | 8 ++++++--
 mitogen/master.py | 6 ++++--
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/mitogen/core.py b/mitogen/core.py
index 6da9a6e6..36661b96 100644
--- a/mitogen/core.py
+++ b/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)
diff --git a/mitogen/master.py b/mitogen/master.py
index 09982e6e..da0983c8 100644
--- a/mitogen/master.py
+++ b/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: