Enhance SSH configuration and security group management

- Added optional SSH security-group CIDR configuration in `env.example` to allow broader access if needed. - Improved `guess_operator_cidr` function to support multiple IP retrieval services and allow overriding CIDR via environment variable. - Refactored security group management to ensure idempotency and proper logging of security group rules. - Introduced helper function `_get_volume_or_none` for better volume management and error handling. - Updated `wait_ssh` function to enhance SSH connection handling with improved timeout settings and logging. - Added state reset logic in `cmd_up` to handle scenarios where the server is not found in the cloud.
2026-08-21 04:12:38 +03:00
parent 28019d1d09
commit d6d247875b
4 changed files with 123 additions and 39 deletions
@@ -12,6 +12,9 @@ GPU_RENT_AZ=ru-7a

 SSH_PRIVATE_KEY_PATH=
 SSH_USER=ubuntu
+# Optional: force SSH security-group CIDR (default = your public IP /32).
+# If WARP/VPN and SSH hangs: GPU_RENT_SSH_CIDR=0.0.0.0/0
+# GPU_RENT_SSH_CIDR=0.0.0.0/0

 BOOT_VOLUME_ID=
 DATA_VOLUME_ID=
@@ -44,13 +44,24 @@ def _oid(obj: Any) -> str:


 def guess_operator_cidr() -> str:
-    try:
-        response = httpx.get("https://ifconfig.me/ip", timeout=5.0)
-        ip = response.text.strip()
-        if ip.count(".") == 3 and all(p.isdigit() for p in ip.split(".")):
-            return f"{ip}/32"
-    except httpx.HTTPError:
-        pass
+    """Public egress /32 for SSH SG. Override: GPU_RENT_SSH_CIDR=0.0.0.0/0."""
+    import os
+
+    override = (os.environ.get("GPU_RENT_SSH_CIDR") or "").strip()
+    if override:
+        return override
+    for url in (
+        "https://ifconfig.me/ip",
+        "https://api.ipify.org",
+        "https://icanhazip.com",
+    ):
+        try:
+            response = httpx.get(url, timeout=5.0, follow_redirects=True)
+            ip = response.text.strip().split()[0]
+            if ip.count(".") == 3 and all(p.isdigit() for p in ip.split(".")):
+                return f"{ip}/32"
+        except httpx.HTTPError:
+            continue
    return "0.0.0.0/0"


@@ -142,27 +153,54 @@ def ensure_network(conn, log: Callable[[str], None]) -> tuple[Any, Any]:

 def ensure_security_group(conn, cidr: str, log: Callable[[str], None]) -> Any:
    sg = conn.network.find_security_group(SG_NAME)
-    if sg:
-        log(f"security group {SG_NAME} уже есть")
-        return sg
+    created = False
+    if not sg:
+        try:
+            sg = conn.network.create_security_group(
+                name=SG_NAME,
+                description="gpu-rent SSH only",
+            )
+            created = True
+        except Exception as exc:
+            raise _wrap(exc, "security group") from exc
+
+    # Ensure TCP/22 from operator CIDR (idempotent; old /32 from WARP may be stale).
+    have_cidr = False
    try:
-        sg = conn.network.create_security_group(
-            name=SG_NAME,
-            description="gpu-rent SSH only",
-        )
-        conn.network.create_security_group_rule(
-            security_group_id=sg.id,
-            direction="ingress",
-            ethertype="IPv4",
-            protocol="tcp",
-            port_range_min=22,
-            port_range_max=22,
-            remote_ip_prefix=cidr,
-        )
+        for rule in conn.network.security_group_rules(security_group_id=sg.id):
+            if (
+                getattr(rule, "direction", None) == "ingress"
+                and getattr(rule, "protocol", None) == "tcp"
+                and int(getattr(rule, "port_range_min", 0) or 0) == 22
+                and int(getattr(rule, "port_range_max", 0) or 0) == 22
+                and (getattr(rule, "remote_ip_prefix", None) or "") == cidr
+            ):
+                have_cidr = True
+                break
+    except Exception:
+        have_cidr = False
+    if not have_cidr:
+        try:
+            conn.network.create_security_group_rule(
+                security_group_id=sg.id,
+                direction="ingress",
+                ethertype="IPv4",
+                protocol="tcp",
+                port_range_min=22,
+                port_range_max=22,
+                remote_ip_prefix=cidr,
+            )
+            log(f"SG {SG_NAME}: + TCP/22 с {cidr}")
+        except Exception as exc:
+            text = str(exc).lower()
+            if "already" not in text and "conflict" not in text and "409" not in text:
+                raise _wrap(exc, "security group rule") from exc
+            log(f"SG {SG_NAME}: TCP/22 с {cidr} уже есть")
+    elif created:
        log(f"SG {SG_NAME}: TCP/22 с {cidr}")
-        return sg
-    except Exception as exc:
-        raise _wrap(exc, "security group") from exc
+    else:
+        log(f"security group {SG_NAME} уже есть (SSH {cidr})")
+    return sg


 def ensure_keypair(conn, public_key: str, log: Callable[[str], None]) -> Any:
@@ -186,6 +224,18 @@ def ensure_keypair(conn, public_key: str, log: Callable[[str], None]) -> Any:
        raise _wrap(exc, "keypair") from exc


+def _get_volume_or_none(conn, volume_id: str) -> Any | None:
+    if not volume_id:
+        return None
+    try:
+        return conn.block_storage.get_volume(volume_id)
+    except Exception as exc:
+        text = str(exc).lower()
+        if "404" in text or "not found" in text or "could not be found" in text:
+            return None
+        raise
+
+
 def ensure_boot_volume(
    conn,
    *,
@@ -197,9 +247,11 @@ def ensure_boot_volume(
    log: Callable[[str], None],
 ) -> Any:
    if existing_id:
-        vol = conn.block_storage.get_volume(existing_id)
-        log(f"boot volume из state {existing_id}")
-        return vol
+        vol = _get_volume_or_none(conn, existing_id)
+        if vol is not None:
+            log(f"boot volume из state {existing_id}")
+            return vol
+        log(f"boot volume {existing_id} нет в облаке — создаём заново")
    found = find_volumes_by_name(conn, BOOT_VOLUME_NAME)
    if found:
        log(f"boot volume {BOOT_VOLUME_NAME} уже в проекте")
@@ -245,9 +297,11 @@ def ensure_data_volume(
    log: Callable[[str], None],
 ) -> Any:
    if existing_id:
-        vol = conn.block_storage.get_volume(existing_id)
-        log(f"data volume из state {existing_id}")
-        return vol
+        vol = _get_volume_or_none(conn, existing_id)
+        if vol is not None:
+            log(f"data volume из state {existing_id}")
+            return vol
+        log(f"data volume {existing_id} нет в облаке — создаём заново")
    found = find_volumes_by_name(conn, DATA_VOLUME_NAME)
    if found:
        log(f"data volume {DATA_VOLUME_NAME} уже в проекте")
@@ -155,6 +155,16 @@ def cmd_up(
                return state
            raise CloudError(f"сервер gpu-rent в статусе {status} — разбери в панели")

+        if state.server_id:
+            log(f"в state был server {state.server_id}, в облаке нет — создаём заново")
+            state.server_id = None
+            state.server_name = None
+            state.floating_ip = None
+            state.floating_ip_id = None
+            state.bootstrapped = False
+            state.phase = "idle"
+            save_state(state)
+
        flavors = list(iter_flavors(conn))
        try:
            picked = resolve_flavor(
@@ -28,12 +28,20 @@ def wait_tcp(host: str, port: int, timeout: float = 300.0) -> None:
    raise CloudError(f"TCP {host}:{port} не открылся за {int(timeout)} с ({last})")


-def wait_ssh(cfg: Config, host: str, timeout: float = 420.0) -> None:
-    wait_tcp(host, 22, timeout=min(timeout, 240))
+def wait_ssh(cfg: Config, host: str, timeout: float = 900.0) -> None:
+    """Wait until sshd accepts our key. Paramiko banner noise is muted."""
+    import logging
+
+    logging.getLogger("paramiko").setLevel(logging.CRITICAL)
+    logging.getLogger("paramiko.transport").setLevel(logging.CRITICAL)
+
+    wait_tcp(host, 22, timeout=min(timeout, 300))
    deadline = time.time() + timeout
    key = str(cfg.ssh_private_key_path)
    last = None
+    attempt = 0
    while time.time() < deadline:
+        attempt += 1
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
@@ -41,9 +49,11 @@ def wait_ssh(cfg: Config, host: str, timeout: float = 420.0) -> None:
                hostname=host,
                username=cfg.ssh_user,
                key_filename=key,
-                timeout=12,
-                banner_timeout=12,
-                auth_timeout=12,
+                timeout=20,
+                banner_timeout=60,
+                auth_timeout=30,
+                allow_agent=False,
+                look_for_keys=False,
            )
            client.close()
            return
@@ -53,8 +63,15 @@ def wait_ssh(cfg: Config, host: str, timeout: float = 420.0) -> None:
                client.close()
            except Exception:
                pass
-            time.sleep(5)
-    raise CloudError(f"SSH {cfg.ssh_user}@{host} не принял ключ ({last})")
+            if attempt == 1 or attempt % 6 == 0:
+                # Progress without paramiko traceback spam.
+                print(f"жду SSH {cfg.ssh_user}@{host}… ({type(exc).__name__})", flush=True)
+            time.sleep(8)
+    raise CloudError(
+        f"SSH {cfg.ssh_user}@{host} не принял ключ за {int(timeout)} с ({last}). "
+        "Часто cloud-init ещё поднимает sshd на GPU-образе. "
+        "Если SG узкий (VPN/WARP): в .env GPU_RENT_SSH_CIDR=0.0.0.0/0 и снова up."
+    )


 def ssh_argv(cfg: Config, host: str, remote: list[str] | None = None) -> list[str]: