diff --git a/mitogen/core.py b/mitogen/core.py index 836c451a..d2d20723 100644 --- a/mitogen/core.py +++ b/mitogen/core.py @@ -168,6 +168,14 @@ def restart(func, *args): def is_blacklisted_import(importer, fullname): + """Return ``True`` if `fullname` is part of a blacklisted package, or if + any packages have been whitelisted and `fullname` is not part of one. + + NB: + - If a package is on both lists, then it is treated as blacklisted. + - If any package is whitelisted, then all non-whitelisted packages are + treated as blacklisted. + """ return ((not any(fullname.startswith(s) for s in importer.whitelist)) or (any(fullname.startswith(s) for s in importer.blacklist))) @@ -412,7 +420,7 @@ class Importer(object): 'utils', ]} self._lock = threading.Lock() - self.whitelist = whitelist or [''] + self.whitelist = list(whitelist) or [''] self.blacklist = list(blacklist) + [ # 2.x generates needless imports for 'builtins', while 3.x does the # same for '__builtin__'. The correct one is built-in, the other diff --git a/tests/importer_test.py b/tests/importer_test.py index 5fe65f45..55cad182 100644 --- a/tests/importer_test.py +++ b/tests/importer_test.py @@ -9,6 +9,8 @@ import pytest import unittest2 import mitogen.core +import mitogen.utils + import testlib @@ -25,6 +27,63 @@ class ImporterMixin(testlib.RouterMixin): super(ImporterMixin, self).tearDown() +class ImporterBlacklist(testlib.TestCase): + def test_is_blacklisted_import_default(self): + importer = mitogen.core.Importer( + router=mock.Mock(), context=None, core_src='', + ) + self.assertIsInstance(importer.whitelist, list) + self.assertIsInstance(importer.blacklist, list) + self.assertFalse(mitogen.core.is_blacklisted_import(importer, 'mypkg')) + self.assertFalse(mitogen.core.is_blacklisted_import(importer, 'mypkg.mod')) + self.assertFalse(mitogen.core.is_blacklisted_import(importer, 'otherpkg')) + self.assertFalse(mitogen.core.is_blacklisted_import(importer, 'otherpkg.mod')) + self.assertTrue(mitogen.core.is_blacklisted_import(importer, '__builtin__')) + self.assertTrue(mitogen.core.is_blacklisted_import(importer, 'builtins')) + + def test_is_blacklisted_import_just_whitelist(self): + importer = mitogen.core.Importer( + router=mock.Mock(), context=None, core_src='', + whitelist=('mypkg',), + ) + self.assertIsInstance(importer.whitelist, list) + self.assertIsInstance(importer.blacklist, list) + self.assertFalse(mitogen.core.is_blacklisted_import(importer, 'mypkg')) + self.assertFalse(mitogen.core.is_blacklisted_import(importer, 'mypkg.mod')) + self.assertTrue(mitogen.core.is_blacklisted_import(importer, 'otherpkg')) + self.assertTrue(mitogen.core.is_blacklisted_import(importer, 'otherpkg.mod')) + self.assertTrue(mitogen.core.is_blacklisted_import(importer, '__builtin__')) + self.assertTrue(mitogen.core.is_blacklisted_import(importer, 'builtins')) + + def test_is_blacklisted_import_just_blacklist(self): + importer = mitogen.core.Importer( + router=mock.Mock(), context=None, core_src='', + blacklist=('mypkg',), + ) + self.assertIsInstance(importer.whitelist, list) + self.assertIsInstance(importer.blacklist, list) + self.assertTrue(mitogen.core.is_blacklisted_import(importer, 'mypkg')) + self.assertTrue(mitogen.core.is_blacklisted_import(importer, 'mypkg.mod')) + self.assertFalse(mitogen.core.is_blacklisted_import(importer, 'otherpkg')) + self.assertFalse(mitogen.core.is_blacklisted_import(importer, 'otherpkg.mod')) + self.assertTrue(mitogen.core.is_blacklisted_import(importer, '__builtin__')) + self.assertTrue(mitogen.core.is_blacklisted_import(importer, 'builtins')) + + def test_is_blacklisted_import_whitelist_and_blacklist(self): + importer = mitogen.core.Importer( + router=mock.Mock(), context=None, core_src='', + whitelist=('mypkg',), blacklist=('mypkg',), + ) + self.assertIsInstance(importer.whitelist, list) + self.assertIsInstance(importer.blacklist, list) + self.assertTrue(mitogen.core.is_blacklisted_import(importer, 'mypkg')) + self.assertTrue(mitogen.core.is_blacklisted_import(importer, 'mypkg.mod')) + self.assertTrue(mitogen.core.is_blacklisted_import(importer, 'otherpkg')) + self.assertTrue(mitogen.core.is_blacklisted_import(importer, 'otherpkg.mod')) + self.assertTrue(mitogen.core.is_blacklisted_import(importer, '__builtin__')) + self.assertTrue(mitogen.core.is_blacklisted_import(importer, 'builtins')) + + class LoadModuleTest(ImporterMixin, testlib.TestCase): data = zlib.compress("data = 1\n\n") path = 'fake_module.py'