Refactored

cli/utils/handler/vault.py

@@ -0,0 +1,50 @@
+import subprocess
+from typing import Any, Dict
+
+from yaml.loader import SafeLoader
+from yaml.dumper import SafeDumper
+
+class VaultScalar(str):
+    """A subclass of str to represent vault-encrypted strings."""
+    pass
+
+def _vault_constructor(loader, node):
+    """Custom constructor to handle !vault tag as plain text."""
+    return node.value
+
+def _vault_representer(dumper, data):
+    """Custom representer to dump VaultScalar as literal blocks."""
+    return dumper.represent_scalar('!vault', data, style='|')
+
+SafeLoader.add_constructor('!vault', _vault_constructor)
+SafeDumper.add_representer(VaultScalar, _vault_representer)
+
+class VaultHandler:
+    def __init__(self, vault_password_file: str):
+        self.vault_password_file = vault_password_file
+
+    def encrypt_string(self, value: str, name: str) -> str:
+        """Encrypt a string using ansible-vault."""
+        cmd = [
+            "ansible-vault", "encrypt_string",
+            value, f"--name={name}",
+            "--vault-password-file", self.vault_password_file
+        ]
+        proc = subprocess.run(cmd, capture_output=True, text=True)
+        if proc.returncode != 0:
+            raise RuntimeError(f"ansible-vault encrypt_string failed:\n{proc.stderr}")
+        return proc.stdout
+
+    def encrypt_leaves(self, branch: Dict[str, Any], vault_pw: str):
+        """Recursively encrypt all leaves (plain text values) under the credentials section."""
+        for key, value in branch.items():
+            if isinstance(value, dict):
+                self.encrypt_leaves(value, vault_pw)  # Recurse into nested dictionaries
+            else:
+                # Skip if already vaulted (i.e., starts with $ANSIBLE_VAULT)
+                if isinstance(value, str) and not value.lstrip().startswith("$ANSIBLE_VAULT"):
+                    snippet = self.encrypt_string(value, key)
+                    lines = snippet.splitlines()
+                    indent = len(lines[1]) - len(lines[1].lstrip())
+                    body = "\n".join(line[indent:] for line in lines[1:])
+                    branch[key] = VaultScalar(body)  # Store encrypted value as VaultScalar

cli/utils/handler/yaml.py

@@ -0,0 +1,23 @@
+import yaml
+from yaml.loader import SafeLoader
+from typing import Any, Dict
+from utils.handler.vault import VaultScalar
+
+class YamlHandler:
+    @staticmethod
+    def load_yaml(path) -> Dict:
+        """Load the YAML file and wrap existing !vault entries."""
+        text = path.read_text()
+        data = yaml.load(text, Loader=SafeLoader) or {}
+        return YamlHandler.wrap_existing_vaults(data)
+
+    @staticmethod
+    def wrap_existing_vaults(node: Any) -> Any:
+        """Recursively wrap any str that begins with '$ANSIBLE_VAULT' in a VaultScalar so it dumps as a literal block."""
+        if isinstance(node, dict):
+            return {k: YamlHandler.wrap_existing_vaults(v) for k, v in node.items()}
+        if isinstance(node, list):
+            return [YamlHandler.wrap_existing_vaults(v) for v in node]
+        if isinstance(node, str) and node.lstrip().startswith("$ANSIBLE_VAULT"):
+            return VaultScalar(node)
+        return node