From 1801f3f2805bf0cd7cad47f7ef56a1d654e00705 Mon Sep 17 00:00:00 2001
From: topher-lo <46541035+topher-lo@users.noreply.github.com>
Date: Sat, 25 Jan 2025 15:19:59 -0800
Subject: [PATCH] fix ansible integration

---
 .../tracecat_registry/integrations/ansible.py | 114 ++++++++++++------
 1 file changed, 78 insertions(+), 36 deletions(-)

diff --git a/registry/tracecat_registry/integrations/ansible.py b/registry/tracecat_registry/integrations/ansible.py
index f2c778d14..9cf68d234 100644
--- a/registry/tracecat_registry/integrations/ansible.py
+++ b/registry/tracecat_registry/integrations/ansible.py
@@ -1,8 +1,6 @@
 """Generic interface to Ansible Python API."""
 
 import asyncio
-import tempfile
-from pathlib import Path
 from typing import Annotated, Any
 
 import orjson
@@ -13,26 +11,33 @@
 
 ansible_secret = RegistrySecret(
     name="ansible",
-    optional_keys=[
+    keys=[
         "ANSIBLE_SSH_KEY",
+    ],
+    optional_keys=[
         "ANSIBLE_PASSWORDS",
     ],
 )
 """Ansible Runner secret.
 
 - name: `ansible`
-- optional_keys:
+- keys:
     - `ANSIBLE_SSH_KEY`
+- optional_keys:
     - `ANSIBLE_PASSWORDS`
 
 `ANSIBLE_SSH_KEY` should be the private key string, not the path to the file.
-`ANSIBLE_PASSWORDS` should be a JSON object mapping password prompts to their responses (e.g. `{"Vault password": "password"}`).
+`ANSIBLE_PASSWORDS` should be a JSON object mapping password prompts to their responses:
+{
+    "SSH password:": "sshpass",
+    "BECOME password": "sudopass",
+}
 """
 
 
 @registry.register(
     default_title="Run playbook",
-    description="Run Ansible playbook given as a list of plays in JSON format.",
+    description="Run Ansible playbook on a single host given a list of plays in JSON format. Supports SSH host-connection mode only.",
     display_group="Ansible",
     doc_url="https://docs.ansible.com/ansible/latest/index.html",
     namespace="tools.ansible",
@@ -42,7 +47,23 @@ async def run_playbook(
     playbook: Annotated[
         list[dict[str, Any]], Field(..., description="List of plays to run")
     ],
-    extra_vars: Annotated[
+    host: Annotated[
+        str,
+        Field(description="Host to SSH into and run the playbook on"),
+    ],
+    host_name: Annotated[
+        str,
+        Field(description="Host name to use in the inventory"),
+    ],
+    user: Annotated[
+        str,
+        Field(description="SSH user to connect as"),
+    ],
+    envvars: Annotated[
+        dict[str, Any] | None,
+        Field(description="Environment variables to pass to the playbook"),
+    ] = None,
+    extravars: Annotated[
         dict[str, Any] | None,
         Field(description="Extra variables to pass to the playbook"),
     ] = None,
@@ -50,40 +71,61 @@ async def run_playbook(
         dict[str, Any] | None,
         Field(description="Additional keyword arguments to pass to the Ansible runner"),
     ] = None,
+    timeout: Annotated[
+        int,
+        Field(description="Timeout for the playbook execution in seconds"),
+    ] = 60,
 ) -> list[dict[str, Any]]:
     ssh_key = secrets.get("ANSIBLE_SSH_KEY")
     passwords = secrets.get("ANSIBLE_PASSWORDS")
 
-    if not ssh_key and not passwords:
+    extravars = extravars or {}
+    runner_kwargs = runner_kwargs or {}
+
+    if "inventory" in runner_kwargs:
         raise ValueError(
-            "Either `ANSIBLE_SSH_KEY` or `ANSIBLE_PASSWORDS` must be provided"
+            "`inventory` is not supported in this integration. Please use the `host` and `user` parameters instead."
         )
 
-    extra_vars = extra_vars or {}
-    runner_kwargs = runner_kwargs or {}
+    local_patterns = ["localhost", "127.0.0.1", "::1", "0:0:0:0:0:0:0:1"]
+    if host in local_patterns:
+        # Block local connections in initial host connection
+        raise ValueError(f"Local connections are not supported. Got host: {host}")
+
+    # Create inventory config
+    inventory = {
+        "all": {
+            "hosts": {
+                host_name: {
+                    "ansible_host": host,
+                    "ansible_user": user,
+                    "ansible_connection": "ssh",
+                }
+            }
+        }
+    }
+
+    if ssh_key:
+        runner_kwargs["ssh_key"] = ssh_key
+    if passwords:
+        runner_kwargs["passwords"] = orjson.loads(passwords)
+
+    loop = asyncio.get_running_loop()
+
+    def run():
+        runner, result = run_async(
+            playbook=playbook,
+            envvars=envvars,
+            extravars=extravars,
+            timeout=timeout,
+            inventory=inventory,
+            **runner_kwargs,
+        )
+        return runner, result
 
-    with tempfile.TemporaryDirectory() as temp_dir:
-        if ssh_key:
-            ssh_key_path = Path(temp_dir) / "id_rsa"
-            with ssh_key_path.open("w") as f:
-                f.write(ssh_key)
-            runner_kwargs["ssh_key"] = str(ssh_key_path.resolve())
-
-        if passwords:
-            runner_kwargs["passwords"] = orjson.loads(passwords)
-
-        loop = asyncio.get_running_loop()
-
-        def run():
-            return run_async(
-                private_data_dir=temp_dir,
-                playbook=playbook,
-                extravars=extra_vars,
-                **runner_kwargs,
-            )
-
-        _, result = await loop.run_in_executor(None, run)
-        if isinstance(result, Runner):
-            return list(result.events)
-        else:
-            raise ValueError("Ansible runner returned no result.")
+    _, result = await loop.run_in_executor(None, run)
+    if isinstance(result, Runner):
+        # Events are a generator, so we need to convert to a list
+        return list(result.events)
+    else:
+        raise ValueError("Ansible runner returned no result.")