Skip to content

Commit

Permalink
feat(integration): S3 getter and Ansible playbook runner (TracecatHQ#573
Browse files Browse the repository at this point in the history
)
  • Loading branch information
topher-lo authored Dec 8, 2024
1 parent e59b066 commit a9db1ef
Showing 5 changed files with 218 additions and 15 deletions.
1 change: 1 addition & 0 deletions registry/pyproject.toml
Original file line number Diff line number Diff line change
@@ -24,6 +24,7 @@ dependencies = [
"adbc-driver-postgresql==1.0.0",
"adbc-driver-snowflake==1.0.0",
"adbc-driver-sqlite==1.0.0",
"ansible-runner==2.4.0",
"aioboto3==13.0.1",
"boto3==1.34.70",
"crowdstrike-falconpy==1.4.4",
85 changes: 85 additions & 0 deletions registry/tracecat_registry/integrations/ansible.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,85 @@
"""Generic interface to Ansible Python API."""

import asyncio
import tempfile
from pathlib import Path
from typing import Annotated, Any

import orjson
from ansible_runner import run_async
from pydantic import Field

from tracecat_registry import RegistrySecret, registry, secrets

ansible_secret = RegistrySecret(
name="ansible",
optional_keys=[
"ANSIBLE_SSH_KEY",
"ANSIBLE_PASSWORDS",
],
)
"""Ansible Runner secret.
- name: `ansible`
- optional_keys:
- `ANSIBLE_SSH_KEY`
- `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"}`).
"""


@registry.register(
default_title="Run Ansible Playbook",
description="Run an Ansible playbook",
display_group="Ansible",
namespace="integrations.ansible",
secrets=[ansible_secret],
)
async def run_ansible_playbook(
playbook: Annotated[
list[dict[str, Any]], Field(..., description="List of plays to run")
],
extra_vars: Annotated[
dict[str, Any],
Field(description="Extra variables to pass to the playbook"),
] = None,
runner_kwargs: Annotated[
dict[str, Any],
Field(description="Additional keyword arguments to pass to the Ansible runner"),
] = None,
) -> dict[str, Any]:
ssh_key = secrets.get("ANSIBLE_SSH_KEY")
passwords = secrets.get("ANSIBLE_PASSWORDS")

if not ssh_key and not passwords:
raise ValueError(
"Either `ANSIBLE_SSH_KEY` or `ANSIBLE_PASSWORDS` must be provided"
)

runner_kwargs = runner_kwargs or {}

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 result is None:
raise ValueError("Ansible runner returned no result.")
return result
33 changes: 18 additions & 15 deletions registry/tracecat_registry/integrations/boto3.py
Original file line number Diff line number Diff line change
@@ -11,29 +11,32 @@

from tracecat_registry import RegistrySecret, logger, registry, secrets

# TODO: Support possible sets of secrets
# e.g. AWS_PROFILE_NAME or AWS_ROLE_ARN
aws_secret = RegistrySecret(
name="aws",
keys=["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION"],
optional_keys=[
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
"AWS_REGION",
"AWS_PROFILE_NAME",
"AWS_ROLE_ARN",
"AWS_ROLE_SESSION_NAME",
],
)
"""AWS secret.
Secret
------
- name: `aws`
- keys:
- `AWS_ACCESS_KEY_ID`
- `AWS_SECRET_ACCESS_KEY`
- `AWS_REGION`
Example Usage
-------------
Environment variables:
>>> secrets.get["AWS_ACCESS_KEY_ID"]
Expression:
>>> ${{ SECRETS.aws.AWS_ACCESS_KEY_ID }}
- optional_keys:
Either:
- `AWS_ACCESS_KEY_ID`
- `AWS_SECRET_ACCESS_KEY`
- `AWS_REGION`
Or:
- `AWS_PROFILE_NAME`
Or:
- `AWS_ROLE_ARN`
- `AWS_ROLE_SESSION_NAME`
"""


80 changes: 80 additions & 0 deletions registry/tracecat_registry/integrations/s3.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
"""S3 integration to download files and return contents as a string."""

import re
from typing import Annotated

from pydantic import Field

from tracecat_registry import RegistrySecret, registry
from tracecat_registry.integrations.boto3 import get_session

# Add this at the top with other constants
BUCKET_REGEX = re.compile(r"^[a-z0-9][a-z0-9.-]*[a-z0-9]$")

s3_secret = RegistrySecret(
name="s3",
optional_keys=[
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
"AWS_REGION",
"AWS_PROFILE_NAME",
"AWS_ROLE_ARN",
"AWS_ROLE_SESSION_NAME",
],
)
"""AWS secret.
Secret
------
- name: `aws`
- optional_keys:
Either:
- `AWS_ACCESS_KEY_ID`
- `AWS_SECRET_ACCESS_KEY`
- `AWS_REGION`
Or:
- `AWS_PROFILE_NAME`
Or:
- `AWS_ROLE_ARN`
- `AWS_ROLE_SESSION_NAME`
"""


@registry.register(
default_title="Parse S3 URI",
description="Parse an S3 URI into a bucket and key.",
display_group="AWS S3",
namespace="integrations.aws_s3",
)
async def parse_uri(uri: str) -> tuple[str, str]:
uri = str(uri).strip()
if not uri.startswith("s3://"):
raise ValueError("S3 URI must start with 's3://'")

uri_path = uri.replace("s3://", "")
uri_paths = uri_path.split("/")
bucket = uri_paths.pop(0)
key = "/".join(uri_paths)

return bucket, key


@registry.register(
default_title="Download S3 Object",
description="Download an object from S3 and return its body as a string.",
display_group="AWS S3",
namespace="integrations.aws_s3",
secrets=[s3_secret],
)
async def download_object(
bucket: Annotated[str, Field(..., description="S3 bucket name.")],
key: Annotated[str, Field(..., description="S3 object key.")],
) -> str:
session = await get_session()
async with session.client("s3") as s3_client:
obj = await s3_client.get_object(Bucket=bucket, Key=key)
body = await obj["Body"].read()
# Defensively handle different types of bodies
if isinstance(body, bytes):
return body.decode("utf-8")
return body
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
type: action
definition:
name: run_playbook_from_s3
namespace: integrations.ansible
title: Run Ansible Playbook from S3
description: Download an Ansible playbook from S3 and run it
display_group: Ansible
expects:
playbook_path:
type: str
description: Path to the playbook in S3
extra_vars:
type: dict[str, any]
description: Extra variables to pass to the playbook
runner_kwargs:
type: dict[str, any]
description: Additional keyword arguments to pass to the Ansible runner
steps:
- ref: parse_uri
action: integrations.aws_s3.parse_uri
args:
uri: ${{ inputs.playbook_path }}
- ref: download_playbook
action: integrations.aws_s3.download_object
args:
bucket: ${{ steps.parse_uri.result.bucket }}
key: ${{ steps.parse_uri.result.key }}
- ref: run_playbook
action: integrations.ansible.run_playbook
args:
playbook: ${{ steps.download_playbook.result }}
extra_vars: ${{ inputs.extra_vars }}
runner_kwargs: ${{ inputs.runner_kwargs }}
returns: ${{ steps.run_playbook.result }}

0 comments on commit a9db1ef

Please sign in to comment.