diff --git a/playbook.yaml b/playbook.yaml index cb26743..5e8f6d4 100644 --- a/playbook.yaml +++ b/playbook.yaml @@ -2,21 +2,17 @@ playbook_name: "test_playbook" registries: - "./registry/directory_actions.py" - "./registry/restic_actions.py" +log_dir: "./" steps: - - name: "backup dir1" + - name: "backup rusty-packages" actions: - pre: "directory_exists" - play: "copy_directory" - post: "verify_copy" + create_repo_dir_if_not_exists: "create_dir_if_not_exists" + create_repo_if_not_exists: "create_repo_if_not_exists" + check_restic_environment: "check_restic_environment" + backup_data_to_restic_repo: "backup_data_to_restic_repo" context: - source_dir: "./dir1" - destiny_dir: "./dir2" - - name: "backup dir2" - actions: - pre: "directory_exists" - play: "copy_directory" - post: "verify_copy" - context: - source_dir: "./dir1" - destiny_dir: "./dir2" + passwordFile: "/home/esilva/Trash/test_borgmatic/pass_file" + repoPath: "/home/esilva/Trash/test_borgmatic/destiny2" + sourcePath: "/home/esilva/Trash/rusty_packages" + serviceTag: "rusty-packages" diff --git a/playbook/__init__.py b/playbook/__init__.py index 199f632..9b80a58 100644 --- a/playbook/__init__.py +++ b/playbook/__init__.py @@ -4,13 +4,13 @@ import os import sys import yaml import json -import time import importlib.util +from datetime import datetime, timezone from typing import Dict, List, Any, Optional from dataclasses import dataclass from abc import ABC, abstractmethod -from playbook.models import Status, StepLog +from playbook.models import Status, StepLog, timed_run from playbook.action_registry import ActionFn, ActionRegistry from playbook.premade_steps_registry import registry @@ -38,22 +38,13 @@ class CustomStep(StepIF): def get_action_names(self) -> List[str]: return [action.__name__ for action in self.__actions] - def __timed_run(self, op: ActionFn) -> StepLog: - """Function that sets the duration_sec field on the log.""" - start_time: int = time.perf_counter_ns() - log: StepLog = op(self.__context, self.name) - log.duration_sec = time.perf_counter_ns() - start_time - - return log - + # NOTE: refactor this method to be more readable def run(self) -> StepLog: """ Run the step. """ - start_time: int = time.perf_counter_ns() - substeps: List[StepLog] = [] status: Status = Status.GOOD for op in self.__actions: - log: StepLog = self.__timed_run(op) + log: StepLog = timed_run(op, self.__context, op.__name__) status = Status.BAD if log.failed else Status.GOOD substeps.append(log) @@ -61,8 +52,6 @@ class CustomStep(StepIF): if status == Status.BAD: break - delta: int = time.perf_counter_ns() - start_time - errors = [] msg: str = "" if status == Status.BAD: @@ -71,7 +60,7 @@ class CustomStep(StepIF): msg = "success" return StepLog( - step_name=self.name, status=status, duration_sec=delta, + step_name=self.name, status=status, msg=msg, error=errors, substeps=substeps, ) @@ -90,6 +79,7 @@ class StepEntry(object): class Play(object): def __init__(self, name : str, registries: Optional[List[ActionRegistry]] = None): self.name: str = name + self.log_dir: str = "" self.__steps: List[StepEntry] = [] self.__registries: List[ActionRegistry] = [registry] + (registries or []) @@ -141,7 +131,34 @@ class Play(object): return json.dumps(data) def play(self) -> StepLog: - start_time: int = time.perf_counter_ns() + log: StepLog = timed_run(self.__play) + + 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(log, file) + except Exception as e: + return f"failed to create log file {e}" + return "" + + def __play(self) -> StepLog: playLog: StepLog = StepLog( step_name=self.name, status=Status.GOOD, @@ -152,18 +169,15 @@ class Play(object): ) for s in self.__steps: - log: StepLog = s.step.run() + log: StepLog = timed_run(s.step.run) playLog.substeps.append(log) if log.failed: - playLog.error.append(f"failed in step: {s.name}") + playLog.error.append({"status": "failed", "output": f"failed in step: {s.name}"}) playLog.status = Status.BAD break - delta: int = time.perf_counter_ns() - start_time - playLog.duration_sec = delta - if playLog.status == Status.GOOD: playLog.msg = "success" @@ -227,16 +241,17 @@ class Play(object): raise PlaybookError(f"Error loading registry '{reg_path}': {e}") from e playbook = cls(data["playbook_name"], registries=custom_registries) + if data["log_dir"]: + playbook.log_dir = data["log_dir"] for idx, step_cfg in enumerate(data.get("steps", [])): step_name = step_cfg.get("name", f"step_{idx}") actions = step_cfg.get("actions", {}) context = step_cfg.get("context", {}) - if not isinstance(actions, dict) or len(actions) != 3: + if not isinstance(actions, dict) or len(actions) <= 0: raise PlaybookError( - f"Step '{step_name}' (index {idx}) must define an 'actions' map " - "containing 'pre', 'play' and 'post'" + f"Step '{step_name}' (index {idx}) must define at least one action " ) try: diff --git a/playbook/models.py b/playbook/models.py index 1fe503b..371de7c 100644 --- a/playbook/models.py +++ b/playbook/models.py @@ -1,6 +1,8 @@ +import time +from enum import StrEnum +from typing import List, Dict, Any, Callable +from datetime import datetime, timezone from dataclasses import dataclass, field -from enum import Enum, StrEnum -from typing import Callable, Dict, List, Any class Status(StrEnum): @@ -13,14 +15,55 @@ class StepLog(object): step_name: str status: Status = Status.BAD # By default set it to bad so the it fails by default # The user needs to be explicit - duration_sec: int = 0 + cmd: str = "" + start_date_timestamp: int = 0 + start_date: str = "" + end_date: str = "" + duration_sec: float = 0.0 + human_readable_duration: str = "" msg: str = "" - error: List[str] = field(default_factory=list) + error: List[Dict[Any, Any]] = field(default_factory=list) substeps: List[StepLog] = field(default_factory=list) + log_file_path: str = "" @property def failed(self) -> bool: """ Checks if the step failed """ return len(self.error) > 0 or self.status == Status.BAD + @classmethod + def ok(cls, name: str, msg: str = "success") -> StepLog: + return cls(step_name=name, status=Status.GOOD, msg=msg) + @classmethod + def fail(cls, name: str, errors: List[Any]) -> StepLog: + return cls(step_name=name, status=Status.BAD, error=errors) + + +def human_readable_date(time: int) -> str: + seconds = time // 1_000_000_000 + nanos = time % 1_000_000_000 + + dt = datetime.fromtimestamp(seconds, tz=timezone.utc).astimezone() + + return dt.strftime(f"%Y-%m-%dT%H:%M:%S.{nanos:09d}%:z") + + +def timed_run(op: Callable[..., StepLog], *args: Any, **kwargs: Any) -> StepLog: + start_date: int = time.time_ns() + start_time: int = time.perf_counter_ns() + + try: + log: StepLog = op(*args, **kwargs) + finally: + delta = time.perf_counter_ns() - start_time + + # NOTE: we could simplyfy this by having a marshalling method + if isinstance(log, StepLog): + end_date: int = time.time_ns() + log.duration_sec = delta / 1_000_000_000.0 + log.start_date = human_readable_date(start_date) + log.start_date_timestamp = start_date + log.end_date = human_readable_date(end_date) + + return log diff --git a/registry/directory_actions.py b/registry/directory_actions.py index 7345a77..d14e5b9 100644 --- a/registry/directory_actions.py +++ b/registry/directory_actions.py @@ -1,34 +1,42 @@ # directory actions registry from playbook.action_registry import ActionRegistry -from playbook.models import StepLog, Status +from playbook.models import StepLog + +import os + diract = ActionRegistry("directory_actions_reg") + @diract.register(name="directory_exists", version="stat1.0") def directory_exists(ctx, name) -> StepLog: - log: StepLog = StepLog( - step_name=name, - status=Status.GOOD, - msg="source directory exists", - ) - return log + if not os.path.exists(ctx["repoPath"]): + return StepLog.fail(name, [ + {"status": "failed", "output": f"dir '{ctx["repoPath"]}' doesn't exist"}] + ) + return StepLog.ok(name, "directory exists") -@diract.register(name="copy_directory", version="cp1.0") -def copy_directory(ctx, name) -> StepLog: - log: StepLog = StepLog( - step_name=name, - status=Status.GOOD, - msg="directory was copied/movied/borged/rsynced", - ) - return log +@diract.register(name="create_dir", version="stat1.0") +def create_dir(ctx, name) -> StepLog: + try: + os.makedirs(ctx["repoPath"]) + except Exception as e: + return StepLog.fail(name, [{"status": "failed", "output": str(e)}]) + + return StepLog.ok(name, "dir create successfully") -@diract.register(name="verify_copy", version="verify1.0") -def verify_copy(ctx, name) -> StepLog: - log: StepLog = StepLog( - step_name=name, - status=Status.GOOD, - msg="copied/movied/borged/rsynced exists and is valid", - ) - return log +@diract.register(name="create_dir_if_not_exists", version="stat1.0") +def create_dir_if_not_exists(ctx, name) -> StepLog: + log: StepLog = directory_exists(ctx, name + f".{directory_exists.__name__}") + msg: str = "" + if log.failed: + dir_creation_log: StepLog = create_dir(ctx, name + f".{create_dir.__name__}") + if dir_creation_log.failed: + return StepLog.fail(name, [{"status": "failed", "output": dir_creation_log}]) + msg = "dir created succesfully" + else: + msg = "dir already exists" + + return StepLog.ok(name, msg) diff --git a/registry/restic_actions.py b/registry/restic_actions.py index 53b854a..d652a6e 100644 --- a/registry/restic_actions.py +++ b/registry/restic_actions.py @@ -1,54 +1,175 @@ -# borg actions registry +# restic actions registry from playbook.action_registry import ActionRegistry -from playbook.models import StepLog, Status +from playbook.models import StepLog +import sys +import json +import subprocess +from typing import List, Dict, Any, Tuple + + +# NOTE: +# most restic run commands seem to be the same pattern of: +# restic command -> parse output -> grab errors +# we can create a class for this I reckon + + +def parse_output(line: str) -> List[Dict[str, Any]]: + errors = [] + for line in line.split('\n'): + if line.startswith("{"): # } treesitter is borked lmao + data = json.loads(line) + errors.append(data) + return errors + + +def find_restic_errors(data: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + errors = [] + for s in data: + if s["message_type"] == "exit_error": + errors.append(s) + return errors + + +def run_restic_command(cmd: List[str]) -> Tuple[int, List[Dict[str, Any]]]: + """ + Executes a restic command, streams output live to terminal, + and returns (returncode, parsed_json_objects). + """ + parsed_json = [] + + # Start process with combined stdout/stderr pipe + process = subprocess.Popen( + cmd, + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1 + ) + + # NOTE: is this what we want? + assert process.stdout is not None + + # Read line-by-line in real-time + for line in iter(process.stdout.readline, ''): + # Write live output directly to terminal screen + sys.stdout.write(line) + sys.stdout.flush() + + # Parse and capture valid JSON objects on the fly + clean_line = line.strip() + if clean_line.startswith('{'): # } + try: + data = json.loads(clean_line) + parsed_json.append(data) + except json.JSONDecodeError: + pass + + process.stdout.close() + returncode = process.wait() + return returncode, parsed_json restic = ActionRegistry("restic_actions_reg") -@restic.register(name="check_repo_exists", version="preborg1.0") -def check_repo_exists(ctx, name) -> StepLog: - passwordFile: str = ctx["passwordFile"] - repoPath: str = ctx["repoPath"] +@restic.register(name="check_restic_environment", version="") +def check_restic_environment(ctx, name: str = "") -> StepLog: + for v in ["passwordFile", "repoPath", "sourcePath", "serviceTag"]: + if not ctx.get(v): + return StepLog.fail(name, [f"{v} variable is empty"]) - if passwordFile == "": - return StepLog( - step_name=name, - status=Status.BAD, - error=["password file variable is empty"], - ) + return StepLog.ok(name, "environment configuration seems to be valid") - if repoPath == "": - return StepLog( - step_name=name, - status=Status.BAD, - error=["password file variable is empty"] - ) - # If the path doesn't exist restic will make sure of warning us about it - - return StepLog( - step_name=name, - status=Status.GOOD, - msg="password file variable is empty", +@restic.register(name="create_restic_repo", version="") +def create_restic_repo(ctx, name: str = "") -> StepLog: + cmd = ["restic", "init", "--json", "-r", ctx["repoPath"], + "--password-file", ctx["passwordFile"]] + + returncode, json_output = run_restic_command(cmd) + + errors: List[Dict[str, Any]] = find_restic_errors(json_output) + + # Handle error messages in JSON output + if len(errors) > 0: + return StepLog.fail(name, errors) + + # Make sure the return code is good too + if returncode != 0: + return StepLog.fail(name, [{"status": "failed", "output": json_output}]) + + return StepLog.ok(name, "repo create succesfully") + + +@restic.register(name="create_repo_if_not_exists", version="") +def create_repo_if_not_exists(ctx, name: str = "") -> StepLog: + log: StepLog = check_if_repo_exists(ctx, "") + + msg: str = "" + if log.failed: + # Repo doesn't exist is error code 10 + if log.error[0]["code"] != 10: + return StepLog.fail(name, [{"status": "failed", "output": log.error}]) + + # Create repo + repoCreateLog: StepLog = create_restic_repo(ctx, name + f".{create_restic_repo}") + if repoCreateLog.failed: + return StepLog.fail(name, [{"status": "failed", "output": repoCreateLog.error}]) + msg = "repo created" + else: + msg = "repo already exists" + + + return StepLog.ok(name, msg) + + +@restic.register(name="check_if_repo_exists", version="") +def check_if_repo_exists(ctx, name: str = "") -> StepLog: + cmd = ["restic", "-r", ctx["repoPath"], "cat", "config", "--json", + "--password-file", ctx["passwordFile"]] + returncode, json_output = run_restic_command(cmd) + + errors: List[Dict[str, Any]] = find_restic_errors(json_output) + + # Handle error messages in JSON output + if len(errors) > 0: + return StepLog.fail(name, errors) + + # Make sure the return code is good too + if returncode != 0: + return StepLog.fail(name, [{"status": "failed", "output": json_output}]) + + return StepLog.ok(name, "repo seems to exist") + + +@restic.register(name="backup_data_to_restic_repo", version="") +def backup_data_to_restic_repo(ctx, name) -> StepLog: + out = subprocess.run( + ["restic", "backup", ctx["sourcePath"], "--json", "--quiet", + "-r", ctx["repoPath"], "--password-file", ctx["passwordFile"]], + text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, ) + json_output: List[Dict[str, Any]] = parse_output(out.stdout) + errors: List[Dict[str, Any]] = find_restic_errors(json_output) -@restic.register(name="backup_data_to_repo", version="playborg1.0") -def backup_data_to_repo(ctx, name) -> StepLog: - log: StepLog = StepLog( - step_name=name, - status=Status.GOOD, - msg="directory was copied/movied/borged/rsynced", - ) - return log + # Handle error messages in JSON output + if len(errors) > 0: + return StepLog.fail(name, errors) + # Make sure the return code is good too + if out.returncode != 0: + return StepLog.fail(name, [{"status": "failed", "output": out.stdout}]) -@restic.register(name="post", version="postborg1.0") -def validate_data_backup(ctx, name) -> StepLog: - log: StepLog = StepLog( - step_name=name, - status=Status.GOOD, - msg="copied/movied/borged/rsynced exists and is valid", - ) - return log + # Return code is 0, so it was successful, output the message + try: + data = json.loads(out.stdout) + msg = data + except: + msg = { + "status": "ran succesfully, but an error occurred parsing the output", + "output": out.stdout, + } + + return StepLog.ok(name, str(msg)) diff --git a/taskrunner.py b/taskrunner.py index 45e9e76..8187ee6 100644 --- a/taskrunner.py +++ b/taskrunner.py @@ -19,15 +19,6 @@ def parse_args() -> argparse.Namespace: help="Path to the playbook yaml file" ) - # NOTE: For now I will only take the registries from the yaml playbook - # # Flag for the custom registry files - # parser.add_argument( - # "-r", "--registry", - # nargs="+", - # default=[], - # help="One or more python files with actions" - # ) - return parser.parse_args() @@ -42,7 +33,7 @@ def main(): traceback.print_exc() exit(1) - print(newPlay.view_playbook()) + # print(newPlay.view_playbook()) print(json.dumps(asdict(newPlay.play()), indent=2))