From 0df917531e35aecae3a30dc146dccd840527e6d7 Mon Sep 17 00:00:00 2001 From: Ben Thomasson Date: Mon, 8 May 2023 15:03:01 -0400 Subject: [PATCH] Add support for custom actions Add action loading from a directory --- ansible_rulebook/app.py | 5 +- ansible_rulebook/builtin.py | 1 + ansible_rulebook/cli.py | 5 ++ ansible_rulebook/collection.py | 21 ++++++++ ansible_rulebook/common.py | 2 + ansible_rulebook/engine.py | 2 + ansible_rulebook/rule_set_runner.py | 55 ++++++++++++++++++++- ansible_rulebook/schema/ruleset_schema.json | 3 ++ 8 files changed, 92 insertions(+), 2 deletions(-) diff --git a/ansible_rulebook/app.py b/ansible_rulebook/app.py index b98a31c0b..0a84dce84 100644 --- a/ansible_rulebook/app.py +++ b/ansible_rulebook/app.py @@ -85,6 +85,8 @@ async def run(parsed_args: argparse.ArgumentParser) -> None: startup_args.controller_url = parsed_args.controller_url startup_args.controller_token = parsed_args.controller_token startup_args.controller_ssl_verify = parsed_args.controller_ssl_verify + startup_args.source_dir = parsed_args.source_dir + startup_args.action_dir = parsed_args.action_dir validate_actions(startup_args) set_controller_params(startup_args) @@ -98,7 +100,7 @@ async def run(parsed_args: argparse.ArgumentParser) -> None: tasks, ruleset_queues = spawn_sources( startup_args.rulesets, startup_args.variables, - [parsed_args.source_dir], + [startup_args.source_dir], parsed_args.shutdown_delay, ) @@ -122,6 +124,7 @@ async def run(parsed_args: argparse.ArgumentParser) -> None: startup_args.inventory, parsed_args, startup_args.project_data_file, + [startup_args.action_dir], ) logger.info("Cancelling event source tasks") diff --git a/ansible_rulebook/builtin.py b/ansible_rulebook/builtin.py index 5bc3364de..42cea6919 100644 --- a/ansible_rulebook/builtin.py +++ b/ansible_rulebook/builtin.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import subprocess import asyncio import concurrent.futures import glob diff --git a/ansible_rulebook/cli.py b/ansible_rulebook/cli.py index f083d9132..6737f94c1 100644 --- a/ansible_rulebook/cli.py +++ b/ansible_rulebook/cli.py @@ -86,6 +86,11 @@ def get_parser() -> argparse.ArgumentParser: "--source-dir", help="Source dir", ) + parser.add_argument( + "-A", + "--action-dir", + help="Action dir", + ) parser.add_argument( "-i", "--inventory", diff --git a/ansible_rulebook/collection.py b/ansible_rulebook/collection.py index a95e3bd64..4f190e9a4 100644 --- a/ansible_rulebook/collection.py +++ b/ansible_rulebook/collection.py @@ -34,6 +34,10 @@ f"{EDA_PATH_PREFIX}/plugins/event_sources", "plugins/event_source", ] + +EDA_ACTION_PATHS = [ + f"{EDA_PATH_PREFIX}/plugins/rule_action", +] logger = logging.getLogger(__name__) @@ -120,6 +124,23 @@ def load_rulebook(collection, rulebook): print(f"Loading rulebook from {location}") return yaml.safe_load(f.read()) +def has_action(collection, action): + return has_object( + collection, + action, + EDA_ACTION_PATHS, + ".py", + ) + + +def find_action(collection, action): + return find_object( + collection, + action, + EDA_ACTION_PATHS, + ".py", + ) + def has_source(collection, source): return has_object( diff --git a/ansible_rulebook/common.py b/ansible_rulebook/common.py index bc880e5df..203d662ef 100644 --- a/ansible_rulebook/common.py +++ b/ansible_rulebook/common.py @@ -27,3 +27,5 @@ class StartupArgs: controller_ssl_verify: str = field(default="") project_data_file: str = field(default="") inventory: str = field(default="") + source_dir: str = field(default="") + action_dir: str = field(default="") diff --git a/ansible_rulebook/engine.py b/ansible_rulebook/engine.py index e2c225064..2f1f584b5 100644 --- a/ansible_rulebook/engine.py +++ b/ansible_rulebook/engine.py @@ -224,6 +224,7 @@ async def run_rulesets( inventory: str = "", parsed_args: argparse.ArgumentParser = None, project_data_file: Optional[str] = None, + action_directories: Optional[List[str]] = None, ): logger.info("run_ruleset") rulesets_queue_plans = rule_generator.generate_rulesets( @@ -265,6 +266,7 @@ async def run_rulesets( project_data_file=project_data_file, parsed_args=parsed_args, broadcast_method=broadcast, + action_directories=action_directories, ) task_name = f"main_ruleset :: {ruleset_queue_plan.ruleset.name}" ruleset_task = asyncio.create_task( diff --git a/ansible_rulebook/rule_set_runner.py b/ansible_rulebook/rule_set_runner.py index 0ee1ca4f2..ecc2884f8 100644 --- a/ansible_rulebook/rule_set_runner.py +++ b/ansible_rulebook/rule_set_runner.py @@ -15,6 +15,8 @@ import asyncio import gc import logging +import runpy +import os from datetime import datetime from pprint import PrettyPrinter, pformat from typing import Dict, List, Optional, Union, cast @@ -26,6 +28,7 @@ MessageObservedException, ) +from ansible_rulebook.collection import find_action, split_collection_name from ansible_rulebook.builtin import actions as builtin_actions from ansible_rulebook.conf import settings from ansible_rulebook.exception import ShutdownException @@ -55,6 +58,7 @@ def __init__( project_data_file: Optional[str] = None, parsed_args=None, broadcast_method=None, + action_directories=None, ): self.action_loop_task = None self.event_log = event_log @@ -68,6 +72,14 @@ def __init__( self.active_actions = set() self.broadcast_method = broadcast_method self.event_counter = 0 + self.action_directories = action_directories + + def find_action(self, action: str): + for action_dir in self.action_directories: + action_plugin_file = os.path.join(action_dir, f"{action}.py") + if os.path.exists(action_plugin_file): + return runpy.run_path(action_plugin_file) + return None async def run_ruleset(self): tasks = [] @@ -307,7 +319,6 @@ async def _call_action( hosts: List, rules_engine_result, ) -> None: - logger.info("call_action %s", action) result = None @@ -422,6 +433,48 @@ async def _call_action( except BaseException as e: logger.error(e) raise + elif action_plugin := self.find_action(action): + try: + result = await action_plugin['main']( + event_log=self.event_log, + inventory=inventory, + hosts=hosts, + variables=variables, + project_data_file=self.project_data_file, + source_ruleset_name=ruleset, + source_ruleset_uuid=ruleset_uuid, + source_rule_name=rule, + source_rule_uuid=rule_uuid, + rule_run_at=rule_run_at, + **action_args, + ) + except Exception as e: + logger.error("Error calling action %s, err %s", action, str(e)) + result = dict(error=e) + except BaseException as e: + logger.error(e) + raise + elif action_plugin := find_action(*split_collection_name(action)): + try: + result = await action_plugin.main( + event_log=self.event_log, + inventory=inventory, + hosts=hosts, + variables=variables, + project_data_file=self.project_data_file, + source_ruleset_name=ruleset, + source_ruleset_uuid=ruleset_uuid, + source_rule_name=rule, + source_rule_uuid=rule_uuid, + rule_run_at=rule_run_at, + **action_args, + ) + except Exception as e: + logger.error("Error calling action %s, err %s", action, str(e)) + result = dict(error=e) + except BaseException as e: + logger.error(e) + raise else: logger.error("Action %s not supported", action) result = dict(error=f"Action {action} not supported") diff --git a/ansible_rulebook/schema/ruleset_schema.json b/ansible_rulebook/schema/ruleset_schema.json index a880375f4..31853f538 100644 --- a/ansible_rulebook/schema/ruleset_schema.json +++ b/ansible_rulebook/schema/ruleset_schema.json @@ -220,6 +220,9 @@ }, { "$ref": "#/$defs/shutdown-action" + }, + { + "type": "object" } ] }