356 lines
12 KiB
Python
356 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
import yaml
|
|
import json
|
|
import importlib.util
|
|
from datetime import datetime, timedelta, timezone
|
|
from typing import Dict, List, Any, Optional, Tuple
|
|
from dataclasses import MISSING, asdict, dataclass
|
|
from abc import ABC, abstractmethod
|
|
|
|
from playbook.models import Status, StepLog, timed_run
|
|
from playbook.action_registry import ActionFn, ActionRegistry
|
|
from playbook.premade_steps_registry import registry
|
|
|
|
|
|
class StepIF(ABC):
|
|
@abstractmethod
|
|
def run(self, ctx: Dict[str, Any]) -> StepLog:
|
|
"""Perform the actions."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def get_action_names(self) -> List[str]:
|
|
""" Return a mapping of action roles (pre, play, post) to function names. """
|
|
pass
|
|
|
|
|
|
class CustomStep(StepIF):
|
|
""" Setup custom steps. """
|
|
|
|
def __init__(self, name: str, actions: List[ActionFn], context: Dict[str, Any]):
|
|
self.name: str = name
|
|
self.__actions = actions
|
|
self.__context = context
|
|
|
|
def get_action_names(self) -> List[str]:
|
|
return [action.__name__ for action in self.__actions]
|
|
|
|
# NOTE: refactor this method to be more readable
|
|
def run(self, ctx) -> StepLog:
|
|
""" Run the step. """
|
|
substeps: List[StepLog] = []
|
|
status: Status = Status.GOOD
|
|
|
|
# NOTE:
|
|
# Use this to move context from the previous step to the next step
|
|
# This is wrong, we are moving away from having a list of steps here to a single
|
|
# step for each operation
|
|
prev_step_ctx: Dict[str, Any] = {}
|
|
|
|
for op in self.__actions:
|
|
try:
|
|
log: StepLog = timed_run(op, ctx | prev_step_ctx | self.__context, self.name)
|
|
except Exception as e:
|
|
log: StepLog = StepLog.fail(self.name, [{"status": "failed", "output": str(e)}])
|
|
|
|
status = Status.BAD if log.failed else Status.GOOD
|
|
prev_step_ctx = log.pipe_ctx
|
|
|
|
substeps.append(log)
|
|
|
|
if status == Status.BAD:
|
|
break
|
|
|
|
errors = []
|
|
msg: str = ""
|
|
if status == Status.BAD:
|
|
# Note, use this to set the standar error output/formatting
|
|
# errors = [{"status": "failed", "output": ""}]
|
|
errors = []
|
|
else:
|
|
msg = "success"
|
|
|
|
return StepLog(
|
|
step_name=self.name, status=status,
|
|
msg=msg, error=errors, substeps=substeps, pipe_ctx=prev_step_ctx
|
|
)
|
|
|
|
|
|
class PlaybookError(Exception):
|
|
""" Base exception for Playbook parsing and execution failures. """
|
|
pass
|
|
|
|
|
|
@dataclass
|
|
class StepEntry(object):
|
|
name: str
|
|
step: StepIF
|
|
|
|
|
|
class Play(object):
|
|
def __init__(self, name : str, registries: Dict[str, ActionRegistry]):
|
|
self.name: str = name
|
|
self.log_dir: str = ""
|
|
self._context: Dict[str, Any] = {}
|
|
self._acts: Dict[str, List[CustomStep]] = {}
|
|
self._steps: List[StepEntry] = []
|
|
self._registries: Dict[str, ActionRegistry] = {registry.name: registry} | registries
|
|
|
|
# NOTE: this should move to the ActionRegistry thing probably
|
|
@staticmethod
|
|
def _load_registry_file(file_path: str) -> List[ActionRegistry]:
|
|
abs_path = os.path.abspath(file_path)
|
|
module_name = os.path.splitext(os.path.basename(abs_path))[0]
|
|
|
|
spec = importlib.util.spec_from_file_location(module_name, abs_path)
|
|
if spec is None or spec.loader is None:
|
|
raise ImportError(f"could not create spec for registry file: {file_path}")
|
|
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[module_name] = module
|
|
spec.loader.exec_module(module)
|
|
|
|
discovered_registries = [
|
|
obj for obj in vars(module).values()
|
|
if isinstance(obj, ActionRegistry)
|
|
]
|
|
|
|
if not discovered_registries:
|
|
raise PlaybookError(f"No ActionRegistry instances were found in '{file_path}'")
|
|
|
|
return discovered_registries
|
|
|
|
def add_step(self, name: str, stepRunner : StepIF):
|
|
self._steps.append(StepEntry(name=name, step=stepRunner))
|
|
|
|
def view_playbook(self) -> str:
|
|
acts_info = {}
|
|
for act_name, step_list in self._acts.items():
|
|
steps_info = []
|
|
for s in step_list:
|
|
steps_info.append({
|
|
"name": s.name,
|
|
"actions": s.get_action_names()
|
|
})
|
|
|
|
new_act = {
|
|
"name": act_name,
|
|
"steps": steps_info
|
|
}
|
|
|
|
acts_info[act_name] = new_act
|
|
|
|
|
|
|
|
registry_data = []
|
|
for reg in self._registries.values():
|
|
registry_data.append(reg.manifest())
|
|
|
|
data = {
|
|
"name": self.name,
|
|
"registries": registry_data,
|
|
"number_of_acts": len(acts_info),
|
|
"acts": acts_info,
|
|
}
|
|
|
|
return json.dumps(data)
|
|
|
|
def play(self) -> StepLog:
|
|
log: StepLog = timed_run(self._play_act)
|
|
|
|
if self.log_dir != "":
|
|
log.log_file_path = os.path.join(
|
|
self.log_dir, self.log_file_name(log.start_date_timestamp)
|
|
)
|
|
self._write_log(log)
|
|
|
|
return log
|
|
|
|
def log_file_name(self, start_date_ns: int) -> str:
|
|
""" Format to ISO 8601. """
|
|
seconds = start_date_ns // 1_000_000_000
|
|
ms = (start_date_ns % 1_000_000_000) // 1_000_000
|
|
|
|
dt = datetime.fromtimestamp(seconds, tz=timezone.utc).astimezone()
|
|
|
|
return dt.strftime(f"{self.name}_%Y-%m-%d_%H-%M-%S.{ms:03d}.log")
|
|
|
|
def _write_log(self, log: StepLog):
|
|
try:
|
|
with open(log.log_file_path, "w") as file:
|
|
json.dump(asdict(log), file)
|
|
except Exception as e:
|
|
raise PlaybookError(f"failed to create log file {e}") from e
|
|
return ""
|
|
|
|
def _play_act(self) -> StepLog:
|
|
playLog: StepLog = StepLog(
|
|
step_name=self.name,
|
|
status=Status.GOOD,
|
|
duration_sec=0,
|
|
msg="",
|
|
error=[],
|
|
substeps=[],
|
|
)
|
|
|
|
for a in self._acts:
|
|
log: StepLog = self._play_steps(self._acts[a])
|
|
|
|
if log.failed:
|
|
playLog.error.append({
|
|
"status": "failed",
|
|
"output": f"failed in step: {log.step_name}"
|
|
})
|
|
playLog.status = Status.BAD
|
|
break
|
|
|
|
if playLog.status == Status.GOOD:
|
|
playLog.msg = "success"
|
|
|
|
return playLog
|
|
|
|
def _play_steps(self, steps: List[CustomStep]) -> StepLog:
|
|
playLog: StepLog = StepLog(
|
|
step_name=self.name,
|
|
status=Status.GOOD,
|
|
duration_sec=0,
|
|
msg="",
|
|
error=[],
|
|
substeps=[],
|
|
)
|
|
|
|
prev_ctx: Dict[str, Any] = {}
|
|
for step in steps:
|
|
log: StepLog = timed_run(step.run, self._context | prev_ctx)
|
|
prev_ctx = log.pipe_ctx
|
|
|
|
playLog.substeps.append(log)
|
|
|
|
if log.failed:
|
|
playLog.error.append({
|
|
"status": "failed",
|
|
"output": f"failed in step: {step.name}"
|
|
})
|
|
playLog.status = Status.BAD
|
|
break
|
|
|
|
if playLog.status == Status.GOOD:
|
|
playLog.msg = "success"
|
|
|
|
return playLog
|
|
|
|
def _get_function_from_registry(self, fn_name) -> ActionFn:
|
|
for reg in self._registries.values():
|
|
try:
|
|
return reg.get(fn_name, "*")
|
|
except ValueError:
|
|
continue
|
|
|
|
raise ValueError(f"no function with name '{fn_name}' was found")
|
|
|
|
def _build_custom_step(self, name, actions, ctx) -> CustomStep:
|
|
functionList: List[ActionFn] = []
|
|
for f in actions:
|
|
try:
|
|
fn = self._get_function_from_registry(actions[f])
|
|
functionList.append(fn)
|
|
except KeyError as e:
|
|
raise PlaybookError(f"Step '{name}' is missing required action field: {e}") from e
|
|
|
|
return CustomStep(name, functionList, ctx)
|
|
|
|
# NOTE: Move to factory
|
|
@classmethod
|
|
def from_yaml(cls, fp: str) -> Play:
|
|
""" Loads the playbook from a given YAML file. """
|
|
try:
|
|
with open(fp, "r") as f:
|
|
file_contents = f.read()
|
|
except FileNotFoundError as e:
|
|
raise PlaybookError(f"Playbook file not found: {fp}") from e
|
|
except Exception as e:
|
|
raise PlaybookError(f"Failed to read file {fp}: {e}") from e
|
|
|
|
yaml_dir = os.path.dirname(os.path.abspath(fp))
|
|
return cls.__from_yaml_str(file_contents, base_dir=yaml_dir)
|
|
|
|
# NOTE: Move to factory
|
|
@classmethod
|
|
def __from_yaml_str_load_registries(cls, dir, data) -> Dict[str, ActionRegistry]:
|
|
custom_registries: Dict[str, ActionRegistry] = {}
|
|
registry_paths = data.get("registries", {})
|
|
|
|
custom_registries: Dict[str, ActionRegistry] = {}
|
|
for reg_path in registry_paths:
|
|
full_path = os.path.join(dir, reg_path) if not os.path.isabs(reg_path) else reg_path
|
|
try:
|
|
regs = cls._load_registry_file(full_path)
|
|
for new_reg in regs:
|
|
custom_registries[new_reg.name] = new_reg
|
|
except Exception as e:
|
|
raise PlaybookError(f"Error loading registry '{reg_path}': {e}") from e
|
|
|
|
return custom_registries
|
|
|
|
# NOTE: yes, I know, I have to refactor this
|
|
def __from_yaml_str_load_acts(self, data):
|
|
for idx, act in enumerate(data.get("acts", [])):
|
|
act_name = act.get("name", f"step_{idx}")
|
|
steps = act.get("steps", [])
|
|
|
|
self._acts[act_name] = []
|
|
try:
|
|
self.__from_yaml_str_load_steps(act_name, steps)
|
|
except Exception as e:
|
|
raise PlaybookError(f"error loading steps for act {act_name}") from e
|
|
|
|
|
|
# NOTE: Move to factory
|
|
def __from_yaml_str_load_steps(self, key, data):
|
|
for step_cfg in data:
|
|
step_name = step_cfg.get("name", f"")
|
|
actions = step_cfg.get("actions", {})
|
|
context = step_cfg.get("context", {})
|
|
|
|
if not isinstance(actions, dict) or len(actions) <= 0:
|
|
raise PlaybookError(
|
|
f"Step '{step_name}' (index ) must define at least one action "
|
|
)
|
|
|
|
try:
|
|
customStep: CustomStep = self._build_custom_step(step_name, actions, context)
|
|
self._acts[key].append(customStep)
|
|
except Exception as e:
|
|
raise PlaybookError(f"Error configuring step '{step_name}': {e}") from e
|
|
|
|
# NOTE: Move to factory
|
|
@classmethod
|
|
def __from_yaml_str(cls, yaml_str: str, base_dir: str = ".") -> Play:
|
|
data = None
|
|
try:
|
|
data = yaml.safe_load(yaml_str)
|
|
except Exception as e:
|
|
raise PlaybookError(f"YAML Syntax error: {e}") from e
|
|
|
|
if not isinstance(data, dict) or "playbook_name" not in data:
|
|
raise PlaybookError("YAML must contain a top-level 'playbook_name' field.")
|
|
|
|
try:
|
|
custom_registries = cls.__from_yaml_str_load_registries(base_dir, data)
|
|
except Exception as e:
|
|
raise PlaybookError(f"registries load error: {e}") from e
|
|
|
|
playbook = cls(data["playbook_name"], registries=custom_registries)
|
|
try:
|
|
playbook.__from_yaml_str_load_acts(data)
|
|
except Exception as e:
|
|
raise PlaybookError(f"steps load error: {e}") from e
|
|
|
|
playbook._context = data.get("global_context", {})
|
|
playbook.log_dir = data.get("log_dir", "")
|
|
|
|
return playbook
|