From f5d48347e55de2f0457835e0b6dcc71c5814c744 Mon Sep 17 00:00:00 2001
From: Correl Roush <correlr@aweber.com>
Date: Tue, 23 Jan 2024 12:59:49 -0500
Subject: [PATCH] Add copy-tree command

---
 ssm.py | 90 ++++++++++++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 78 insertions(+), 12 deletions(-)

diff --git a/ssm.py b/ssm.py
index 6818b16..aa2e0cc 100644
--- a/ssm.py
+++ b/ssm.py
@@ -10,6 +10,7 @@ import botocore.client
 import botocore.session
 import rich.console
 import rich.table
+import rich.text
 import typer
 
 app = typer.Typer()
@@ -39,6 +40,22 @@ def get_client(profile: str) -> botocore.client.BaseClient:
         sys.exit(1)
 
 
+def get_parameters(
+    client: botocore.client.BaseClient, path: SSMPath, recursive: bool
+) -> typing.Iterable[dict]:
+    ...
+    results = client.get_paginator("get_parameters_by_path").paginate(
+        Path=str(path),
+        Recursive=recursive,
+        WithDecryption=True,
+    )
+    for parameter in (
+        result for chunk in results for result in chunk.get("Parameters")
+    ):
+        if parameter.get("Name"):
+            yield parameter
+
+
 @app.command()
 def profiles() -> None:
     """List available AWS profiles."""
@@ -46,8 +63,8 @@ def profiles() -> None:
         print(profile)
 
 
-@app.command()
-def list(
+@app.command("list")
+def list_parameters(
     profile: ProfileArg,
     path: pathlib.Path,
     recursive: bool = True,
@@ -63,21 +80,17 @@ def list(
         "Type",
         title=f"{root} ({profile})",
     )
-    results = client.get_paginator("get_parameters_by_path").paginate(
-        Path=str(root),
-        Recursive=recursive,
-        WithDecryption=True,
-    )
-    for parameter in (
-        result for chunk in results for result in chunk.get("Parameters")
-    ):
-        if name := parameter.get("Name"):
+    try:
+        for parameter in get_parameters(client, root, recursive):
             table.add_row(
-                str(pathlib.Path(name).relative_to(root)),
+                str(pathlib.Path(parameter["Name"]).relative_to(root)),
                 parameter.get("Value"),
                 parameter.get("Description"),
                 parameter.get("Type"),
             )
+    except client.exceptions.ClientError as e:
+        stderr.print(str(e), style="bold red")
+        sys.exit(1)
     console.print(table)
 
 
@@ -122,5 +135,58 @@ def unset(profile: ProfileArg, path: str) -> None:
         sys.exit(1)
 
 
+@app.command()
+def copy_tree(
+    source_profile: ProfileArg,
+    source_path: str,
+    dest_profile: ProfileArg,
+    dest_path: str,
+    replacement_pairs: list[str] = typer.Option([], "--replace", "-r"),
+    recursive: bool = True,
+):
+    """Copy parameters from SRC_PATH to DEST_PATH."""
+
+    def parse_replacements(values: list[str]) -> list[tuple[str, str]]:
+        pairs: list[tuple[str, str]] = []
+        for value in values:
+            a, b = value.split("=", 1)
+            pairs.append((a, b))
+        return pairs
+
+    replacements = parse_replacements(replacement_pairs)
+
+    def replace(value: str) -> str:
+        for old, new in replacements:
+            value = value.replace(old, new)
+        return value
+
+    source = get_client(source_profile)
+    destination = get_client(dest_profile)
+    sources = {
+        str(pathlib.Path(p["Name"]).relative_to(SSMPath(source_path))): p
+        for p in get_parameters(source, SSMPath(source_path), recursive=recursive)
+    }
+    targets = {
+        str(pathlib.Path(p["Name"]).relative_to(SSMPath(source_path))): p
+        for p in get_parameters(destination, SSMPath(dest_path), recursive=recursive)
+    }
+    table = rich.table.Table("Path", "Old Value", "New Value")
+    for name, param in sources.items():
+        old = targets[name]["Value"] if name in targets else None
+        new = replace(param["Value"])
+        table.add_row(
+            name,
+            rich.text.Text(old, style="red")
+            if old
+            else rich.text.Text("Not defined", style="bright_black italic"),
+            rich.text.Text(new, style="green"),
+        )
+    stdout.print(table)
+    confirmed = typer.confirm("Are you sure you want to apply the above changes?")
+    if not confirmed:
+        stderr.print("No changes applied", style="yellow italic")
+        sys.exit(1)
+
+
 if __name__ == "__main__":
     app()