From c3da2d24c31c9505dfa1d5521904aeac54cac951 Mon Sep 17 00:00:00 2001 From: Wilfried Roset Date: Fri, 26 Apr 2024 09:51:25 +0200 Subject: [PATCH] Add support for mitogen Signed-off-by: Wilfried Roset --- lib.py | 51 +++++++++++++++- sshwrapper.py | 38 ++++++++---- tests.py | 158 +++++++++++++++++++++++++++++++++++++++++++++++--- 3 files changed, 224 insertions(+), 23 deletions(-) diff --git a/lib.py b/lib.py index 77aa2db..0eaf207 100644 --- 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 index 5b841d2..45d40a3 100755 --- 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,32 @@ 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( - os.environ.get("BASTION_CONF_FILE", default_configuration_file), - bastion_host, - bastion_port, - bastion_user, + 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,9 +86,15 @@ 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) @@ -91,6 +102,9 @@ def main(): 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 +119,7 @@ def main(): bastion_host, "-T", ] - + argv + + options + [ "--", "-q", diff --git a/tests.py b/tests.py index 643e6e0..06d827b 100644 --- a/tests.py +++ b/tests.py @@ -7,52 +7,73 @@ 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 +154,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