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
@@ -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} уже в проекте")