"""Various with-statement context utilities.""" import errno import inspect import os import pickle import signal import subprocess import sys import threading import traceback from contextlib import AbstractContextManager, contextmanager from importlib import import_module from multiprocessing.connection import Pipe from .cli.exceptions import UserException from .process import namespaces from .sequences import predicate_split # Ideas and code for SplitExec have been borrowed from withhacks # (https://pypi.python.org/pypi/withhacks) governed by the MIT license found # below. # # Copyright (c) 2010 Ryan Kelly # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN # THE SOFTWARE. class SplitExec: """Context manager separating code execution across parent/child processes. This is done by forking and doing some magic on the stack so the contents of the context are executed only on the forked child. Exceptions are pickled and passed back to the parent. """ def __init__(self): self.__trace_lock = threading.Lock() self.__orig_sys_trace = None self.__orig_trace_funcs = {} self.__injected_trace_funcs = {} self.__pipe = None self.childpid = None self.exit_status = -1 self.locals = {} def _parent_handler(self, signum, frame): """Signal handler for the parent process. By default this runs the parent cleanup and then resends the original signal to the parent process. """ self._cleanup() signal.signal(signum, signal.SIG_DFL) os.kill(os.getpid(), signum) def _parent_setup(self): """Initialization for parent process.""" try: signal.signal(signal.SIGINT, self._parent_handler) signal.signal(signal.SIGTERM, self._parent_handler) except ValueError: # skip if we're not in the main thread pass def _child_setup(self): """Initialization for child process.""" def _cleanup(self): """Parent process clean up on termination of the child.""" def _exception_cleanup(self): """Parent process clean up after the child throws an exception.""" self._cleanup() def _child_exit(self, exc): # pass back changed local scope vars from the child that can be pickled frame = self.__get_context_frame() local_vars = {} for k, v in frame.f_locals.items(): if k not in self.__child_orig_locals or v != self.__child_orig_locals[k]: try: pickle.dumps(v) local_vars[k] = v except (AttributeError, TypeError, pickle.PicklingError): continue exc._locals = local_vars try: self.__pipe.send(exc) except BrokenPipeError as e: if e.errno in (errno.EPIPE, errno.ESHUTDOWN): pass else: raise if isinstance(exc, SystemExit) and exc.code is not None: code = exc.code else: code = 0 os._exit(code) # pylint: disable=W0212 def __enter__(self): parent_pipe, child_pipe = Pipe() if pid := os.fork(): self.childpid = pid self._parent_setup() self.__pipe = parent_pipe frame = self.__get_context_frame() self.__inject_trace_func(frame, self.__exit_context) return self else: frame = self.__get_context_frame() self.__child_orig_locals = dict(frame.f_locals) self.__pipe = child_pipe try: self._child_setup() # pylint: disable=W0703 # need to catch all exceptions here since we are passing them to # the parent process except Exception as exc: exc.__traceback_list__ = traceback.format_exc() self._child_exit(exc) return self def __exit__(self, exc_type, exc_value, exc_traceback): if self.childpid is not None: # make sure system tracing function is reset self.__revert_tracing(inspect.currentframe()) # re-raise unknown exceptions from the parent if exc_type is not self.ParentException: raise exc_value # get exception from the child try: exc = self.__pipe.recv() self.locals = exc._locals except EOFError as e: exc = SystemExit(e) # handle child exiting abnormally if not isinstance(exc, SystemExit): os.waitpid(self.childpid, 0) self._exception_cleanup() sys.excepthook = self.__excepthook raise exc else: if exc_value is not None: exc = exc_value # Unfortunately, traceback objects can't be pickled so the relevant # traceback from the code executing within the chroot context is # placed in the __traceback_list__ attribute and printed by a # custom exception hook. exc.__traceback_list__ = traceback.format_exc() else: exc = SystemExit() self._child_exit(exc) # wait for child process to exit _pid, exit_status = os.waitpid(self.childpid, 0) self.exit_status = exit_status >> 8 self._cleanup() return True @staticmethod def __excepthook(_exc_type, exc_value, exc_traceback): """Output the proper traceback information from the chroot context.""" if hasattr(exc_value, "__traceback_list__"): sys.stderr.write(exc_value.__traceback_list__) else: traceback.print_tb(exc_traceback) @staticmethod def __dummy_sys_trace(frame, event, arg): """Dummy trace function used to enable tracing.""" class ParentException(Exception): """Exception used to detect when the child terminates.""" def __enable_tracing(self): """Enable system-wide tracing via a dummy method.""" self.__orig_sys_trace = sys.gettrace() sys.settrace(self.__dummy_sys_trace) def __revert_tracing(self, frame=None): """Revert to previous system trace setting.""" sys.settrace(self.__orig_sys_trace) if frame is not None: frame.f_trace = self.__orig_sys_trace def __exit_context(self, frame, event, arg): """Simple function to throw a ParentException.""" raise self.ParentException() def __inject_trace_func(self, frame, func): """Inject a trace function for a frame. The given trace function will be executed immediately when the frame's execution resumes. """ with self.__trace_lock: if frame.f_trace is not self.__invoke_trace_funcs: self.__orig_trace_funcs[frame] = frame.f_trace frame.f_trace = self.__invoke_trace_funcs self.__injected_trace_funcs[frame] = [] if len(self.__orig_trace_funcs) == 1: self.__enable_tracing() self.__injected_trace_funcs[frame].append(func) def __invoke_trace_funcs(self, frame, event, arg): """Invoke all trace funcs that have been injected. Once the injected functions have been executed all trace hooks are removed in order to minimize overhead. """ try: for func in self.__injected_trace_funcs[frame]: func(frame, event, arg) finally: del self.__injected_trace_funcs[frame] with self.__trace_lock: if len(self.__orig_trace_funcs) == 1: self.__revert_tracing() frame.f_trace = self.__orig_trace_funcs.pop(frame) def __get_context_frame(self): """Get the frame object for the with-statement context. This is designed to work from within superclass method call. It finds the first frame where the local variable "self" doesn't exist. """ try: return self.__frame except AttributeError: # an offset of two accounts for this method and its caller frame = inspect.stack(0)[2][0] while frame.f_locals.get("self") is self: frame = frame.f_back self.__frame = frame # pylint: disable=W0201 return frame class Namespace(SplitExec): """Context manager that provides Linux namespace support.""" def __init__( self, mount=False, uts=True, ipc=False, net=False, pid=False, user=False, hostname=None, ): self._hostname = hostname self._namespaces = { "mount": mount, "uts": uts, "ipc": ipc, "net": net, "pid": pid, "user": user, } super().__init__() def _child_setup(self): namespaces.simple_unshare(hostname=self._hostname, **self._namespaces) class GitStash(AbstractContextManager): """Context manager for stashing untracked or modified/uncommitted files.""" def __init__(self, path, pathspecs=None, staged=False): self.path = path self.pathspecs = ["--"] + pathspecs if pathspecs else [] self._staged = ["--keep-index"] if staged else [] self._stashed = False def __enter__(self): """Stash all untracked or modified files in working tree.""" # check for untracked or modified/uncommitted files try: p = subprocess.run( ["git", "status", "--porcelain=1", "-u"] + self.pathspecs, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, cwd=self.path, encoding="utf8", check=True, ) except subprocess.CalledProcessError: raise ValueError(f"not a git repo: {self.path}") # split file changes into unstaged vs staged unstaged, staged = predicate_split(lambda x: x[1] == " ", p.stdout.splitlines()) # don't stash when no relevant changes exist if self._staged: if not unstaged: return elif not p.stdout: return # stash all existing untracked or modified/uncommitted files try: stash_cmd = ["git", "stash", "push", "-u", "-m", "pkgcheck scan --commits"] subprocess.run( stash_cmd + self._staged + self.pathspecs, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, cwd=self.path, check=True, encoding="utf8", ) except subprocess.CalledProcessError as e: error = e.stderr.splitlines()[0] raise UserException(f"git failed stashing files: {error}") self._stashed = True def __exit__(self, _exc_type, _exc_value, _traceback): """Apply any previously stashed files back to the working tree.""" if self._stashed: try: subprocess.run( ["git", "stash", "pop"], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, cwd=self.path, check=True, encoding="utf8", ) except subprocess.CalledProcessError as e: error = e.stderr.splitlines()[0] raise UserException(f"git failed applying stash: {error}") @contextmanager def chdir(path): """Context manager that changes the current working directory. On exiting the context, the current working directory is switched back to its original value. :param path: The directory path to change the working directory to. """ orig_cwd = os.getcwd() os.chdir(path) try: yield finally: os.chdir(orig_cwd) @contextmanager def syspath(path: str, condition: bool = True, position: int = 0): """Context manager that mangles ``sys.path`` and then reverts on exit. :param path: The directory path to add to ``sys.path``. :param condition: Optional boolean that decides whether ``sys.path`` is mangled or not, defaults to being enabled. :param position: Optional integer that is the place where the path is inserted in ``sys.path``, defaults to prepending. """ syspath = sys.path[:] if condition: sys.path.insert(position, path) try: yield finally: sys.path = syspath @contextmanager def os_environ(*remove, **update): """Mangle the ``os.environ`` dictionary and revert on exit. :param remove: variables to remove :param update: variable -> value mapping to add or alter """ env = os.environ update = update or {} remove = remove or [] # variables being updated or removed changed = (set(update) | set(remove)) & set(env) # variables to restore on exit update_after = {k: env[k] for k in changed} # variables to remove on exit remove_after = frozenset(k for k in update if k not in env) try: env.update(update) for k in remove: env.pop(k, None) yield finally: env.update(update_after) for k in remove_after: env.pop(k) # Ideas and code for the patch context manager have been borrowed from mock # (https://github.com/testing-cabal/mock) governed by the BSD-2 license found # below. # # Copyright (c) 2007-2013, Michael Foord & the mock team # All rights reserved. # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions are # met: # # * Redistributions of source code must retain the above copyright # notice, this list of conditions and the following disclaimer. # # * Redistributions in binary form must reproduce the above # copyright notice, this list of conditions and the following # disclaimer in the documentation and/or other materials provided # with the distribution. # # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR # A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT # OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, # DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. @contextmanager def patch(target, new): """Simplified module monkey patching via context manager. :param target: Target class or object. :param new: Object or value to replace the target with. """ def _import_module(target): components = target.split(".") import_path = components.pop(0) module = import_module(import_path) for comp in components: try: module = getattr(module, comp) except AttributeError: import_path += ".%s" % comp module = import_module(import_path) return module def _get_target(target): if isinstance(target, str): try: module, attr = target.rsplit(".", 1) except (TypeError, ValueError): raise TypeError(f"invalid target: {target!r}") module = _import_module(module) return module, attr else: try: obj, attr = target except (TypeError, ValueError): raise TypeError(f"invalid target: {target!r}") return obj, attr obj, attr = _get_target(target) orig_attr = getattr(obj, attr) setattr(obj, attr, new) try: yield finally: setattr(obj, attr, orig_attr)