From 16b64392e2092cad084c38356aed5a6c574baeb2 Mon Sep 17 00:00:00 2001 From: David Wilson Date: Sun, 1 Apr 2018 18:19:34 +0100 Subject: [PATCH] issue #106: support WANT_JSON modules. --- ansible_mitogen/mixins.py | 2 + ansible_mitogen/planner.py | 61 +++++++++++++- ansible_mitogen/runner.py | 157 +++++++++++++++++++++++++++---------- 3 files changed, 173 insertions(+), 47 deletions(-) diff --git a/ansible_mitogen/mixins.py b/ansible_mitogen/mixins.py index 831d77eb..013e6b81 100644 --- a/ansible_mitogen/mixins.py +++ b/ansible_mitogen/mixins.py @@ -310,6 +310,8 @@ class ActionModuleMixin(ansible.plugins.action.ActionBase): connection=self._connection, module_name=mitogen.utils.cast(module_name), module_args=mitogen.utils.cast(module_args), + task_vars=task_vars, + templar=self._templar, env=mitogen.utils.cast(env), wrap_async=wrap_async, ) diff --git a/ansible_mitogen/planner.py b/ansible_mitogen/planner.py index 11b9d7a6..ec14fdcc 100644 --- a/ansible_mitogen/planner.py +++ b/ansible_mitogen/planner.py @@ -55,13 +55,41 @@ import ansible_mitogen.services LOG = logging.getLogger(__name__) +def parse_script_interpreter(source): + """ + Extract the script interpreter and its sole argument from the module + source code. + + :returns: + Tuple of `(interpreter, arg)`, where `intepreter` is the script + interpreter and `arg` is its solve argument if present, otherwise + :py:data:`None`. + """ + # Linux requires first 2 bytes with no whitespace, pretty sure it's the + # same everywhere. See binfmt_script.c. + if not source.startswith('#!'): + return None, None + + # Find terminating newline. Assume last byte of binprm_buf if absent. + nl = source.find('\n', 0, 128) + if nl == -1: + nl = min(128, len(source)) + + # Split once on the first run of whitespace. If no whitespace exists, + # bits just contains the interpreter filename. + bits = source[2:nl].strip().split(None, 1) + if len(bits) == 1: + return bits[0], None + return bits[0], bits[1] + + class Invocation(object): """ Collect up a module's execution environment then use it to invoke helpers.run_module() or helpers.run_module_async() in the target context. """ def __init__(self, action, connection, module_name, module_args, - env, wrap_async): + task_vars, templar, env, wrap_async): #: ActionBase instance invoking the module. Required to access some #: output postprocessing methods that don't belong in ActionBase at #: all. @@ -73,6 +101,10 @@ class Invocation(object): self.module_name = module_name #: Final module arguments. self.module_args = module_args + #: Task variables, needed to extract ansible_*_interpreter. + self.task_vars = task_vars + #: Templar, needed to extract ansible_*_interpreter. + self.templar = templar #: Final module environment. self.env = env #: Boolean, if :py:data:`True`, launch the module asynchronously. @@ -129,6 +161,27 @@ class BinaryPlanner(Planner): } +class ScriptPlanner(BinaryPlanner): + """ + Common functionality for script module planners -- handle interpreter + detection and rewrite. + """ + def plan(self, invocation): + kwargs = super(ScriptPlanner, self).plan(invocation) + interpreter, arg = parse_script_interpreter(invocation.module_source) + shebang, _ = module_common._get_shebang( + interpreter=interpreter, + task_vars=invocation.task_vars, + templar=invocation.templar, + ) + if shebang: + interpreter = shebang[2:] + + kwargs['interpreter'] = interpreter + kwargs['interpreter_arg'] = arg + return kwargs + + class ReplacerPlanner(BinaryPlanner): """ The Module Replacer framework is the original framework implementing @@ -159,7 +212,7 @@ class ReplacerPlanner(BinaryPlanner): return module_common.REPLACER in invocation.module_source -class JsonArgsPlanner(BinaryPlanner): +class JsonArgsPlanner(ScriptPlanner): """ Script that has its interpreter directive and the task arguments substituted into its source as a JSON string. @@ -170,7 +223,7 @@ class JsonArgsPlanner(BinaryPlanner): return module_common.REPLACER_JSONARGS in invocation.module_source -class WantJsonPlanner(BinaryPlanner): +class WantJsonPlanner(ScriptPlanner): """ If a module has the string WANT_JSON in it anywhere, Ansible treats it as a non-native module that accepts a filename as its only command line @@ -224,7 +277,7 @@ class NativePlanner(Planner): _planners = [ # JsonArgsPlanner, - # WantJsonPlanner, + WantJsonPlanner, # ReplacerPlanner, BinaryPlanner, NativePlanner, diff --git a/ansible_mitogen/runner.py b/ansible_mitogen/runner.py index ad94ef7c..cf970a6b 100644 --- a/ansible_mitogen/runner.py +++ b/ansible_mitogen/runner.py @@ -37,6 +37,7 @@ how to build arguments for it, preseed related data, etc. from __future__ import absolute_import import json +import logging import os import tempfile @@ -52,6 +53,9 @@ import ansible.module_utils.basic ansible.module_utils.basic._ANSIBLE_ARGS = '{}' +LOG = logging.getLogger(__name__) + + class Runner(object): """ Ansible module runner. After instantiation (with kwargs supplied by the @@ -247,70 +251,52 @@ class NativeRunner(Runner): } -class BinaryRunner(Runner): +class ProgramRunner(Runner): def __init__(self, path, service_context, **kwargs): - print 'derp', kwargs - super(BinaryRunner, self).__init__(**kwargs) + super(ProgramRunner, self).__init__(**kwargs) self.path = path self.service_context = service_context def setup(self): - super(BinaryRunner, self).setup() + super(ProgramRunner, self).setup() self._setup_program() - self._setup_args() - - def _get_program(self): - """ - Fetch the module binary from the master if necessary. - """ - return ansible_mitogen.helpers.get_file( - context=self.service_context, - path=self.path, - ) - - def _get_args(self): - """ - Return the module arguments formatted as JSON. - """ - return json.dumps(self.args) def _setup_program(self): """ Create a temporary file containing the program code. The code is fetched via :meth:`_get_program`. """ - self.bin_fp = tempfile.NamedTemporaryFile( + self.program_fp = tempfile.NamedTemporaryFile( prefix='ansible_mitogen', suffix='-binary', ) - self.bin_fp.write(self._get_program()) - self.bin_fp.flush() - os.chmod(self.bin_fp.name, int('0700', 8)) + self.program_fp.write(self._get_program()) + self.program_fp.flush() + os.chmod(self.program_fp.name, int('0700', 8)) - def _setup_args(self): + def _get_program(self): """ - Create a temporary file containing the module's arguments. The - arguments are formatted via :meth:`_get_args`. + Fetch the module binary from the master if necessary. """ - self.args_fp = tempfile.NamedTemporaryFile( - prefix='ansible_mitogen', - suffix='-args', + return ansible_mitogen.helpers.get_file( + context=self.service_context, + path=self.path, ) - self.args_fp.write(self._get_args()) - self.args_fp.flush() + + def _get_program_args(self): + return [self.program_fp.name] def revert(self): """ - Delete the temporary binary and argument files. + Delete the temporary program file. """ - self.args_fp.close() - self.bin_fp.close() - super(BinaryRunner, self).revert() + super(ProgramRunner, self).revert() + self.program_fp.close() def _run(self): try: rc, stdout, stderr = ansible_mitogen.helpers.exec_args( - args=[self.bin_fp.name, self.args_fp.name], + args=self._get_program_args(), ) except Exception, e: return { @@ -326,15 +312,100 @@ class BinaryRunner(Runner): } -class WantJsonRunner(BinaryRunner): +class ArgsFileRunner(Runner): + def setup(self): + super(ArgsFileRunner, self).setup() + self._setup_args() + + def _setup_args(self): + """ + Create a temporary file containing the module's arguments. The + arguments are formatted via :meth:`_get_args`. + """ + self.args_fp = tempfile.NamedTemporaryFile( + prefix='ansible_mitogen', + suffix='-args', + ) + self.args_fp.write(self._get_args_contents()) + self.args_fp.flush() + + def _get_args_contents(self): + """ + Return the module arguments formatted as JSON. + """ + return json.dumps(self.args) + + def _get_program_args(self): + return [self.program_fp.name, self.args_fp.name] + + def revert(self): + """ + Delete the temporary argument file. + """ + super(ArgsFileRunner, self).revert() + self.args_fp.close() + + +class BinaryRunner(ArgsFileRunner, ProgramRunner): + pass + + +class ScriptRunner(ProgramRunner): + def __init__(self, interpreter, interpreter_arg, **kwargs): + super(ScriptRunner, self).__init__(**kwargs) + self.interpreter = interpreter + self.interpreter_arg = interpreter_arg + + b_ENCODING_STRING = b'# -*- coding: utf-8 -*-' + def _get_program(self): - s = super(WantJsonRunner, self)._get_program() - # fix up shebang. - return s + return self._rewrite_source( + super(ScriptRunner, self)._get_program() + ) + + def _rewrite_source(self, s): + """ + Mutate the source according to the per-task parameters. + """ + # Couldn't find shebang, so let shell run it, because shell assumes + # executables like this are just shell scripts. + LOG.debug('++++++++++++++ %s', self.interpreter) + if not self.interpreter: + return s + + shebang = '#!' + self.interpreter + if self.interpreter_arg: + shebang += ' ' + self.interpreter_arg + + new = [shebang] + if os.path.basename(self.interpreter).startswith('python'): + new.append(self.b_ENCODING_STRING) + + _, _, rest = s.partition('\n') + new.append(rest) + return '\n'.join(new) + + +class JsonArgsFileRunner(ScriptRunner): + JSON_ARGS = '<>' + + def _get_args_contents(self): + return json.dump(self.args) + + def _rewrite_source(self, s): + return ( + super(JsonArgsFileRunner, self)._rewrite_source(s) + .replace(self.JSON_ARGS, self._get_args_contents()) + ) + + +class WantJsonRunner(ArgsFileRunner, ScriptRunner): + pass + -class OldStyleRunner(BinaryRunner): - def _get_args(self): +class OldStyleRunner(ScriptRunner): + def _get_args_contents(self): """ Mimic the argument formatting behaviour of ActionBase._execute_module().