From d0d0c258d65a29e674e8cbe7530b4fa45da6c4c7 Mon Sep 17 00:00:00 2001
From: Darko Poljak <darko.poljak@gmail.com>
Date: Fri, 18 Mar 2016 23:34:02 +0100
Subject: [PATCH] Add ssh mux options by default if available.

---
 docs/man/cdist-reference.text.sh |  6 ++++
 docs/man/man1/cdist.text         |  5 +++
 scripts/cdist                    | 60 ++++++++++++++++++--------------
 3 files changed, 44 insertions(+), 27 deletions(-)

diff --git a/docs/man/cdist-reference.text.sh b/docs/man/cdist-reference.text.sh
index 0a9d76ef..0aaaec0a 100755
--- a/docs/man/cdist-reference.text.sh
+++ b/docs/man/cdist-reference.text.sh
@@ -249,6 +249,12 @@ CDIST_OVERRIDE::
 CDIST_ORDER_DEPENDENCY::
     Create dependencies based on the execution order (see cdist-manifest(7))
 
+CDIST_REMOTE_EXEC::
+    Use this command for remote execution (should behave like ssh)
+
+CDIST_REMOTE_COPY::
+    Use this command for remote copy (should behave like scp)
+
 SEE ALSO
 --------
 - cdist(1)
diff --git a/docs/man/man1/cdist.text b/docs/man/man1/cdist.text
index c09d8f41..e29ae3ae 100644
--- a/docs/man/man1/cdist.text
+++ b/docs/man/man1/cdist.text
@@ -138,6 +138,11 @@ CDIST_LOCAL_SHELL::
 CDIST_REMOTE_SHELL::
     Selects shell for remote scirpt execution, defaults to /bin/sh
 
+CDIST_REMOTE_EXEC::
+    Use this command for remote execution (should behave like ssh)
+
+CDIST_REMOTE_COPY::
+    Use this command for remote copy (should behave like scp)
 
 EXIT STATUS
 -----------
diff --git a/scripts/cdist b/scripts/cdist
index 8aa998d2..f5629641 100755
--- a/scripts/cdist
+++ b/scripts/cdist
@@ -21,33 +21,27 @@
 #
 #
 
-def inspect_ssh_mux_opts():
-    import subprocess
-
+def inspect_ssh_mux_opts(control_path_dir="~/.ssh/"):
     """Inspect whether or not ssh supports multiplexing options"""
+    import subprocess
+    import os
+
+    control_path = os.path.join(control_path_dir, "cdist.master-%l-%r@%h:%p")
     wanted_mux_opts = {
-        "ControlPath":"~/.ssh/master-%l-%r@%h:%p",
-        "ControlMaster":"auto",
-        "ControlPersist":"125",
+        "ControlPath": control_path,
+        "ControlMaster": "auto",
+        "ControlPersist": "125",
     }
-    # if checked key option is present then this assumes
-    # all options in value are present
-    check = {
-        "ControlMaster": ("ControlMaster", "ControlPath"),
-        "ControlPersist": ("ControlPersist",),
-    }
-    mux_opts = {}
-    for x in check:
-        try:
-            subprocess.check_output("ssh -o {}".format(x),
-                    stderr=subprocess.STDOUT, shell=True)
-        except subprocess.CalledProcessError as e:
-            foo = e.output.decode().lower()
-            if not "bad configuration option" in foo:
-                for o in check[x]:
-                    mux_opts[o] = wanted_mux_opts[o]
-    foo = [" -o {}={}".format(x, mux_opts[x]) for x in mux_opts]
-    return " ".join(foo)
+    mux_opts = " ".join([" -o {}={}".format(x,
+        wanted_mux_opts[x]) for x in wanted_mux_opts])
+    try:
+        subprocess.check_output("ssh {}".format(mux_opts),
+                stderr=subprocess.STDOUT, shell=True)
+    except subprocess.CalledProcessError as e:
+        subproc_output = e.output.decode().lower()
+        if "bad configuration option" in subproc_output:
+            return ""
+    return mux_opts
 
 def commandline():
     """Parse command line"""
@@ -56,6 +50,9 @@ def commandline():
     import cdist.banner
     import cdist.config
     import cdist.shell
+    import tempfile
+    import shutil
+    import os
 
     # Construct parser others can reuse
     parser = {}
@@ -83,9 +80,18 @@ def commandline():
 
     # Config
     # inspect multiplexing options for default remote copy/exec scp/ssh
-    MUX_OPTS = inspect_ssh_mux_opts()
-    cdist.REMOTE_EXEC += MUX_OPTS
-    cdist.REMOTE_COPY += MUX_OPTS
+    # but not if env vars are present
+    has_env_remote_exec = "CDIST_REMOTE_EXEC" in os.environ
+    has_env_remote_copy = "CDIST_REMOTE_COPY" in os.environ
+    if not has_env_remote_exec or not has_env_remote_copy:
+        control_path_dir = tempfile.mkdtemp(prefix="cdist.control.path")
+        import atexit
+        atexit.register(lambda: shutil.rmtree(control_path_dir))
+        MUX_OPTS = inspect_ssh_mux_opts(control_path_dir)
+        if not has_env_remote_exec:
+            cdist.REMOTE_EXEC += MUX_OPTS
+        if not has_env_remote_copy:
+            cdist.REMOTE_COPY += MUX_OPTS
 
     parser['config'] = parser['sub'].add_parser('config',
         parents=[parser['loglevel']])