Add support for mitogen

Compared to the Vanilla Ansible, Mitogen calls `sshwrapper.py` differently. This imply a different parsing to extract the options, host and command. Moreover, Mitogen does not passes the `remote_user` therefore we must set as well `BASTION_ANSIBLE_REMOTE_USER` env var. Apart from setting the `BASTION_ANSIBLE_REMOTE_USER` the wrapper works for both vanilla and mitogen connection. Users should refer to mitogen to learn about the installation process. The following commit has been tested with: * Test case #1 * Mitogen 0.3.7 * Ansible 210.8 * Python 3.9.2 * Debian 11.9 * Test case #2 * Mitogen 0.3.7 * Ansible 2.16.6 * Python 3.12 Signed-off-by: Wilfried Roset <[email protected]>
ovh · Apr 26, 2024 · 0364e20 · 0364e20
1 parent ac65f1b
commit 0364e20
Show file tree

Hide file tree

Showing 3 changed files with 236 additions and 20 deletions.
diff --git a/lib.py b/lib.py
@@ -3,6 +3,7 @@
 import os
 import subprocess
 import time
+from shlex import quote
 
 from yaml import YAMLError, safe_load
 
@@ -119,7 +120,9 @@ def get_hostvars(host) -> dict:
     return {}
 
 
-def manage_conf_file(conf_file, bastion_host, bastion_port, bastion_user):
+def manage_conf_file(
+    conf_file, bastion_host, bastion_port, bastion_user, bastion_ansible_remote_user
+):
     """Fetch the bastion vars from a config file.
 
     There will be set if not already defined, and before looking in the ansible inventory
@@ -137,11 +140,15 @@ def manage_conf_file(conf_file, bastion_host, bastion_port, bastion_user):
                     bastion_port = yaml_conf.get("bastion_port")
                 if not bastion_user:
                     bastion_user = yaml_conf.get("bastion_user")
+                if not bastion_ansible_remote_user:
+                    bastion_ansible_remote_user = yaml_conf.get(
+                        "bastion_ansible_remote_user"
+                    )
 
         except (YAMLError, IOError) as e:
             print("Error loading yaml file: {}".format(e))
 
-    return bastion_host, bastion_port, bastion_user
+    return bastion_host, bastion_port, bastion_user, bastion_ansible_remote_user
 
 
 def get_var_within(my_value, hostvar, check_list=None):
@@ -250,3 +257,43 @@ def get_bastion_vars(host_vars):
         "bastion_port": bastion_port,
         "bastion_user": bastion_user,
     }
+
+
+def parse_ansible_command(args):
+    options = []
+    cmd = ""
+    host = ""
+    i = 0
+    cmd_is_next = False
+    while i < len(args):
+        # -o options: Can be used to give options in the format used in the configuration file, followed by 1 argument
+        # -l login name, followed by 1 argument
+        if args[i] in ["-o", "-l"]:
+            options.append(args[i])
+            options.append(args[i + 1])
+            i = i + 2
+        # Collect everything option as standalone
+        # Example:
+        # -C Requests compression of all data, not followed by an argument
+        elif args[i].startswith("-"):
+            options.append(args[i])
+            i = i + 1
+        elif not host:
+            host = args[i]
+            i = i + 1
+            cmd_is_next = True
+        elif cmd_is_next:
+            # The cmd is the last elem but in two forms
+            # vanilla ansible passes the cmd as a single string
+            # mitogen passes the command as an array
+            # let's normalize
+            cmd = args[i:]
+            if isinstance(cmd, list):
+                cmd = " ".join(cmd)
+            break
+
+    # the wrapper expects a shell
+    if not cmd.startswith("/bin/sh"):
+        # note that we efficiently quote the cmd
+        cmd = " ".join(["/bin/sh", "-c", quote(cmd)])
+    return options, cmd, host
diff --git a/sshwrapper.py b/sshwrapper.py
@@ -11,6 +11,7 @@
     get_hostvars,
     get_var_within,
     manage_conf_file,
+    parse_ansible_command,
 )
 
 
@@ -20,12 +21,12 @@ def main():
     bastion_user = None
     bastion_host = None
     bastion_port = None
+    bastion_ansible_remote_user = None
     remote_user = None
     remote_port = 22
     default_configuration_file = "/etc/ovh/bastion/config.yml"
 
-    cmd = argv.pop()
-    host = argv.pop()
+    options, cmd, host = parse_ansible_command(argv)
 
     # check if bastion_vars are passed as env vars in the playbook
     # may be usefull if the ansible controller manage many bastions
@@ -38,28 +39,35 @@ def main():
     #     BASTION_PORT: "{{ bastion_port }}"
     #
     # will result as : ... '/bin/sh -c '"'"'BASTION_USER=my_bastion_user BASTION_HOST=my_bastion_host BASTION_PORT=22 /usr/bin/python3 && sleep 0'"'"''
-    for i in list(cmd.split(" ")):
+    for i in list(cmd):
         if "bastion_user" in i.lower():
             bastion_user = i.split("=")[1]
         elif "bastion_host" in i.lower():
             bastion_host = i.split("=")[1]
         elif "bastion_port" in i.lower():
             bastion_port = i.split("=")[1]
+        elif "bastion_ansible_remote_user" in i.lower():
+            bastion_ansible_remote_user = i.split("=")[1]
 
     # in some cases (AWX in a non containerised environment for instance), the environment is overridden by the job
     # so we are not able to get the BASTION vars
     # if some vars are still undefined, try to load them from a configuration file
-    bastion_host, bastion_port, bastion_user = manage_conf_file(
+    (
+        bastion_host,
+        bastion_port,
+        bastion_user,
+        bastion_ansible_remote_user,
+    ) = manage_conf_file(
         os.environ.get("BASTION_CONF_FILE", default_configuration_file),
         bastion_host,
         bastion_port,
         bastion_user,
+        bastion_ansible_remote_user,
     )
 
     # lookup on the inventory may take some time, depending on the source, so use it only if not defined elsewhere
     # it seems like some module like template does not send env vars too...
     if not bastion_host or not bastion_port or not bastion_user:
-
         # check if running on AWX, we'll get the vars in a different way
         awx_inventory_file = awx_get_inventory_file()
         if os.path.exists(awx_inventory_file):
@@ -81,16 +89,25 @@ def main():
         bastion_host = get_var_within(
             hostvar.get("bastion_host", os.environ.get("BASTION_HOST")), hostvar
         )
+        bastion_ansible_remote_user = get_var_within(
+            hostvar.get(
+                "bastion_ansible_remote_user",
+                os.environ.get("BASTION_ANSIBLE_REMOTE_USER"),
+            ),
+            hostvar,
+        )
 
-    for i, e in enumerate(argv):
-
+    for i, e in enumerate(options):
         if e.startswith("User="):
             remote_user = e.split("=")[-1]
             argv[i] = "User={}".format(bastion_user)
         elif e.startswith("Port="):
             remote_port = e.split("=")[-1]
             argv[i] = "Port={}".format(bastion_port)
 
+    if not remote_user:
+        remote_user = bastion_ansible_remote_user
+
     # syscall exec
     args = (
         [
@@ -105,7 +122,7 @@ def main():
             bastion_host,
             "-T",
         ]
-        + argv
+        + options
         + [
             "--",
             "-q",

diff --git a/tests.py b/tests.py
@@ -7,52 +7,85 @@
     get_bastion_vars,
     get_var_within,
     manage_conf_file,
+    parse_ansible_command,
 )
 
 BASTION_HOST = "my_bastion"
 BASTION_PORT = 22
 BASTION_USER = "my_bastion_user"
+BASTION_ANSIBLE_REMOTE_USER = "my_ansible_remote_user"
 BASTION_CONF_FILE = "/tmp/test_bastion_conf_file.yml"
 
 
 def test_manage_conf_file_bastion_host_undefined():
-    bastion_host, bastion_port, bastion_user = manage_conf_file(
-        BASTION_CONF_FILE, None, BASTION_PORT, BASTION_USER
+    (
+        bastion_host,
+        bastion_port,
+        bastion_user,
+        bastion_ansible_remote_user,
+    ) = manage_conf_file(
+        BASTION_CONF_FILE,
+        None,
+        BASTION_PORT,
+        BASTION_USER,
+        BASTION_ANSIBLE_REMOTE_USER,
     )
     assert bastion_host == BASTION_HOST
 
 
 def test_manage_conf_file_bastion_port_undefined():
-    bastion_host, bastion_port, bastion_user = manage_conf_file(
-        BASTION_CONF_FILE, BASTION_HOST, None, BASTION_USER
+    (
+        bastion_host,
+        bastion_port,
+        bastion_user,
+        bastion_ansible_remote_user,
+    ) = manage_conf_file(
+        BASTION_CONF_FILE,
+        BASTION_HOST,
+        None,
+        BASTION_USER,
+        BASTION_ANSIBLE_REMOTE_USER,
     )
     assert bastion_port == BASTION_PORT
 
 
 def test_manage_conf_file_bastion_user_undefined():
-    bastion_host, bastion_port, bastion_user = manage_conf_file(
-        BASTION_CONF_FILE, BASTION_HOST, BASTION_PORT, None
+    (
+        bastion_host,
+        bastion_port,
+        bastion_user,
+        bastion_ansible_remote_user,
+    ) = manage_conf_file(
+        BASTION_CONF_FILE,
+        BASTION_HOST,
+        BASTION_PORT,
+        None,
+        BASTION_ANSIBLE_REMOTE_USER,
     )
     assert bastion_user == BASTION_USER
 
 
 def test_manage_conf_file_bastion_all_undefined():
     write_conf_file(BASTION_CONF_FILE)
-    bastion_host, bastion_port, bastion_user = manage_conf_file(
-        BASTION_CONF_FILE, None, None, None
-    )
+    (
+        bastion_host,
+        bastion_port,
+        bastion_user,
+        bastion_ansible_remote_user,
+    ) = manage_conf_file(BASTION_CONF_FILE, None, None, None, None)
     assert bastion_user == BASTION_USER
     assert bastion_port == BASTION_PORT
     assert bastion_host == BASTION_HOST
+    assert bastion_ansible_remote_user == BASTION_ANSIBLE_REMOTE_USER
 
 
 def write_conf_file(conf_file):
     with open(conf_file, "w") as f:
-
         data = {
             "bastion_host": BASTION_HOST,
             "bastion_port": BASTION_PORT,
             "bastion_user": BASTION_USER,
+            "bastion_ansible_remote_user": BASTION_ANSIBLE_REMOTE_USER,
         }
 
         dump(data, f)
@@ -133,3 +166,122 @@ def test_get_bastion_vars_not_full():
     host_vars = {"bastion_port": BASTION_PORT, "bastion_user": BASTION_USER}
     bastion_vars = get_bastion_vars(host_vars)
     assert not bastion_vars["bastion_host"]
+
+
+def test_parse_ansible_command():
+    cases = [
+        # Ansible
+        (
+            [
+                "-C",
+                "-o",
+                "ControlMaster=auto",
+                "-o",
+                "ControlPersist=60s",
+                "-o",
+                "StrictHostKeyChecking=no",
+                "-o",
+                "KbdInteractiveAuthentication=no",
+                "-o",
+                "PreferredAuthentications=gssapi-with-mic,gssapi-keyex,hostbased,publickey",
+                "-o",
+                "PasswordAuthentication=no",
+                "-o",
+                'User="root"',
+                "-o",
+                "ConnectTimeout=10",
+                "-o",
+                'ControlPath="/some/control/path"',
+                "my-secured-host",
+                "/bin/sh -c '/usr/bin/python3 && sleep 0'",
+            ],
+            [
+                "-C",
+                "-o",
+                "ControlMaster=auto",
+                "-o",
+                "ControlPersist=60s",
+                "-o",
+                "StrictHostKeyChecking=no",
+                "-o",
+                "KbdInteractiveAuthentication=no",
+                "-o",
+                "PreferredAuthentications=gssapi-with-mic,gssapi-keyex,hostbased,publickey",
+                "-o",
+                "PasswordAuthentication=no",
+                "-o",
+                'User="root"',
+                "-o",
+                "ConnectTimeout=10",
+                "-o",
+                'ControlPath="/some/control/path"',
+            ],
+            "/bin/sh -c '/usr/bin/python3 && sleep 0'",
+            "my-secured-host",
+        ),
+        # Mitogen
+        (
+            [
+                "-o",
+                "LogLevel ERROR",
+                "-l",
+                "root",
+                "-o",
+                "Compression yes",
+                "-o",
+                "ServerAliveInterval 30",
+                "-o",
+                "ServerAliveCountMax 10",
+                "-o",
+                "BatchMode yes",
+                "-o",
+                "StrictHostKeyChecking no",
+                "-o",
+                "UserKnownHostsFile /dev/null",
+                "-o",
+                "GlobalKnownHostsFile /dev/null",
+                "-C",
+                "-o",
+                "ControlMaster=auto",
+                "-o",
+                "ControlPersist=60s",
+                "my-secured-host",
+                "/usr/bin/python3",
+                "-c",
+                "'import sys;sys.path=[p for p in sys.path if p];import binascii,os,zlib;exec(zlib.decompress(binascii.a2b_base64(foobar)))'",
+            ],
+            [
+                "-o",
+                "LogLevel ERROR",
+                "-l",
+                "root",
+                "-o",
+                "Compression yes",
+                "-o",
+                "ServerAliveInterval 30",
+                "-o",
+                "ServerAliveCountMax 10",
+                "-o",
+                "BatchMode yes",
+                "-o",
+                "StrictHostKeyChecking no",
+                "-o",
+                "UserKnownHostsFile /dev/null",
+                "-o",
+                "GlobalKnownHostsFile /dev/null",
+                "-C",
+                "-o",
+                "ControlMaster=auto",
+                "-o",
+                "ControlPersist=60s",
+            ],
+            "/bin/sh -c '/usr/bin/python3 -c '\"'\"'import sys;sys.path=[p for p in sys.path if p];import binascii,os,zlib;exec(zlib.decompress(binascii.a2b_base64(foobar)))'\"'\"''",
+            "my-secured-host",
+        ),
+    ]
+
+    for args, expected_options, expected_cmd, expected_host in cases:
+        options, cmd, host = parse_ansible_command(args)
+        assert options == expected_options
+        assert cmd == expected_cmd
+        assert host == expected_host