import os import re from contextlib import contextmanager from itertools import cycle from os import PathLike from typing import ( TYPE_CHECKING, Any, Generator, Iterator, List, Optional, Union, ) from unittest.mock import Mock from .config import Config, DataProxy from .exceptions import Failure, AuthFailure, ResponseNotAccepted from .runners import Result from .watchers import FailingResponder if TYPE_CHECKING: from invoke.runners import Runner class Context(DataProxy): """ Context-aware API wrapper & state-passing object. `.Context` objects are created during command-line parsing (or, if desired, by hand) and used to share parser and configuration state with executed tasks (see :ref:`why-context`). Specifically, the class offers wrappers for core API calls (such as `.run`) which take into account CLI parser flags, configuration files, and/or changes made at runtime. It also acts as a proxy for its `~.Context.config` attribute - see that attribute's documentation for details. Instances of `.Context` may be shared between tasks when executing sub-tasks - either the same context the caller was given, or an altered copy thereof (or, theoretically, a brand new one). .. versionadded:: 1.0 """ def __init__(self, config: Optional[Config] = None) -> None: """ :param config: `.Config` object to use as the base configuration. Defaults to an anonymous/default `.Config` instance. """ #: The fully merged `.Config` object appropriate for this context. #: #: `.Config` settings (see their documentation for details) may be #: accessed like dictionary keys (``c.config['foo']``) or object #: attributes (``c.config.foo``). #: #: As a convenience shorthand, the `.Context` object proxies to its #: ``config`` attribute in the same way - e.g. ``c['foo']`` or #: ``c.foo`` returns the same value as ``c.config['foo']``. config = config if config is not None else Config() self._set(_config=config) #: A list of commands to run (via "&&") before the main argument to any #: `run` or `sudo` calls. Note that the primary API for manipulating #: this list is `prefix`; see its docs for details. command_prefixes: List[str] = list() self._set(command_prefixes=command_prefixes) #: A list of directories to 'cd' into before running commands with #: `run` or `sudo`; intended for management via `cd`, please see its #: docs for details. command_cwds: List[str] = list() self._set(command_cwds=command_cwds) @property def config(self) -> Config: # Allows Context to expose a .config attribute even though DataProxy # otherwise considers it a config key. return self._config @config.setter def config(self, value: Config) -> None: # NOTE: mostly used by client libraries needing to tweak a Context's # config at execution time; i.e. a Context subclass that bears its own # unique data may want to be stood up when parameterizing/expanding a # call list at start of a session, with the final config filled in at # runtime. self._set(_config=value) def run(self, command: str, **kwargs: Any) -> Optional[Result]: """ Execute a local shell command, honoring config options. Specifically, this method instantiates a `.Runner` subclass (according to the ``runner`` config option; default is `.Local`) and calls its ``.run`` method with ``command`` and ``kwargs``. See `.Runner.run` for details on ``command`` and the available keyword arguments. .. versionadded:: 1.0 """ runner = self.config.runners.local(self) return self._run(runner, command, **kwargs) # NOTE: broken out of run() to allow for runner class injection in # Fabric/etc, which needs to juggle multiple runner class types (local and # remote). def _run( self, runner: "Runner", command: str, **kwargs: Any ) -> Optional[Result]: command = self._prefix_commands(command) return runner.run(command, **kwargs) def sudo(self, command: str, **kwargs: Any) -> Optional[Result]: """ Execute a shell command via ``sudo`` with password auto-response. **Basics** This method is identical to `run` but adds a handful of convenient behaviors around invoking the ``sudo`` program. It doesn't do anything users could not do themselves by wrapping `run`, but the use case is too common to make users reinvent these wheels themselves. .. note:: If you intend to respond to sudo's password prompt by hand, just use ``run("sudo command")`` instead! The autoresponding features in this method will just get in your way. Specifically, `sudo`: * Places a `.FailingResponder` into the ``watchers`` kwarg (see :doc:`/concepts/watchers`) which: * searches for the configured ``sudo`` password prompt; * responds with the configured sudo password (``sudo.password`` from the :doc:`configuration `); * can tell when that response causes an authentication failure (e.g. if the system requires a password and one was not configured), and raises `.AuthFailure` if so. * Builds a ``sudo`` command string using the supplied ``command`` argument, prefixed by various flags (see below); * Executes that command via a call to `run`, returning the result. **Flags used** ``sudo`` flags used under the hood include: - ``-S`` to allow auto-responding of password via stdin; - ``-p `` to explicitly state the prompt to use, so we can be sure our auto-responder knows what to look for; - ``-u `` if ``user`` is not ``None``, to execute the command as a user other than ``root``; - When ``-u`` is present, ``-H`` is also added, to ensure the subprocess has the requested user's ``$HOME`` set properly. **Configuring behavior** There are a couple of ways to change how this method behaves: - Because it wraps `run`, it honors all `run` config parameters and keyword arguments, in the same way that `run` does. - Thus, invocations such as ``c.sudo('command', echo=True)`` are possible, and if a config layer (such as a config file or env var) specifies that e.g. ``run.warn = True``, that too will take effect under `sudo`. - `sudo` has its own set of keyword arguments (see below) and they are also all controllable via the configuration system, under the ``sudo.*`` tree. - Thus you could, for example, pre-set a sudo user in a config file; such as an ``invoke.json`` containing ``{"sudo": {"user": "someuser"}}``. :param str password: Runtime override for ``sudo.password``. :param str user: Runtime override for ``sudo.user``. .. versionadded:: 1.0 """ runner = self.config.runners.local(self) return self._sudo(runner, command, **kwargs) # NOTE: this is for runner injection; see NOTE above _run(). def _sudo( self, runner: "Runner", command: str, **kwargs: Any ) -> Optional[Result]: prompt = self.config.sudo.prompt password = kwargs.pop("password", self.config.sudo.password) user = kwargs.pop("user", self.config.sudo.user) env = kwargs.get("env", {}) # TODO: allow subclassing for 'get the password' so users who REALLY # want lazy runtime prompting can have it easily implemented. # TODO: want to print a "cleaner" echo with just 'sudo '; but # hard to do as-is, obtaining config data from outside a Runner one # holds is currently messy (could fix that), if instead we manually # inspect the config ourselves that duplicates logic. NOTE: once we # figure that out, there is an existing, would-fail-if-not-skipped test # for this behavior in test/context.py. # TODO: once that is done, though: how to handle "full debug" output # exactly (display of actual, real full sudo command w/ -S and -p), in # terms of API/config? Impl is easy, just go back to passing echo # through to 'run'... user_flags = "" if user is not None: user_flags = "-H -u {} ".format(user) env_flags = "" if env: env_flags = "--preserve-env='{}' ".format(",".join(env.keys())) command = self._prefix_commands(command) cmd_str = "sudo -S -p '{}' {}{}{}".format( prompt, env_flags, user_flags, command ) watcher = FailingResponder( pattern=re.escape(prompt), response="{}\n".format(password), sentinel="Sorry, try again.\n", ) # Ensure we merge any user-specified watchers with our own. # NOTE: If there are config-driven watchers, we pull those up to the # kwarg level; that lets us merge cleanly without needing complex # config-driven "override vs merge" semantics. # TODO: if/when those semantics are implemented, use them instead. # NOTE: config value for watchers defaults to an empty list; and we # want to clone it to avoid actually mutating the config. watchers = kwargs.pop("watchers", list(self.config.run.watchers)) watchers.append(watcher) try: return runner.run(cmd_str, watchers=watchers, **kwargs) except Failure as failure: # Transmute failures driven by our FailingResponder, into auth # failures - the command never even ran. # TODO: wants to be a hook here for users that desire "override a # bad config value for sudo.password" manual input # NOTE: as noted in #294 comments, we MAY in future want to update # this so run() is given ability to raise AuthFailure on its own. # For now that has been judged unnecessary complexity. if isinstance(failure.reason, ResponseNotAccepted): # NOTE: not bothering with 'reason' here, it's pointless. error = AuthFailure(result=failure.result, prompt=prompt) raise error # Reraise for any other error so it bubbles up normally. else: raise # TODO: wonder if it makes sense to move this part of things inside Runner, # which would grow a `prefixes` and `cwd` init kwargs or similar. The less # that's stuffed into Context, probably the better. def _prefix_commands(self, command: str) -> str: """ Prefixes ``command`` with all prefixes found in ``command_prefixes``. ``command_prefixes`` is a list of strings which is modified by the `prefix` context manager. """ prefixes = list(self.command_prefixes) current_directory = self.cwd if current_directory: prefixes.insert(0, "cd {}".format(current_directory)) return " && ".join(prefixes + [command]) @contextmanager def prefix(self, command: str) -> Generator[None, None, None]: """ Prefix all nested `run`/`sudo` commands with given command plus ``&&``. Most of the time, you'll want to be using this alongside a shell script which alters shell state, such as ones which export or alter shell environment variables. For example, one of the most common uses of this tool is with the ``workon`` command from `virtualenvwrapper `_:: with c.prefix('workon myvenv'): c.run('./manage.py migrate') In the above snippet, the actual shell command run would be this:: $ workon myvenv && ./manage.py migrate This context manager is compatible with `cd`, so if your virtualenv doesn't ``cd`` in its ``postactivate`` script, you could do the following:: with c.cd('/path/to/app'): with c.prefix('workon myvenv'): c.run('./manage.py migrate') c.run('./manage.py loaddata fixture') Which would result in executions like so:: $ cd /path/to/app && workon myvenv && ./manage.py migrate $ cd /path/to/app && workon myvenv && ./manage.py loaddata fixture Finally, as alluded to above, `prefix` may be nested if desired, e.g.:: with c.prefix('workon myenv'): c.run('ls') with c.prefix('source /some/script'): c.run('touch a_file') The result:: $ workon myenv && ls $ workon myenv && source /some/script && touch a_file Contrived, but hopefully illustrative. .. versionadded:: 1.0 """ self.command_prefixes.append(command) try: yield finally: self.command_prefixes.pop() @property def cwd(self) -> str: """ Return the current working directory, accounting for uses of `cd`. .. versionadded:: 1.0 """ if not self.command_cwds: # TODO: should this be None? Feels cleaner, though there may be # benefits to it being an empty string, such as relying on a no-arg # `cd` typically being shorthand for "go to user's $HOME". return "" # get the index for the subset of paths starting with the last / or ~ for i, path in reversed(list(enumerate(self.command_cwds))): if path.startswith("~") or path.startswith("/"): break # TODO: see if there's a stronger "escape this path" function somewhere # we can reuse. e.g., escaping tildes or slashes in filenames. paths = [path.replace(" ", r"\ ") for path in self.command_cwds[i:]] return str(os.path.join(*paths)) @contextmanager def cd(self, path: Union[PathLike, str]) -> Generator[None, None, None]: """ Context manager that keeps directory state when executing commands. Any calls to `run`, `sudo`, within the wrapped block will implicitly have a string similar to ``"cd && "`` prefixed in order to give the sense that there is actually statefulness involved. Because use of `cd` affects all such invocations, any code making use of the `cwd` property will also be affected by use of `cd`. Like the actual 'cd' shell builtin, `cd` may be called with relative paths (keep in mind that your default starting directory is your user's ``$HOME``) and may be nested as well. Below is a "normal" attempt at using the shell 'cd', which doesn't work since all commands are executed in individual subprocesses -- state is **not** kept between invocations of `run` or `sudo`:: c.run('cd /var/www') c.run('ls') The above snippet will list the contents of the user's ``$HOME`` instead of ``/var/www``. With `cd`, however, it will work as expected:: with c.cd('/var/www'): c.run('ls') # Turns into "cd /var/www && ls" Finally, a demonstration (see inline comments) of nesting:: with c.cd('/var/www'): c.run('ls') # cd /var/www && ls with c.cd('website1'): c.run('ls') # cd /var/www/website1 && ls .. note:: Space characters will be escaped automatically to make dealing with such directory names easier. .. versionadded:: 1.0 .. versionchanged:: 1.5 Explicitly cast the ``path`` argument (the only argument) to a string; this allows any object defining ``__str__`` to be handed in (such as the various ``Path`` objects out there), and not just string literals. """ path = str(path) self.command_cwds.append(path) try: yield finally: self.command_cwds.pop() class MockContext(Context): """ A `.Context` whose methods' return values can be predetermined. Primarily useful for testing Invoke-using codebases. .. note:: This class wraps its ``run``, etc methods in `unittest.mock.Mock` objects. This allows you to easily assert that the methods (still returning the values you prepare them with) were actually called. .. note:: Methods not given `Results <.Result>` to yield will raise ``NotImplementedError`` if called (since the alternative is to call the real underlying method - typically undesirable when mocking.) .. versionadded:: 1.0 .. versionchanged:: 1.5 Added ``Mock`` wrapping of ``run`` and ``sudo``. """ def __init__(self, config: Optional[Config] = None, **kwargs: Any) -> None: """ Create a ``Context``-like object whose methods yield `.Result` objects. :param config: A Configuration object to use. Identical in behavior to `.Context`. :param run: A data structure indicating what `.Result` objects to return from calls to the instantiated object's `~.Context.run` method (instead of actually executing the requested shell command). Specifically, this kwarg accepts: - A single `.Result` object. - A boolean; if True, yields a `.Result` whose ``exited`` is ``0``, and if False, ``1``. - An iterable of the above values, which will be returned on each subsequent call to ``.run`` (the first item on the first call, the second on the second call, etc). - A dict mapping command strings or compiled regexen to the above values (including an iterable), allowing specific call-and-response semantics instead of assuming a call order. :param sudo: Identical to ``run``, but whose values are yielded from calls to `~.Context.sudo`. :param bool repeat: A flag determining whether results yielded by this class' methods repeat or are consumed. For example, when a single result is indicated, it will normally only be returned once, causing ``NotImplementedError`` afterwards. But when ``repeat=True`` is given, that result is returned on every call, forever. Similarly, iterable results are normally exhausted once, but when this setting is enabled, they are wrapped in `itertools.cycle`. Default: ``True``. :raises: ``TypeError``, if the values given to ``run`` or other kwargs aren't of the expected types. .. versionchanged:: 1.5 Added support for boolean and string result values. .. versionchanged:: 1.5 Added support for regex dict keys. .. versionchanged:: 1.5 Added the ``repeat`` keyword argument. .. versionchanged:: 2.0 Changed ``repeat`` default value from ``False`` to ``True``. """ # Set up like any other Context would, with the config super().__init__(config) # Pull out behavioral kwargs self._set("__repeat", kwargs.pop("repeat", True)) # The rest must be things like run/sudo - mock Context method info for method, results in kwargs.items(): # For each possible value type, normalize to iterable of Result # objects (possibly repeating). singletons = (Result, bool, str) if isinstance(results, dict): for key, value in results.items(): results[key] = self._normalize(value) elif isinstance(results, singletons) or hasattr( results, "__iter__" ): results = self._normalize(results) # Unknown input value: cry else: err = "Not sure how to yield results from a {!r}" raise TypeError(err.format(type(results))) # Save results for use by the method self._set("__{}".format(method), results) # Wrap the method in a Mock self._set(method, Mock(wraps=getattr(self, method))) def _normalize(self, value: Any) -> Iterator[Any]: # First turn everything into an iterable if not hasattr(value, "__iter__") or isinstance(value, str): value = [value] # Then turn everything within into a Result results = [] for obj in value: if isinstance(obj, bool): obj = Result(exited=0 if obj else 1) elif isinstance(obj, str): obj = Result(obj) results.append(obj) # Finally, turn that iterable into an iteratOR, depending on repeat return cycle(results) if getattr(self, "__repeat") else iter(results) # TODO: _maybe_ make this more metaprogrammy/flexible (using __call__ etc)? # Pretty worried it'd cause more hard-to-debug issues than it's presently # worth. Maybe in situations where Context grows a _lot_ of methods (e.g. # in Fabric 2; though Fabric could do its own sub-subclass in that case...) def _yield_result(self, attname: str, command: str) -> Result: try: obj = getattr(self, attname) # Dicts need to try direct lookup or regex matching if isinstance(obj, dict): try: obj = obj[command] except KeyError: # TODO: could optimize by skipping this if not any regex # objects in keys()? for key, value in obj.items(): if hasattr(key, "match") and key.match(command): obj = value break else: # Nope, nothing did match. raise KeyError # Here, the value was either never a dict or has been extracted # from one, so we can assume it's an iterable of Result objects due # to work done by __init__. result: Result = next(obj) # Populate Result's command string with what matched unless # explicitly given if not result.command: result.command = command return result except (AttributeError, IndexError, KeyError, StopIteration): # raise_from(NotImplementedError(command), None) raise NotImplementedError(command) def run(self, command: str, *args: Any, **kwargs: Any) -> Result: # TODO: perform more convenience stuff associating args/kwargs with the # result? E.g. filling in .command, etc? Possibly useful for debugging # if one hits unexpected-order problems with what they passed in to # __init__. return self._yield_result("__run", command) def sudo(self, command: str, *args: Any, **kwargs: Any) -> Result: # TODO: this completely nukes the top-level behavior of sudo(), which # could be good or bad, depending. Most of the time I think it's good. # No need to supply dummy password config, etc. # TODO: see the TODO from run() re: injecting arg/kwarg values return self._yield_result("__sudo", command) def set_result_for( self, attname: str, command: str, result: Result ) -> None: """ Modify the stored mock results for given ``attname`` (e.g. ``run``). This is similar to how one instantiates `MockContext` with a ``run`` or ``sudo`` dict kwarg. For example, this:: mc = MockContext(run={'mycommand': Result("mystdout")}) assert mc.run('mycommand').stdout == "mystdout" is functionally equivalent to this:: mc = MockContext() mc.set_result_for('run', 'mycommand', Result("mystdout")) assert mc.run('mycommand').stdout == "mystdout" `set_result_for` is mostly useful for modifying an already-instantiated `MockContext`, such as one created by test setup or helper methods. .. versionadded:: 1.0 """ attname = "__{}".format(attname) heck = TypeError( "Can't update results for non-dict or nonexistent mock results!" ) # Get value & complain if it's not a dict. # TODO: should we allow this to set non-dict values too? Seems vaguely # pointless, at that point, just make a new MockContext eh? try: value = getattr(self, attname) except AttributeError: raise heck if not isinstance(value, dict): raise heck # OK, we're good to modify, so do so. value[command] = self._normalize(result)