diff --git a/README.md b/README.md index 68b7c68..09a59ef 100644 --- a/README.md +++ b/README.md @@ -178,6 +178,9 @@ dropkit destroy my-droplet 3. Snapshots are tagged with `owner:` and `size:` for tracking 4. After waking, you're prompted to delete the snapshot (default: yes) +Wake waits up to 15 minutes by default. Use `dropkit wake my-droplet --timeout 1800` +to wait up to 30 minutes. + **Note:** Snapshots are billed at $0.06/GB/month, which is typically much cheaper than keeping a droplet running. ### Cloud-Init Customization diff --git a/dropkit/api.py b/dropkit/api.py index 47e09b3..d3c93a6 100644 --- a/dropkit/api.py +++ b/dropkit/api.py @@ -68,6 +68,10 @@ def __init__(self, message: str, status_code: int | None = None): super().__init__(message) +class DropletWaitTimeoutError(DigitalOceanAPIError): + """The local wait expired; droplet creation may still be in progress.""" + + PROTECTED_TAGS = {"owner", "firewall"} @@ -405,12 +409,13 @@ def wait_for_droplet_active( Raises: ValueError: If droplet_id is not positive - DigitalOceanAPIError: If timeout is reached or droplet enters error state + DropletWaitTimeoutError: If the local wait expires + DigitalOceanAPIError: If the droplet enters an error state or an API request fails """ import time self._validate_positive_int(droplet_id, "droplet_id") - start_time = time.time() + start_time = time.monotonic() while True: droplet = self.get_droplet(droplet_id) @@ -423,9 +428,9 @@ def wait_for_droplet_active( f"Droplet entered error state: {droplet.get('name', droplet_id)}" ) - elapsed = time.time() - start_time - if elapsed > timeout: - raise DigitalOceanAPIError( + elapsed = time.monotonic() - start_time + if elapsed >= timeout: + raise DropletWaitTimeoutError( f"Timeout waiting for droplet to become active (waited {elapsed:.0f}s)" ) diff --git a/dropkit/main.py b/dropkit/main.py index 7b31fa1..ad5561d 100644 --- a/dropkit/main.py +++ b/dropkit/main.py @@ -16,7 +16,7 @@ from rich.prompt import Confirm, Prompt from rich.table import Table -from dropkit.api import DigitalOceanAPI, DigitalOceanAPIError +from dropkit.api import DigitalOceanAPI, DigitalOceanAPIError, DropletWaitTimeoutError from dropkit.cloudinit import render_cloud_init from dropkit.config import DEFAULT_IMAGE, DEFAULT_REGION, DEFAULT_SIZE, Config, DropkitConfig from dropkit.lock import requires_lock @@ -4036,17 +4036,21 @@ def wake( droplet_name: str = typer.Argument( ..., autocompletion=complete_snapshot_name, help="Name of the hibernated droplet to restore" ), + timeout: int = typer.Option( + 900, "--timeout", min=1, help="Seconds to wait for restoration (default: 15 minutes)" + ), no_tailscale: bool = typer.Option(False, "--no-tailscale", help="Skip Tailscale VPN re-setup"), ): """ Wake a hibernated droplet (restore from snapshot). This will create a new droplet from the hibernated snapshot. + Run again after a timeout to resume setup of the same restored droplet. After successful restoration, you'll be prompted to delete the snapshot. If the original droplet had Tailscale lockdown enabled, this command will re-setup Tailscale after the droplet becomes active. Use --no-tailscale to - skip this and keep public SSH access. + skip this. When resuming, existing Tailscale SSH access is preserved. Use 'dropkit destroy ' to delete a hibernated snapshot without restoring. """ @@ -4064,10 +4068,6 @@ def wake( # Check if a droplet with this name already exists existing_droplet, _ = find_user_droplet(api, droplet_name) - if existing_droplet: - console.print(f"[red]Error: A droplet named '{droplet_name}' already exists.[/red]") - console.print("[dim]Destroy or rename the existing droplet first.[/dim]") - raise typer.Exit(1) # Find the hibernated snapshot snapshot_name = get_snapshot_name(droplet_name) @@ -4079,6 +4079,13 @@ def wake( if not snapshot: console.print(f"[red]Error: No hibernated snapshot found for '{droplet_name}'[/red]") console.print(f"[dim]Expected snapshot name: {snapshot_name}[/dim]") + if existing_droplet: + console.print( + f"[dim]The droplet already exists. Check it with: dropkit info {droplet_name}[/dim]" + ) + console.print( + f"[dim]To configure SSH, run: dropkit config-ssh {droplet_name}[/dim]" + ) raise typer.Exit(1) snapshot_id_str = snapshot.get("id") @@ -4087,6 +4094,23 @@ def wake( raise typer.Exit(1) snapshot_id = int(snapshot_id_str) # API returns string, convert to int + if existing_droplet: + # A matching name alone is insufficient: only resume a restore of this snapshot. + source_image_id = existing_droplet.get("image", {}).get("id") + if str(source_image_id) != str(snapshot_id): + console.print( + f"[red]Error: A droplet named '{droplet_name}' already exists, " + "but was not restored from this snapshot.[/red]" + ) + console.print(f"[dim]Inspect it with: dropkit info {droplet_name}[/dim]") + raise typer.Exit(1) + status = existing_droplet.get("status") + if status not in {"new", "active"}: + console.print(f"[red]Cannot resume wake: droplet status is '{status}'.[/red]") + if status == "off": + console.print(f"[dim]Power it on first: dropkit on {droplet_name}[/dim]") + raise typer.Exit(1) + # Get snapshot details size_gb = snapshot.get("size_gigabytes", 0) regions = snapshot.get("regions", []) @@ -4120,33 +4144,60 @@ def wake( ) console.print() - # Create droplet from snapshot - console.print(f"[dim]Creating droplet '{droplet_name}' from snapshot...[/dim]") - - # Build tags for new droplet - tags_list = build_droplet_tags(username, list(config.defaults.extra_tags)) - - droplet = api.create_droplet_from_snapshot( - name=droplet_name, - region=original_region, - size=original_size, - snapshot_id=snapshot_id, - tags=tags_list, - ssh_keys=config.cloudinit.ssh_key_ids, - ) + if existing_droplet: + droplet = existing_droplet + console.print(f"[dim]Resuming wake for existing droplet '{droplet_name}'...[/dim]") + else: + # Create droplet from snapshot + console.print(f"[dim]Creating droplet '{droplet_name}' from snapshot...[/dim]") + + # Build tags for new droplet + tags_list = build_droplet_tags(username, list(config.defaults.extra_tags)) + + droplet = api.create_droplet_from_snapshot( + name=droplet_name, + region=original_region, + size=original_size, + snapshot_id=snapshot_id, + tags=tags_list, + ssh_keys=config.cloudinit.ssh_key_ids, + ) droplet_id = droplet.get("id") if not droplet_id: console.print("[red]Error: Failed to get droplet ID from API response[/red]") raise typer.Exit(1) - console.print(f"[green]✓[/green] Droplet created (ID: [cyan]{droplet_id}[/cyan])") + console.print(f"[green]✓[/green] Droplet ID: [cyan]{droplet_id}[/cyan]") # Wait for droplet to become active console.print("[dim]Waiting for droplet to become active...[/dim]") - with console.status("[cyan]Waiting...[/cyan]"): - active_droplet = api.wait_for_droplet_active(droplet_id) + try: + with console.status("[cyan]Waiting...[/cyan]"): + active_droplet = api.wait_for_droplet_active(droplet_id, timeout=timeout) + except DropletWaitTimeoutError: + console.print( + f"[yellow]Stopped waiting after {timeout}s. DigitalOcean may still be " + f"restoring droplet '{droplet_name}' (ID: {droplet_id}).[/yellow]" + ) + console.print("[dim]The droplet and snapshot have been kept.[/dim]") + console.print(f"[dim]Check status: dropkit info {droplet_name}[/dim]") + console.print( + f"[dim]Resume setup: dropkit wake {droplet_name} --timeout {timeout}" + f"{' --no-tailscale' if no_tailscale else ''}[/dim]" + ) + raise typer.Exit(1) + except DigitalOceanAPIError: + console.print( + f"[yellow]Could not check readiness of droplet {droplet_id}. " + "The droplet and snapshot have been kept.[/yellow]" + ) + console.print( + f"[dim]Check status with dropkit info {droplet_name}, " + "then rerun wake to resume setup.[/dim]" + ) + raise # Get IP address networks = active_droplet.get("networks", {}) @@ -4162,23 +4213,34 @@ def wake( console.print(f"[green]✓[/green] Droplet is active (IP: [cyan]{ip_address}[/cyan])") else: console.print("[green]✓[/green] Droplet is active") - console.print("[yellow]⚠[/yellow] Could not determine IP address") + console.print("[red]Could not determine public IP address; snapshot kept.[/red]") + console.print(f"[dim]Retry setup: dropkit wake {droplet_name}[/dim]") + raise typer.Exit(1) # Add SSH config entry if ip_address and config.ssh.auto_update: try: console.print("[dim]Configuring SSH...[/dim]") ssh_hostname = get_ssh_hostname(droplet_name) - add_ssh_host( - config_path=config.ssh.config_path, - host_name=ssh_hostname, - hostname=ip_address, - user=username, - identity_file=config.ssh.identity_file, - ) - console.print("[green]✓[/green] SSH config updated") + if existing_droplet and is_droplet_tailscale_locked(config, droplet_name): + # A previous wake may have blocked public SSH before being interrupted. + console.print("[green]✓[/green] Preserved Tailscale SSH config") + else: + add_ssh_host( + config_path=config.ssh.config_path, + host_name=ssh_hostname, + hostname=ip_address, + user=username, + identity_file=config.ssh.identity_file, + ) + console.print("[green]✓[/green] SSH config updated") except Exception as e: - console.print(f"[yellow]⚠[/yellow] Could not update SSH config: {e}") + console.print(f"[red]Could not update SSH config: {e}[/red]") + console.print( + f"[dim]Snapshot kept. Fix the SSH configuration error, " + f"then retry: dropkit wake {droplet_name}[/dim]" + ) + raise typer.Exit(1) # Handle Tailscale re-setup if the original droplet had Tailscale lockdown if was_tailscale_locked and ip_address: @@ -4186,10 +4248,7 @@ def wake( if no_tailscale: console.print() console.print("[yellow]⚠[/yellow] Original droplet had Tailscale lockdown enabled.") - console.print( - "[dim]Skipping Tailscale setup (--no-tailscale). " - "Public SSH access available.[/dim]" - ) + console.print("[dim]Skipping Tailscale setup (--no-tailscale).[/dim]") console.print( f"[dim]Enable Tailscale later with: " f"[cyan]dropkit enable-tailscale {droplet_name}[/cyan][/dim]" @@ -4207,10 +4266,7 @@ def wake( tailscale_ip = setup_tailscale(ssh_hostname, username, config) if not tailscale_ip: - console.print( - "[yellow]⚠[/yellow] Tailscale setup incomplete. " - "Public SSH access remains available." - ) + console.print("[yellow]⚠[/yellow] Tailscale setup incomplete.") console.print( f"[dim]Complete setup later with: " f"[cyan]dropkit enable-tailscale {droplet_name}[/cyan][/dim]" diff --git a/tests/test_wake.py b/tests/test_wake.py new file mode 100644 index 0000000..9f41fec --- /dev/null +++ b/tests/test_wake.py @@ -0,0 +1,285 @@ +"""Regression tests for interrupted snapshot restores.""" + +from contextlib import ExitStack +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +from typer.testing import CliRunner + +from dropkit.api import DigitalOceanAPI, DigitalOceanAPIError, DropletWaitTimeoutError +from dropkit.main import app, setup_tailscale +from dropkit.ssh_config import add_ssh_host, get_ssh_host_ip + +runner = CliRunner() + + +@pytest.fixture +def wake_env(tmp_path): + """Isolate all cloud, SSH, VPN, and lock effects.""" + api = MagicMock() + api.get_username.return_value = "testuser" + snapshot = { + "id": "456", + "regions": ["lon1"], + "tags": ["owner:testuser", "size:s-8vcpu-16gb"], + } + droplet = { + "id": 123, + "name": "test-droplet", + "image": {"id": 456}, + "status": "active", + "networks": {"v4": [{"type": "public", "ip_address": "192.0.2.1"}]}, + } + api.get_snapshot_by_name.return_value = snapshot + api.create_droplet_from_snapshot.return_value = droplet + api.wait_for_droplet_active.return_value = droplet + config = SimpleNamespace( + defaults=SimpleNamespace(extra_tags=[], size="s-8vcpu-16gb"), + cloudinit=SimpleNamespace(ssh_key_ids=[789]), + ssh=SimpleNamespace( + auto_update=True, config_path="unused", identity_file="~/.ssh/test_key" + ), + ) + with ExitStack() as stack: + stack.enter_context(patch("dropkit.lock.LOCK_FILE", tmp_path / "dropkit.lock")) + stack.enter_context( + patch( + "dropkit.main.load_config_and_api", + return_value=(SimpleNamespace(config=config), api), + ) + ) + find = stack.enter_context( + patch("dropkit.main.find_user_droplet", return_value=(None, "testuser")) + ) + ssh = stack.enter_context(patch("dropkit.main.add_ssh_host")) + vpn = stack.enter_context(patch("dropkit.main.setup_tailscale", return_value="100.64.0.1")) + stack.enter_context(patch("dropkit.main.time.sleep")) + yield SimpleNamespace( + api=api, + snapshot=snapshot, + droplet=droplet, + find=find, + ssh=ssh, + vpn=vpn, + config=config, + ) + + +def test_timeout_then_retry_finishes_same_droplet(wake_env): + """A slow restore must retain its resources and resume without creating another VM.""" + env = wake_env + env.api.wait_for_droplet_active.side_effect = DropletWaitTimeoutError("timed out") + result = runner.invoke(app, ["wake", "test-droplet", "--timeout", "300", "--no-tailscale"]) + assert result.exit_code == 1 + assert "DigitalOcean may still be restoring" in result.output + assert "dropkit info test-droplet" in result.output + assert "dropkit wake test-droplet --timeout 300 --no-tailscale" in result.output + env.ssh.assert_not_called() + env.api.delete_snapshot.assert_not_called() + env.api.wait_for_droplet_active.assert_called_once_with(123, timeout=300) + + env.find.return_value = (env.droplet, "testuser") + env.api.wait_for_droplet_active.side_effect = None + result = runner.invoke(app, ["wake", "test-droplet"], input="no\n") + assert result.exit_code == 0, result.output + assert "Resuming wake" in result.output + assert "ssh dropkit.test-droplet" in result.output + env.api.create_droplet_from_snapshot.assert_called_once() + env.api.wait_for_droplet_active.assert_called_with(123, timeout=900) + env.ssh.assert_called_once_with( + config_path="unused", + host_name="dropkit.test-droplet", + hostname="192.0.2.1", + user="testuser", + identity_file="~/.ssh/test_key", + ) + env.api.delete_snapshot.assert_not_called() + + +@pytest.mark.parametrize("status", ["new", "active"]) +def test_resume_retains_tailscale_setup_and_snapshot_prompt(wake_env, status): + env = wake_env + env.droplet["status"] = status + env.find.return_value = (env.droplet, "testuser") + env.snapshot["tags"].append("tailscale-lockdown") + result = runner.invoke(app, ["wake", "test-droplet", "--timeout", "1800"], input="yes\n") + assert result.exit_code == 0, result.output + env.api.create_droplet_from_snapshot.assert_not_called() + env.api.wait_for_droplet_active.assert_called_once_with(123, timeout=1800) + env.vpn.assert_called_once() + env.api.delete_snapshot.assert_called_once_with(456) + + +def test_resume_no_tailscale(wake_env): + env = wake_env + env.find.return_value = (env.droplet, "testuser") + env.snapshot["tags"].append("tailscale-lockdown") + result = runner.invoke(app, ["wake", "test-droplet", "--no-tailscale"], input="no\n") + assert result.exit_code == 0, result.output + env.vpn.assert_not_called() + + +@pytest.mark.parametrize("no_tailscale", [False, True]) +def test_retry_after_firewall_lockdown_preserves_tailscale_ssh(wake_env, tmp_path, no_tailscale): + env = wake_env + env.config.ssh.config_path = str(tmp_path / "ssh_config") + env.config.tailscale = SimpleNamespace(lock_down_firewall=True) + env.snapshot["tags"].append("tailscale-lockdown") + env.ssh.side_effect = add_ssh_host + env.vpn.side_effect = setup_tailscale + locked = False + addresses = [] + + def connect(ssh_hostname, verbose): + """Simulate an authenticated node that rejects public SSH after lockdown.""" + address = get_ssh_host_ip(env.config.ssh.config_path, ssh_hostname) + addresses.append(address) + assert address == ("100.64.0.1" if locked else "192.0.2.1") + + def lockdown(ssh_hostname, verbose): + nonlocal locked + assert get_ssh_host_ip(env.config.ssh.config_path, ssh_hostname) == "100.64.0.1" + locked = True + return True + + with ( + patch("dropkit.main.run_tailscale_up", side_effect=connect), + patch("dropkit.main.wait_for_tailscale_ip", return_value="100.64.0.1"), + patch("dropkit.main.check_local_tailscale", return_value=True), + patch("dropkit.main.lock_down_to_tailscale", side_effect=lockdown), + patch("dropkit.main.verify_tailscale_ssh", return_value=True), + ): + with patch("dropkit.main.Prompt.ask", side_effect=KeyboardInterrupt): + result = runner.invoke(app, ["wake", "test-droplet"]) + assert result.exit_code == 130, result.output + assert locked + env.api.delete_snapshot.assert_not_called() + assert get_ssh_host_ip(env.config.ssh.config_path, "dropkit.test-droplet") == "100.64.0.1" + + env.find.return_value = (env.droplet, "testuser") + args = ["wake", "test-droplet"] + if no_tailscale: + args.append("--no-tailscale") + result = runner.invoke(app, args, input="yes\n") + + assert result.exit_code == 0, result.output + assert addresses == (["192.0.2.1"] if no_tailscale else ["192.0.2.1", "100.64.0.1"]) + assert get_ssh_host_ip(env.config.ssh.config_path, "dropkit.test-droplet") == "100.64.0.1" + env.api.create_droplet_from_snapshot.assert_called_once() + env.api.delete_snapshot.assert_called_once_with(456) + + +def test_fresh_restore_replaces_stale_tailscale_ssh(wake_env, tmp_path): + env = wake_env + env.config.ssh.config_path = str(tmp_path / "ssh_config") + add_ssh_host(env.config.ssh.config_path, "dropkit.test-droplet", "100.64.0.2", "testuser") + env.ssh.side_effect = add_ssh_host + + result = runner.invoke(app, ["wake", "test-droplet"], input="no\n") + + assert result.exit_code == 0, result.output + assert get_ssh_host_ip(env.config.ssh.config_path, "dropkit.test-droplet") == "192.0.2.1" + + +@pytest.mark.parametrize("image", [{"id": 999}, {}, {"id": None}]) +def test_refuse_unrelated_existing_droplet(wake_env, image): + env = wake_env + env.droplet["image"] = image + env.find.return_value = (env.droplet, "testuser") + result = runner.invoke(app, ["wake", "test-droplet"]) + assert result.exit_code == 1 + assert "was not restored from this snapshot" in " ".join(result.output.split()) + assert "Destroy or rename" not in result.output + env.api.create_droplet_from_snapshot.assert_not_called() + env.api.wait_for_droplet_active.assert_not_called() + env.api.delete_snapshot.assert_not_called() + env.ssh.assert_not_called() + + +@pytest.mark.parametrize("status", ["off", "error", "archive"]) +def test_refuse_resume_when_droplet_cannot_become_ready(wake_env, status): + env = wake_env + env.droplet["status"] = status + env.find.return_value = (env.droplet, "testuser") + result = runner.invoke(app, ["wake", "test-droplet"]) + assert result.exit_code == 1 + assert "Cannot resume wake" in result.output + if status == "off": + assert "dropkit on test-droplet" in result.output + env.api.wait_for_droplet_active.assert_not_called() + env.api.create_droplet_from_snapshot.assert_not_called() + + +def test_missing_snapshot_gives_existing_droplet_recovery_commands(wake_env): + env = wake_env + env.find.return_value = (env.droplet, "testuser") + env.api.get_snapshot_by_name.return_value = None + result = runner.invoke(app, ["wake", "test-droplet"]) + assert result.exit_code == 1 + assert "dropkit config-ssh test-droplet" in result.output + env.api.create_droplet_from_snapshot.assert_not_called() + env.ssh.assert_not_called() + + +def test_api_error_during_wait_keeps_resources(wake_env): + env = wake_env + env.api.wait_for_droplet_active.side_effect = DigitalOceanAPIError("Network error") + result = runner.invoke(app, ["wake", "test-droplet"]) + assert result.exit_code == 1 + assert "Network error" in result.output + assert "rerun wake to resume setup" in result.output + env.ssh.assert_not_called() + env.api.delete_snapshot.assert_not_called() + + +@pytest.mark.parametrize("failure", ["missing_ip", "ssh_config"]) +def test_setup_failure_keeps_snapshot_for_retry(wake_env, failure): + env = wake_env + if failure == "missing_ip": + env.droplet["networks"]["v4"] = [] + else: + env.ssh.side_effect = OSError("read-only config") + result = runner.invoke(app, ["wake", "test-droplet"]) + assert result.exit_code == 1 + assert "snapshot kept" in result.output.lower() + env.api.delete_snapshot.assert_not_called() + + +@pytest.mark.parametrize("timeout", ["0", "-1"]) +def test_timeout_must_be_positive(wake_env, timeout): + result = runner.invoke(app, ["wake", "test-droplet", "--timeout", timeout]) + assert result.exit_code == 2 + wake_env.api.create_droplet_from_snapshot.assert_not_called() + + +def test_wait_times_out_at_deadline(): + api = DigitalOceanAPI("test-token") + with ( + patch.object(api, "get_droplet", return_value={"status": "new"}), + patch("time.monotonic", side_effect=[0, 0, 300]), + patch("time.sleep"), + pytest.raises(DropletWaitTimeoutError, match="waited 300s"), + ): + api.wait_for_droplet_active(123) + + +def test_wait_allows_restore_longer_than_five_minutes(): + api = DigitalOceanAPI("test-token") + ready = {"status": "active"} + with ( + patch.object(api, "get_droplet", side_effect=[{"status": "new"}, ready]), + patch("time.monotonic", side_effect=[0, 395]), + patch("time.sleep"), + ): + assert api.wait_for_droplet_active(123, timeout=900) == ready + + +def test_error_state_is_not_a_timeout(): + api = DigitalOceanAPI("test-token") + with ( + patch.object(api, "get_droplet", return_value={"status": "error"}), + pytest.raises(DigitalOceanAPIError, match="error state") as exc, + ): + api.wait_for_droplet_active(123) + assert not isinstance(exc.value, DropletWaitTimeoutError)