|
1 | 1 | #!/usr/bin/env python3 |
| 2 | +""" |
| 3 | +SSH to a cluster host using connection properties from Ansible inventory. |
2 | 4 |
|
3 | | -# This tool allows you to ssh into a host using the ansible inventory. |
4 | | -# Example: ansible-ssh compute[0] -o GlobalKnownHostsFile=/dev/null -o |
5 | | -# UserKnownHostsFile=/dev/null |
| 5 | +Usage: |
| 6 | + ansible-ssh [-n] PATTERN |
| 7 | +
|
| 8 | +where PATTERN can be any of: |
| 9 | + - host e.g. mycluster-login-0 |
| 10 | + - group e.g. login |
| 11 | + - group[index] e.g. compute[0] |
| 12 | +options: |
| 13 | + -n Disable strict host key checking and do not store accepted keys |
| 14 | +""" |
6 | 15 |
|
7 | 16 | import json |
8 | 17 | import os |
9 | 18 | import shlex |
10 | 19 | import subprocess |
11 | 20 | import sys |
12 | | -from collections import defaultdict |
13 | | - |
14 | | - |
15 | | -def _optional_arg(prototype, *values): |
16 | | - # returns empty string if any of the values are falsey |
17 | | - filtered = [value for value in values if value] |
18 | | - return prototype.format(*values) if len(values) == len(filtered) else "" |
19 | 21 |
|
| 22 | +NO_HOSTKEY_CHECK_OPTS = [ |
| 23 | + '-o', 'StrictHostKeyChecking=no', |
| 24 | + '-o', 'GlobalKnownHostsFile=/dev/null', |
| 25 | + '-o', 'UserKnownHostsFile=/dev/null' |
| 26 | +] |
20 | 27 |
|
21 | 28 | if __name__ == "__main__": |
22 | | - if len(sys.argv) < 2: |
23 | | - msg = ( |
24 | | - f"Usage: {sys.argv[0]} <inventory_hostname> [args to pass to ssh]") |
25 | | - print(msg, file=sys.stderr) |
| 29 | + no_known_hosts = '-n' in sys.argv |
| 30 | + if sys.argv[-1] in (sys.argv[0], '-n'): |
| 31 | + print(__doc__, file=sys.stderr) |
26 | 32 | sys.exit(-1) |
27 | 33 |
|
28 | | - # Quote to prevent shell injection |
29 | | - host = shlex.quote(sys.argv[1]) |
| 34 | + pattern = sys.argv[-1] |
30 | 35 |
|
| 36 | + template_str = ("ssh " |
| 37 | + "{% if ansible_ssh_port | default(false) %}-p {{ ansible_ssh_port }} {% endif %}" |
| 38 | + "{% if ansible_ssh_private_key_file | default(false) %} -i {{ ansible_ssh_private_key_file }} {% endif %}" |
| 39 | + "{{ ansible_ssh_common_args | default('') }} " |
| 40 | + "{{ ansible_user }}@{{ ansible_host }}") |
| 41 | + module_args = json.dumps({'msg':template_str}) |
| 42 | + ansible_cmd = ['ansible', pattern, '-o', '-m', 'debug', '-a', module_args] |
31 | 43 | try: |
32 | | - output = subprocess.check_output( |
33 | | - f'ansible-inventory --host {host}', shell=True) |
34 | | - except (subprocess.CalledProcessError) as e: |
35 | | - msg = (f"[ERROR]: Is {host} missing from the inventory?") |
| 44 | + output = subprocess.check_output(ansible_cmd, text=True) |
| 45 | + except (subprocess.CalledProcessError): |
| 46 | + msg = ("[ERROR]: ansible exited in error") |
36 | 47 | print(msg, file=sys.stderr) |
37 | 48 | sys.exit(-1) |
| 49 | + # output looks like e.g. |
| 50 | + # stg-login-0 | SUCCESS => { "changed": false, ... |
| 51 | + # one line per host |
| 52 | + if not output: |
| 53 | + sys.exit(1) |
| 54 | + result = output.splitlines()[0].split('>', 1)[-1] |
| 55 | + expanded = json.loads(result)['msg'] |
| 56 | + # can assume ansible_host exists b/c defined by terraform |
| 57 | + # can assume ansible_user exists b/c defined in common inventory |
38 | 58 |
|
39 | | - meta = defaultdict(str, json.loads(output)) |
40 | | - |
41 | | - ansible_ssh_host = meta['ansible_ssh_host'] or meta['ansible_host'] |
42 | | - ansible_ssh_user = meta['ansible_ssh_user'] or meta['ansible_user'] |
43 | | - ansible_ssh_port = meta['ansible_ssh_port'] |
44 | | - ansible_ssh_private_key_file = meta['ansible_ssh_private_key_file'] |
45 | | - |
46 | | - port = _optional_arg("-p {}", ansible_ssh_port) |
47 | | - identity = _optional_arg("-i {}", ansible_ssh_private_key_file) |
48 | | - host = _optional_arg("{}@{}", ansible_ssh_user, ansible_ssh_host) |
49 | | - opts = meta['ansible_ssh_common_args'] |
50 | | - |
51 | | - # Handle case where user is not set |
52 | | - if not host: |
53 | | - host = ansible_ssh_host |
54 | | - |
55 | | - if not host: |
56 | | - # if we get here, "ansible_ssh_host" is not set. |
57 | | - msg = f"Could not determine the host" |
58 | | - print(msg, file=sys.stderr) |
59 | | - sys.exit(-1) |
| 59 | + cmd = shlex.split(expanded) |
| 60 | + if no_known_hosts: |
| 61 | + cmd = cmd[0:1] + NO_HOSTKEY_CHECK_OPTS + cmd[1:] |
60 | 62 |
|
61 | | - base = shlex.split(f'ssh {port} {identity} {opts}') |
62 | | - extras = sys.argv[2:] |
63 | | - cmd = base + extras + [host] |
64 | 63 | print(f"[INFO]: Running: {subprocess.list2cmdline(cmd)}") |
65 | 64 | os.execvp(cmd[0], cmd) |
0 commit comments