def _custom_import(self, name, globals=None, locals=None, fromlist=(), level=-1):
"""Custom '__import__' function.
Returns safe version of a module if configured in `_safe_module_config`.
Otherwise, returns standard module if the module is whitelisted.
Otherwise, blocks importing other modules.
"""
if not fromlist:
# Return the top-level package if 'fromlist' is empty (e.g. 'os' for 'os.path'),
# which is how '__import__' works.
name = name.split('.')[0]
frame = util.get_caller_frame(skip=[__name__])
filename = inspect.getframeinfo(frame).filename
# The import will be always allowed if it was not called from a project file.
if name in self._import_whitelist or not self._path_predicate(filename):
# Importing a module may cause more '__import__' calls if the module uses other
# modules. Such calls should not be blocked if the top-level import was allowed.
with self.allow_unsafe_import():
return ORIGINAL_IMPORT(name, globals, locals, fromlist, level)
# Return safe version of the module if possible
if name in self._safe_modules_config:
return self._get_safe_module(name)
raise ImportError(
'Importing module {0} is forbidden. '
'If you really need to import this module, read about '
'the allow_unsafe_import() function documented at: '
'https://buckbuild.com/function/allow_unsafe_import.html'.format(name))
def _get_safe_module(self, name):
"""Returns a safe version of the module."""
assert name in self._safe_modules_config, (
"Safe version of module %s is not configured." % name)
# Return the safe version of the module if already created
if name in self._safe_modules:
return self._safe_modules[name]
# Get the normal module, non-empty 'fromlist' prevents returning top-level package
# (e.g. 'os' would be returned for 'os.path' without it)
with self.allow_unsafe_import():
mod = ORIGINAL_IMPORT(name, fromlist=[''])
# Build a new module for the safe version
safe_mod = imp.new_module(name)
# Install whitelisted parts of the module, block the rest to produce errors
# informing about the safe version.
self._install_whitelisted_parts(mod, safe_mod, self._safe_modules_config[name])
# Store the safe version of the module
self._safe_modules[name] = safe_mod
return safe_mod