feat: add specific code path for AWX (#26)

* feat: add specific code path for AWX

The ansible bastion wrapper default lookup to get bastion vars, is to run "ansible-inventory --list"
A downside of this way of doing, is that the whole inventory is going to be evaluated, and if we are using some custom vars plugins, there are executed to.
It can end up being very time consuming.

When using AWX, the whole inventory is available in the AWX Execution Environment as a file.
So it is much easier and faster to get the host associated to the requested ip (the ip sent by Ansible to the ssh wrapper).
Then, we can look for the bastion vars in the host vars, and if not found, execute "ansible-inventory --host" on the specific host, instead of the whole inventory.

* remove isort, as it conflicts with black to maange import with style

---------

Co-authored-by: Damien Cavagnini <[email protected]>

.pre-commit-config.yaml

@@ -10,10 +10,6 @@ repos:
         args: ['--remove']
       - id: requirements-txt-fixer
       - id: trailing-whitespace
-  - repo: https://github.com/PyCQA/isort
-    rev: 5.12.0
-    hooks:
-      - id: isort
   - repo: https://github.com/psf/black
     rev: 22.3.0
     hooks:

README.md

@@ -124,6 +124,26 @@ You may define multiple inventories sources in an ENV var. Example:
 export BASTION_ANSIBLE_INV_OPTIONS='-i my_first_inventory_source -i my_second_inventory_source'
 ```
 
+## Using the bastion wrapper with AWX
+
+When using AWX, the inventory is available as a file in the AWX Execution Environment.
+It is then easy and much faster to get the appropriate host from the IP sent by Ansible to the bastion wrapper.
+
+When AWX usage is detected, the bastion wrapper is going to:
+- lookup in the inventory file for the appropriate host
+- lookup for the bastion vars in the host_vars
+- if not found, run an inventory lookup on the host to get the group_vars too (and execute eventual vars plugins)
+
+The AWX usage is detected by looking for the inventory file, the default path being "/runner/inventory/hosts"
+The path may be changed y setting an "AWX_RUN_DIR" environment variable on the AWX worker.
+Ex on a AWX k8s instance group:
+```
+      env:
+      - name: "AWX_RUN_DIR"
+        value: "/my_folder/my_sub_folder"
+```
+The inventory file will be looked up at "/my_folder/my_sub_folder/inventory/hosts"
+
 ## Connection via SSH
 
 The wrapper can be configured using `ansible.cfg` file as follow:

lib.py

@@ -56,24 +56,11 @@ def get_inventory():
     inventory_options = os.environ.get("BASTION_ANSIBLE_INV_OPTIONS", "")
 
     command = "{} {} --list".format(inventory_cmd, inventory_options)
-    p = subprocess.Popen(
-        command,
-        shell=True,
-        stdin=subprocess.PIPE,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE,
-    )
-    output, error = p.communicate()
-    if type(output) is bytes:
-        output = output.decode()
-    if not p.returncode:
-        inventory = json.loads(output)
-        if cache_file:
-            write_inventory_to_cache(cache_file=cache_file, inventory=inventory)
-        return inventory
-    else:
-        logging.error(error)
-        raise Exception("failed to query ansible-inventory")
+    inventory = get_inv_from_command(command)
+    if cache_file:
+        write_inventory_to_cache(cache_file=cache_file, inventory=inventory)
+
+    return inventory
 
 
 def get_inventory_from_cache(cache_file, cache_timeout):
@@ -114,7 +101,7 @@ def write_inventory_to_cache(cache_file, inventory):
         json.dump(cache, fd)
 
 
-def get_hostvars(host):
+def get_hostvars(host) -> dict:
     """Fetch hostvars for the given host
 
     Ansible either uses the "ansible_host" inventory variable or the hostname.
@@ -192,3 +179,74 @@ def get_var_within(my_value, hostvar, check_list=None):
             return ""
 
     return my_value
+
+
+def get_inv_from_command(command):
+    p = subprocess.Popen(
+        command,
+        shell=True,
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+    output, error = p.communicate()
+    if isinstance(output, bytes):
+        output = output.decode()
+    if not p.returncode:
+        inventory = json.loads(output)
+        return inventory
+    else:
+        logging.error(error)
+        raise Exception("failed to get inventory")
+
+
+def awx_get_inventory_file():
+    # awx execution environment run dir, where the project and inventory are copied
+    default_run_dir = "/runner"
+    run_dir = os.environ.get("AWX_RUN_DIR", default_run_dir)
+    return "{}/inventory/hosts".format(run_dir)
+
+
+def awx_get_vars(host_ip, inventory_file):
+    # the inventory file is a script that print the inventory in json format
+    inv = get_inv_from_command(inventory_file)
+
+    # the ssh command sent only the IP to the ansible bastion wrapper.
+    # We are looking for the host which "ansible_host" has the same ip, then try to fetch the required vars from
+    # its host_vars
+    host = None
+    for k, v in inv.get("_meta", {}).get("hostvars", {}).items():
+        if v.get("ansible_host") == host_ip:
+            host = k
+            host_vars = v
+            break
+
+    # this should not happen
+    if not host:
+        return {}
+
+    bastion_vars = get_bastion_vars(host_vars)
+
+    if None not in [
+        bastion_vars.get("bastion_host"),
+        bastion_vars.get("bastion_port"),
+        bastion_vars.get("bastion_user"),
+    ]:
+        return bastion_vars
+
+    # if some bastion vars are missing, maybe they are defined as group_vars.
+    # We do an inventory lookup to get them.
+    # With AWX no need to list the whole inventory, we already know the host
+    command = "ansible-inventory -i {} --host {}".format(inventory_file, host)
+    return get_inv_from_command(command)
+
+
+def get_bastion_vars(host_vars):
+    bastion_host = host_vars.get("bastion_host")
+    bastion_user = host_vars.get("bastion_user")
+    bastion_port = host_vars.get("bastion_port")
+    return {
+        "bastion_host": bastion_host,
+        "bastion_port": bastion_port,
+        "bastion_user": bastion_user,
+    }

sshwrapper.py

@@ -4,7 +4,14 @@
 import os
 import sys
 
-from lib import find_executable, get_hostvars, get_var_within, manage_conf_file
+from lib import (
+    awx_get_inventory_file,
+    awx_get_vars,
+    find_executable,
+    get_hostvars,
+    get_var_within,
+    manage_conf_file,
+)
 
 
 def main():
@@ -52,8 +59,16 @@ def main():
     # 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:
-        hostvar = get_hostvars(host)  # dict
 
+        # 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):
+            hostvar = awx_get_vars(host, awx_inventory_file)
+        else:
+            hostvar = get_hostvars(host)  # dict
+
+        # manage the case where a bastion var is defined from another var
+        # Ex: bastion_host = {{ my_bastion_host }}
         bastion_port = get_var_within(
             hostvar.get("bastion_port", os.environ.get("BASTION_PORT", 22)), hostvar
         )

tests.py

@@ -1,6 +1,13 @@
+import os
+
 from yaml import dump
 
-from lib import get_var_within, manage_conf_file
+from lib import (
+    awx_get_inventory_file,
+    get_bastion_vars,
+    get_var_within,
+    manage_conf_file,
+)
 
 BASTION_HOST = "my_bastion"
 BASTION_PORT = 22
@@ -95,3 +102,34 @@ def test_get_var_not_a_string():
     hostvars = {"bastion_host": 68}
     bastion_host = get_var_within(hostvars["bastion_host"], hostvars)
     assert bastion_host == hostvars["bastion_host"]
+
+
+def test_awx_get_inventory_file_default():
+    assert awx_get_inventory_file() == "/runner/inventory/hosts"
+
+
+def test_awx_get_inventory_file_env_defined():
+    env_path = "/my_awx"
+    os.environ["AWX_RUN_DIR"] = env_path
+    assert awx_get_inventory_file() == f"{env_path}/inventory/hosts"
+    os.environ.pop("AWX_RUN_DIR")
+
+
+def test_get_bastion_vars():
+    host_vars = {
+        "bastion_port": BASTION_PORT,
+        "bastion_host": BASTION_HOST,
+        "bastion_user": BASTION_USER,
+    }
+    bastion_vars = get_bastion_vars(host_vars)
+    assert (
+        bastion_vars["bastion_port"] == BASTION_PORT
+        and bastion_vars["bastion_host"] == BASTION_HOST
+        and bastion_vars["bastion_user"] == BASTION_USER
+    )
+
+
+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"]