Inject cloud-init SSH rule on wake for Tailscale-locked snapshots

When waking a droplet hibernated with Tailscale lockdown, the snapshot's
UFW rules block SSH on eth0 while Tailscale is logged out — creating a
deadlock where SSH is needed to re-setup Tailscale but is blocked by the
firewall. Fix by injecting a cloud-init user_data script that opens SSH
on eth0 at first boot, which is cleaned up when lock_down_to_tailscale()
resets UFW during Tailscale re-setup.

Also adds wait_for_ssh() to verify SSH reachability before attempting
Tailscale setup (replacing a blind sleep), and enhances the security
warning when --no-tailscale is used with a previously locked snapshot.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

dropkit/api.py

@@ -930,6 +930,7 @@ class DigitalOceanAPI:
         snapshot_id: int,
         tags: list[str],
         ssh_keys: list[int] | None = None,
+        user_data: str | None = None,
     ) -> dict[str, Any]:
         """
         Create a new droplet from a snapshot image.
@@ -941,6 +942,7 @@ class DigitalOceanAPI:
             snapshot_id: Snapshot ID to restore from
             tags: List of tags to apply
             ssh_keys: List of SSH key IDs for root access (optional)
+            user_data: Cloud-init user-data script to run on first boot (optional)
 
         Returns:
             Droplet object from API response
@@ -961,5 +963,8 @@ class DigitalOceanAPI:
         if ssh_keys:
             payload["ssh_keys"] = ssh_keys
 
+        if user_data is not None:
+            payload["user_data"] = user_data
+
         response = self._request("POST", "/droplets", json=payload)
         return response.get("droplet", {})

dropkit/main.py

@@ -969,6 +969,63 @@ def is_droplet_tailscale_locked(config: DropkitConfig, droplet_name: str) -> boo
     return is_tailscale_ip(current_ip)
 
 
+def wait_for_ssh(
+    ssh_hostname: str,
+    timeout: int = 60,
+    interval: int = 5,
+) -> bool:
+    """
+    Poll until SSH is reachable on a host.
+
+    Runs a lightweight SSH command (true) with a short ConnectTimeout,
+    retrying until success or the overall timeout is exceeded.
+
+    Returns:
+        True if SSH became reachable, False if timed out.
+    """
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        try:
+            result = subprocess.run(
+                [
+                    "ssh",
+                    "-o",
+                    "StrictHostKeyChecking=no",
+                    "-o",
+                    "UserKnownHostsFile=/dev/null",
+                    "-o",
+                    "ConnectTimeout=5",
+                    "-o",
+                    "BatchMode=yes",
+                    ssh_hostname,
+                    "true",
+                ],
+                capture_output=True,
+                timeout=10,
+            )
+            if result.returncode == 0:
+                return True
+        except (subprocess.TimeoutExpired, subprocess.SubprocessError, OSError):
+            pass
+        remaining = deadline - time.monotonic()
+        if remaining > 0:
+            time.sleep(min(interval, remaining))
+    return False
+
+
+def build_wake_user_data(was_tailscale_locked: bool) -> str | None:
+    """Build cloud-init user_data script for wake, or None if not needed."""
+    if not was_tailscale_locked:
+        return None
+    return (
+        "#!/bin/bash\n"
+        "set -euo pipefail\n"
+        "echo 'dropkit-wake: opening SSH on eth0' | logger -t dropkit\n"
+        "ufw allow in on eth0 to any port 22\n"
+        "echo 'dropkit-wake: SSH rule applied successfully' | logger -t dropkit\n"
+    )
+
+
 def add_temporary_ssh_rule(ssh_hostname: str, verbose: bool = False) -> bool:
     """
     Add temporary UFW rule to allow SSH on public interface (eth0).
@@ -1008,7 +1065,9 @@ def add_temporary_ssh_rule(ssh_hostname: str, verbose: bool = False) -> bool:
 
         if result.returncode != 0:
             stderr_msg = result.stderr.strip()[:200] if result.stderr else "no error output"
-            console.print(f"[dim]SSH rule command failed (exit {result.returncode}): {stderr_msg}[/dim]")
+            console.print(
+                f"[dim]SSH rule command failed (exit {result.returncode}): {stderr_msg}[/dim]"
+            )
             return False
         return True
     except (subprocess.TimeoutExpired, subprocess.SubprocessError) as e:
@@ -4051,6 +4110,13 @@ def wake(
         # Build tags for new droplet
         tags_list = build_droplet_tags(username, list(config.defaults.extra_tags))
 
+        # Ensure SSH access on eth0 for Tailscale-locked snapshots.
+        # Older snapshots may have been created without the temporary SSH rule.
+        # The eth0 rule is cleaned up by lock_down_to_tailscale() during Tailscale re-setup.
+        wake_user_data = build_wake_user_data(was_tailscale_locked)
+        if wake_user_data:
+            console.print("[dim]Injecting cloud-init script to open SSH access on eth0...[/dim]")
+
         droplet = api.create_droplet_from_snapshot(
             name=droplet_name,
             region=original_region,
@@ -4058,6 +4124,7 @@ def wake(
             snapshot_id=snapshot_id,
             tags=tags_list,
             ssh_keys=config.cloudinit.ssh_key_ids,
+            user_data=wake_user_data,
         )
 
         droplet_id = droplet.get("id")
@@ -4110,13 +4177,16 @@ def wake(
             ssh_hostname = get_ssh_hostname(droplet_name)
             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]"
+                    "[yellow]⚠ Security notice:[/yellow] "
+                    "Original droplet had Tailscale lockdown enabled."
                 )
                 console.print(
-                    f"[dim]Enable Tailscale later with: "
+                    "[yellow]SSH port 22 is now open on the public interface "
+                    "and will remain so until Tailscale is re-enabled.[/yellow]"
+                )
+                console.print(
+                    f"[dim]Re-enable Tailscale lockdown with: "
                     f"[cyan]dropkit enable-tailscale {droplet_name}[/cyan][/dim]"
                 )
             else:
@@ -4124,22 +4194,30 @@ def wake(
                 console.print(
                     "[dim]Original droplet had Tailscale lockdown - re-setting up Tailscale...[/dim]"
                 )
-                # Wait a bit for droplet to be fully ready for SSH
-                console.print("[dim]Waiting for droplet to be ready for SSH...[/dim]")
-                time.sleep(10)
-
-                # Re-setup Tailscale (clean state from hibernate logout)
-                tailscale_ip = setup_tailscale(ssh_hostname, username, config, verbose)
-
-                if not tailscale_ip:
+                console.print("[dim]Waiting for SSH to become reachable...[/dim]")
+                if not wait_for_ssh(ssh_hostname, timeout=60):
+                    console.print("[red]Error: SSH is not reachable on the public IP.[/red]")
                     console.print(
-                        "[yellow]⚠[/yellow] Tailscale setup incomplete. "
-                        "Public SSH access remains available."
+                        "[dim]The cloud-init script may have failed. "
+                        "Check /var/log/cloud-init-output.log via the DigitalOcean console.[/dim]"
                     )
                     console.print(
-                        f"[dim]Complete setup later with: "
+                        f"[dim]Once resolved, run: "
                         f"[cyan]dropkit enable-tailscale {droplet_name}[/cyan][/dim]"
                     )
+                else:
+                    # Re-setup Tailscale (clean state from hibernate logout)
+                    tailscale_ip = setup_tailscale(ssh_hostname, username, config, verbose)
+
+                    if not tailscale_ip:
+                        console.print(
+                            "[yellow]⚠[/yellow] Tailscale setup incomplete. "
+                            "Public SSH access remains available."
+                        )
+                        console.print(
+                            f"[dim]Complete setup later with: "
+                            f"[cyan]dropkit enable-tailscale {droplet_name}[/cyan][/dim]"
+                        )
 
         # Prompt to delete snapshot
         console.print()

tests/test_api.py

@@ -162,3 +162,42 @@ class TestUntagResource:
         api = DigitalOceanAPI("fake-token")
         with pytest.raises(ValueError, match="Cannot remove protected tag: firewall"):
             api.untag_resource("firewall", "12345", "droplet")
+
+
+class TestCreateDropletFromSnapshot:
+    """Tests for create_droplet_from_snapshot method."""
+
+    @patch.object(DigitalOceanAPI, "_request")
+    def test_with_user_data(self, mock_request):
+        """Test that user_data is included in the API payload when provided."""
+        mock_request.return_value = {"droplet": {"id": 123}}
+        api = DigitalOceanAPI("fake-token")
+
+        api.create_droplet_from_snapshot(
+            name="test",
+            region="nyc3",
+            size="s-1vcpu-1gb",
+            snapshot_id=456,
+            tags=["owner:test"],
+            user_data="#!/bin/bash\nufw allow in on eth0 to any port 22\n",
+        )
+
+        payload = mock_request.call_args[1]["json"]
+        assert payload["user_data"] == "#!/bin/bash\nufw allow in on eth0 to any port 22\n"
+
+    @patch.object(DigitalOceanAPI, "_request")
+    def test_without_user_data(self, mock_request):
+        """Test that user_data is omitted from payload when not provided."""
+        mock_request.return_value = {"droplet": {"id": 123}}
+        api = DigitalOceanAPI("fake-token")
+
+        api.create_droplet_from_snapshot(
+            name="test",
+            region="nyc3",
+            size="s-1vcpu-1gb",
+            snapshot_id=456,
+            tags=["owner:test"],
+        )
+
+        payload = mock_request.call_args[1]["json"]
+        assert "user_data" not in payload

tests/test_main_helpers.py

@@ -11,6 +11,7 @@ from dropkit.main import (
     _resize_hibernated_snapshot,
     add_temporary_ssh_rule,
     build_droplet_tags,
+    build_wake_user_data,
     complete_droplet_or_snapshot_name,
     find_snapshot_action,
     get_droplet_name_from_snapshot,
@@ -19,6 +20,7 @@ from dropkit.main import (
     get_user_tag,
     is_droplet_tailscale_locked,
     prepare_for_hibernate,
+    wait_for_ssh,
 )
 
 
@@ -548,3 +550,66 @@ class TestCompleteDropletOrSnapshotName:
         """Test that live droplet names appear before snapshot names."""
         result = complete_droplet_or_snapshot_name("")
         assert result.index("live-vm") < result.index("snap-vm")
+
+
+class TestBuildWakeUserData:
+    """Tests for build_wake_user_data function."""
+
+    def test_returns_script_when_tailscale_locked(self):
+        """Test that a cloud-init script is returned for Tailscale-locked snapshots."""
+        result = build_wake_user_data(was_tailscale_locked=True)
+        assert result is not None
+        assert result.startswith("#!/bin/bash\n")
+        assert "ufw allow in on eth0 to any port 22" in result
+
+    def test_returns_none_when_not_tailscale_locked(self):
+        """Test that None is returned when snapshot was not Tailscale-locked."""
+        result = build_wake_user_data(was_tailscale_locked=False)
+        assert result is None
+
+
+class TestWaitForSsh:
+    """Tests for wait_for_ssh function."""
+
+    @patch("dropkit.main.subprocess.run")
+    @patch("dropkit.main.time.monotonic")
+    @patch("dropkit.main.time.sleep")
+    def test_returns_true_on_immediate_success(self, mock_sleep, mock_monotonic, mock_run):
+        """Test that wait_for_ssh returns True when SSH is immediately reachable."""
+        mock_monotonic.side_effect = [0, 0]  # start, first check
+        mock_run.return_value = MagicMock(returncode=0)
+        assert wait_for_ssh("dropkit.test", timeout=10) is True
+
+    @patch("dropkit.main.subprocess.run")
+    @patch("dropkit.main.time.monotonic")
+    @patch("dropkit.main.time.sleep")
+    def test_returns_true_after_retries(self, mock_sleep, mock_monotonic, mock_run):
+        """Test that wait_for_ssh retries and succeeds on second attempt."""
+        mock_monotonic.side_effect = [0, 0, 5, 5]  # start, fail check, sleep check, success check
+        mock_run.side_effect = [
+            MagicMock(returncode=255),  # first attempt fails
+            MagicMock(returncode=0),  # second attempt succeeds
+        ]
+        assert wait_for_ssh("dropkit.test", timeout=60) is True
+
+    @patch("dropkit.main.subprocess.run")
+    @patch("dropkit.main.time.monotonic")
+    @patch("dropkit.main.time.sleep")
+    def test_returns_false_on_timeout(self, mock_sleep, mock_monotonic, mock_run):
+        """Test that wait_for_ssh returns False when timeout is exceeded."""
+        # deadline calc, while check, remaining check, while check (past deadline)
+        mock_monotonic.side_effect = [0, 0, 5, 61]
+        mock_run.return_value = MagicMock(returncode=255)
+        assert wait_for_ssh("dropkit.test", timeout=60) is False
+
+    @patch("dropkit.main.subprocess.run")
+    @patch("dropkit.main.time.monotonic")
+    @patch("dropkit.main.time.sleep")
+    def test_handles_timeout_exception(self, mock_sleep, mock_monotonic, mock_run):
+        """Test that subprocess.TimeoutExpired is caught and retried."""
+        import subprocess
+
+        # deadline calc, while check, remaining check, while check (past deadline)
+        mock_monotonic.side_effect = [0, 0, 5, 61]
+        mock_run.side_effect = subprocess.TimeoutExpired("ssh", 10)
+        assert wait_for_ssh("dropkit.test", timeout=60) is False