deleted keys fix

Author: arcbtc

__init__.py

@@ -5,7 +5,7 @@
 from loguru import logger
 
 from .crud import db
-from .tasks import wait_for_paid_invoices
+from .tasks import cleanup_tunnels_hourly, wait_for_paid_invoices
 from .views import reverse_proxy_generic_router
 from .views_api import reverse_proxy_api_router
 
@@ -37,6 +37,10 @@ def reverse_proxy_stop():
 def reverse_proxy_start():
     task = create_permanent_unique_task("ext_reverse_proxy", wait_for_paid_invoices)
     scheduled_tasks.append(task)
+    hourly_cleanup = create_permanent_unique_task(
+        "ext_reverse_proxy_cleanup", cleanup_tunnels_hourly
+    )
+    scheduled_tasks.append(hourly_cleanup)
 
 
 __all__ = [
@@ -45,4 +49,4 @@ def reverse_proxy_start():
     "reverse_proxy_start",
     "reverse_proxy_static_files",
     "reverse_proxy_stop",
-]
+]

crud.py

@@ -56,6 +56,13 @@ async def update_extension_settings(user_id: str, data: ExtensionSettings) -> Ex
     return settings
 
 
+async def get_all_extension_settings() -> list[UserExtensionSettings]:
+    return await db.fetchall(
+        "SELECT * FROM reverse_proxy.extension_settings",
+        model=UserExtensionSettings,
+    )
+
+
 ########################### Proxies ############################
 async def create_proxies(
     user_id: str,

migrations.py

@@ -67,4 +67,17 @@ async def m004_extension_settings(db):
             updated_at TIMESTAMP NOT NULL DEFAULT {db.timestamp_now}
         );
     """
-    )
+    )
+
+
+async def m005_add_ssh_public_key(db):
+    """
+    Store the client's SSH public key so it can be re-authorized after grace period.
+    """
+
+    await db.execute(
+        """
+        ALTER TABLE reverse_proxy.proxies
+        ADD COLUMN ssh_public_key TEXT;
+    """
+    )

models.py

@@ -31,6 +31,7 @@ class CreateProxies(BaseModel):
     subdomain: str
     remote_port: int
     ssh_user: str
+    ssh_public_key: str | None = None
     expires_at: datetime
     status: str
     public_id: str
@@ -44,6 +45,7 @@ class Proxies(BaseModel):
     subdomain: str
     remote_port: int
     ssh_user: str
+    ssh_public_key: str | None = None
     expires_at: datetime
     status: str
     public_id: str

services.py

@@ -278,6 +278,10 @@ async def cleanup_expired_proxies_and_sync_caddy(
         model=Proxies,
     )
     for proxy in expired:
+        if not proxy.ssh_public_key:
+            captured_key = _extract_authorized_key(settings, proxy)
+            if captured_key:
+                proxy.ssh_public_key = captured_key
         proxy.status = "expired"
         await update_proxies(proxy)
         _remove_authorized_key(settings, proxy)
@@ -309,6 +313,15 @@ async def cleanup_expired_proxies_and_sync_caddy(
         model=Proxies,
     )
 
+    # 2b) Ensure authorized_keys contains entries for active tunnels
+    for proxy in active:
+        if not proxy.ssh_public_key:
+            proxy.ssh_public_key = _extract_authorized_key(settings, proxy)
+            if proxy.ssh_public_key:
+                await update_proxies(proxy)
+        if proxy.ssh_public_key:
+            _add_authorized_key(settings, proxy, proxy.ssh_public_key)
+
     # 3) Rewrite Caddy tunnels file
     try:
         os.makedirs(os.path.dirname(settings.caddy_tunnels_file), exist_ok=True)
@@ -437,6 +450,9 @@ async def payment_received_for_payments(payment: Payment) -> bool:
         logger.warning("reverse_proxy: invalid price_per_day_sats in settings")
         return False
 
+    if not proxy.ssh_public_key:
+        proxy.ssh_public_key = _extract_authorized_key(settings, proxy)
+
     days_to_extend = (payment.amount / 1000) / price_per_day
 
     now = datetime.now(timezone.utc)
@@ -447,6 +463,9 @@ async def payment_received_for_payments(payment: Payment) -> bool:
     proxy.status = "active"
     await update_proxies(proxy)
 
+    if proxy.ssh_public_key:
+        _add_authorized_key(settings, proxy, proxy.ssh_public_key)
+
     payments.paid = True
     await update_payments(payments)
 
@@ -495,13 +514,16 @@ async def _process_payment_with_setup(payment: Payment, setup: dict) -> bool:
                 subdomain=setup.get("subdomain", random_subdomain()),
                 remote_port=remote_port,
                 ssh_user=setup.get("ssh_user", settings.ssh_user),
+                ssh_public_key=pubkey,
                 expires_at=expires_at,
                 status="pending",
                 public_id=public_id,
                 client_note=setup.get("client_note", ""),
             ),
             proxies_id=proxy_id,
         )
+    if pubkey:
+        proxy.ssh_public_key = pubkey
     if getattr(proxy, "status", "") != "active":
         proxy.status = "active"
 
@@ -544,6 +566,31 @@ def _authorized_keys_path(settings: ExtensionSettings) -> str:
     return os.path.join("/home", settings.ssh_user, ".ssh", "authorized_keys")
 
 
+def _extract_authorized_key(settings: ExtensionSettings, proxy: Proxies) -> str | None:
+    """
+    Try to read the stored public key for this proxy from authorized_keys.
+    """
+    ak_path = _authorized_keys_path(settings)
+    marker = f"reverse_proxy:{proxy.id}"
+    if not os.path.exists(ak_path):
+        return None
+
+    try:
+        with open(ak_path, "r", encoding="utf-8") as f:
+            for line in f:
+                if marker not in line:
+                    continue
+                parts = line.strip().split()
+                if len(parts) < 3 or parts[-1] != marker:
+                    continue
+                return " ".join(parts[1:-1])
+    except Exception as exc:
+        logger.error(
+            f"reverse_proxy: failed to read authorized_key for {proxy.id}: {exc}"
+        )
+    return None
+
+
 def _add_authorized_key(settings: ExtensionSettings, proxy: Proxies, pubkey: str) -> None:
     """
     Append a restricted SSH key for this proxy.

tasks.py

@@ -1,10 +1,14 @@
 import asyncio
 
 from lnbits.core.models import Payment
-from lnbits.tasks import register_invoice_listener
+from lnbits.tasks import register_invoice_listener, run_interval
 from loguru import logger
 
-from .services import payment_received_for_payments
+from .crud import get_all_extension_settings
+from .services import (
+    cleanup_expired_proxies_and_sync_caddy,
+    payment_received_for_payments,
+)
 
 
 async def wait_for_paid_invoices():
@@ -26,3 +30,18 @@ async def on_invoice_paid(payment: Payment) -> None:
         await payment_received_for_payments(payment)
     except Exception as e:
         logger.error(f"Error processing payment for reverse_proxy: {e}")
+
+
+async def _cleanup_all_tunnels_once():
+    settings_list = await get_all_extension_settings()
+    for settings in settings_list:
+        try:
+            await cleanup_expired_proxies_and_sync_caddy(settings.id, settings)
+        except Exception as exc:
+            logger.error(
+                f"reverse_proxy: cleanup failed for user {settings.id}: {exc}"
+            )
+
+
+# Runs every hour while the app is running
+cleanup_tunnels_hourly = run_interval(3600, _cleanup_all_tunnels_once)