issue #106: support WANT_JSON modules.

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,
             )

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,

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 = '<<INCLUDE_ANSIBLE_MODULE_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().