From cc1b6dca718c6b0652bd55e5f454d946184554a1 Mon Sep 17 00:00:00 2001
From: josibake <josibake@protonmail.com>
Date: Tue, 10 Sep 2024 13:49:34 +0200
Subject: [PATCH] add snapshot cli command

allow a user to snapshot a single tank (with a filter) or
snapshot all tanks with a filter. this creates a zip that can
then be loaded via the initContainer in the first commit
---
 src/warnet/control.py | 95 +++++++++++++++++++++++++++++++++++++++++++
 src/warnet/main.py    |  3 +-
 2 files changed, 97 insertions(+), 1 deletion(-)

diff --git a/src/warnet/control.py b/src/warnet/control.py
index 203130f80..cb8b06d06 100644
--- a/src/warnet/control.py
+++ b/src/warnet/control.py
@@ -18,6 +18,7 @@
     get_default_namespace,
     get_mission,
     get_pods,
+    snapshot_bitcoin_datadir,
 )
 from .process import run_command, stream_command
 
@@ -291,3 +292,97 @@ def logs(pod_name: str, follow: bool):
             print(f"Please consider waiting for the pod to become available. Encountered: {e}")
     else:
         pass  # cancelled by user
+
+
+@click.command()
+@click.argument("tank_name", required=False)
+@click.option("--all", "snapshot_all", is_flag=True, help="Snapshot all running tanks")
+@click.option(
+    "--output", "-o", type=click.Path(), default=".", help="Output directory for snapshots"
+)
+@click.option(
+    "--filter",
+    "-f",
+    type=str,
+    help="Comma-separated list of directories and/or files to include in the snapshot",
+)
+def snapshot(tank_name, snapshot_all, output, filter):
+    """Create a snapshot of a tank's Bitcoin data or snapshot all tanks"""
+    tanks = get_mission("tank")
+
+    if not tanks:
+        console.print("[bold red]No active tanks found.[/bold red]")
+        return
+
+    filter_list = [f.strip() for f in filter.split(",")] if filter else None
+    if snapshot_all:
+        snapshot_all_tanks(tanks, output, filter_list)
+    elif tank_name:
+        snapshot_single_tank(tank_name, tanks, output, filter_list)
+    else:
+        select_and_snapshot_tank(tanks, output, filter_list)
+
+
+def find_tank_by_name(tanks, tank_name):
+    for tank in tanks:
+        if tank.metadata.name == tank_name:
+            return tank
+    return None
+
+
+def snapshot_all_tanks(tanks, output_dir, filter_list):
+    with console.status("[bold yellow]Snapshotting all tanks...[/bold yellow]"):
+        for tank in tanks:
+            tank_name = tank.metadata.name
+            chain = tank.metadata.labels["chain"]
+            snapshot_tank(tank_name, chain, output_dir, filter_list)
+    console.print("[bold green]All tank snapshots completed.[/bold green]")
+
+
+def snapshot_single_tank(tank_name, tanks, output_dir, filter_list):
+    tank = find_tank_by_name(tanks, tank_name)
+    if tank:
+        chain = tank.metadata.labels["chain"]
+        snapshot_tank(tank_name, chain, output_dir, filter_list)
+    else:
+        console.print(f"[bold red]No active tank found with name: {tank_name}[/bold red]")
+
+
+def select_and_snapshot_tank(tanks, output_dir, filter_list):
+    table = Table(title="Active Tanks", show_header=True, header_style="bold magenta")
+    table.add_column("Number", style="cyan", justify="right")
+    table.add_column("Tank Name", style="green")
+
+    for idx, tank in enumerate(tanks, 1):
+        table.add_row(str(idx), tank.metadata.name)
+
+    console.print(table)
+
+    choices = [str(i) for i in range(1, len(tanks) + 1)] + ["q"]
+    choice = Prompt.ask(
+        "[bold yellow]Enter the number of the tank to snapshot, or 'q' to quit[/bold yellow]",
+        choices=choices,
+        show_choices=False,
+    )
+
+    if choice == "q":
+        console.print("[bold blue]Operation cancelled.[/bold blue]")
+        return
+
+    selected_tank = tanks[int(choice) - 1]
+    tank_name = selected_tank.metadata.name
+    chain = selected_tank.metadata.labels["chain"]
+    snapshot_tank(tank_name, chain, output_dir, filter_list)
+
+
+def snapshot_tank(tank_name, chain, output_dir, filter_list):
+    try:
+        output_path = Path(output_dir).resolve()
+        snapshot_bitcoin_datadir(tank_name, chain, str(output_path), filter_list)
+        console.print(
+            f"[bold green]Successfully created snapshot for tank: {tank_name}[/bold green]"
+        )
+    except Exception as e:
+        console.print(
+            f"[bold red]Failed to create snapshot for tank {tank_name}: {str(e)}[/bold red]"
+        )
diff --git a/src/warnet/main.py b/src/warnet/main.py
index 28b933ef2..341d5757e 100644
--- a/src/warnet/main.py
+++ b/src/warnet/main.py
@@ -2,7 +2,7 @@
 
 from .admin import admin
 from .bitcoin import bitcoin
-from .control import down, logs, run, stop
+from .control import down, logs, run, snapshot, stop
 from .deploy import deploy
 from .graph import create, graph
 from .image import image
@@ -28,6 +28,7 @@ def cli():
 cli.add_command(new)
 cli.add_command(run)
 cli.add_command(setup)
+cli.add_command(snapshot)
 cli.add_command(status)
 cli.add_command(stop)
 cli.add_command(create)