Skip to content

Commit

Permalink
Add support for mitogen
Browse files Browse the repository at this point in the history 
Compared to the Vanilla Ansible, [Mitogen](https://github.com/mitogen-hq/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
  * MacOS 14.2

This playbook works as expected with and without mitogen enabled
```
❯ cat test.yaml
---
- name: test
  hosts: test
  gather_facts: false
  tasks:
    - name: Run the equivalent of "apt-get update" as a separate step
      ansible.builtin.apt:
        update_cache: true

    - name: Create files with copy content module
      copy:
        content: |
          test file {{ item }}
        dest: /tmp/file_{{item}}
      with_sequence: start=1 end=10

    - name: demo template
      ansible.builtin.template:
        src: demo.txt.j2
        dest: /tmp/demo.txt
        mode: 0640
```
Here is the ansible.cfg
```
[defaults]
interpreter_python = /usr/bin/python3
host_key_checking = False
deprecation_warnings = False
syslog_facility = LOG_USER
bin_ansible_callbacks = True

gathering = explicit
callbacks_enabled = ansible.posix.profile_tasks

strategy_plugins = ./mitogen/ansible_mitogen/plugins/strategy/
strategy = mitogen_linear

[ssh_connection]
scp_if_ssh = False
pipelining = True
transfer_method = sftp
ssh_executable = ./bastion-ansible-wrapper/sshwrapper.py
sftp_executable = ./bastion-ansible-wrapper/sftpbastion.sh
retries = 1
```

Signed-off-by: Wilfried Roset <[email protected]>
  • Loading branch information
@wilfriedroset
wilfriedroset committed Apr 26, 2024
1 parent 5c7c446 commit df3bf51
Show file tree
Hide file tree
Showing 3 changed files with 236 additions and 20 deletions.
51 changes: 49 additions & 2 deletions lib.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import os
import subprocess
import time
from shlex import quote

from yaml import YAMLError, safe_load

Expand Down Expand Up @@ -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
Expand All @@ -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):
Expand Down Expand Up @@ -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
33 changes: 25 additions & 8 deletions sshwrapper.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
get_hostvars,
get_var_within,
manage_conf_file,
parse_ansible_command,
)


Expand All @@ -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
Expand All @@ -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):
Expand All @@ -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 = (
[
Expand All @@ -105,7 +122,7 @@ def main():
bastion_host,
"-T",
]
+ argv
+ options
+ [
"--",
"-q",
Expand Down
172 changes: 162 additions & 10 deletions tests.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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

0 comments on commit df3bf51

Please sign in to comment.