feat: Added reverse SSH tunnel and more

Author: Karol-G

example.py

@@ -3,8 +3,8 @@
 
 
 def main():
-    img = (np.random.rand(512, 512) * 4).astype("uint8")
-    send([img, img])
+    img = (np.random.rand(128, 128))
+    send(img)
 
 
 if __name__ == "__main__":

pyproject.toml

@@ -31,6 +31,7 @@ requires-python = ">=3.10"
 dependencies = [
     "numpy",
     "pyzmq",
+    "medvol",
 ]
 
 [project.optional-dependencies]
@@ -53,6 +54,9 @@ dev = [
 [project.entry-points."napari.manifest"]
 napari-stream = "napari_stream:napari.yaml"
 
+[project.scripts]
+napari_stream = "napari_stream.sender:main"
+
 [build-system]
 requires = ["setuptools>=42.0.0", "setuptools_scm"]
 build-backend = "setuptools.build_meta"

src/napari_stream/_receiver_widget.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 import traceback
-from typing import Optional
+import shutil
+import subprocess
+from typing import Optional, Tuple
 import numpy as np
 from qtpy.QtCore import QThread
 from qtpy.QtGui import QGuiApplication
@@ -12,9 +14,12 @@
     QLabel,
     QHBoxLayout,
     QCheckBox,
+    QRadioButton,
+    QButtonGroup,
 )
 
-from ._listener import ZMQImageListener, bind_endpoint_for_public, default_endpoint
+from ._listener import ZMQImageListener, bind_endpoint_for_public
+from ._utils import DEFAULT_TCP_PORT, default_endpoint
 
 try:
     from napari.types import ImageData
@@ -32,32 +37,52 @@ def __init__(self, napari_viewer: Viewer):
         self._thread: Optional[QThread] = None
         self._worker: Optional[ZMQImageListener] = None
 
-        self._last_auto_endpoint = default_endpoint()
-        self.endpoint_edit = QLineEdit(self._last_auto_endpoint)
+        # Mode-specific endpoints
+        self._endpoint_local = f"tcp://127.0.0.1:{DEFAULT_TCP_PORT}"
+        self._endpoint_private = default_endpoint(public=True)
+        self._endpoint_tunnel = ""
+
+        self._tunnel_proc: Optional[subprocess.Popen] = None
+        self._last_mode = "local"
+
+        self.endpoint_edit = QLineEdit(self._endpoint_local)
         self.status_label = QLabel("Idle")
-        self.public_access = QCheckBox("Enable public Access")
+        self.mode_local = QRadioButton("Local only")
+        self.mode_private = QRadioButton("Private network")
+        self.mode_tunnel = QRadioButton("Private network (Reverse SSH Tunnel)")
+        self.mode_local.setChecked(True)
+        self.mode_group = QButtonGroup(self)
+        for btn in (self.mode_local, self.mode_private, self.mode_tunnel):
+            self.mode_group.addButton(btn)
+
         self.autocontrast = QCheckBox("Auto-contrast on new images")
         self.autocontrast.setChecked(True)
+        self.ignore_affine = QCheckBox("Ignore affine")
 
         self.btn_run = QPushButton("Start")
         self.btn_copy = QPushButton("Copy Endpoint")
 
         top = QVBoxLayout(self)
-        row = QHBoxLayout()
-        row.addWidget(QLabel("Endpoint:"))
-        row.addWidget(self.endpoint_edit)
-        row.addWidget(self.btn_copy)
-        top.addLayout(row)
-        top.addWidget(self.public_access)
+        top.addWidget(QLabel("Endpoint:"))
+        top.addWidget(self.endpoint_edit)
+        top.addWidget(self.btn_copy)
+        top.addWidget(self.mode_local)
+        top.addWidget(self.mode_private)
+        top.addWidget(self.mode_tunnel)
         top.addWidget(self.autocontrast)
+        top.addWidget(self.ignore_affine)
         top.addWidget(self.status_label)
         row2 = QHBoxLayout()
         row2.addWidget(self.btn_run)
         top.addLayout(row2)
 
         self.btn_run.clicked.connect(self._on_toggle_clicked)
         self.btn_copy.clicked.connect(self._copy_endpoint)
-        self.public_access.stateChanged.connect(self._on_public_toggled)
+        self.mode_group.buttonClicked.connect(self._on_mode_changed)
+        self.destroyed.connect(lambda *_: self._stop_tunnel())
+        app = QGuiApplication.instance()
+        if app is not None:
+            app.aboutToQuit.connect(lambda *_: self._stop_tunnel())
 
     def _on_toggle_clicked(self):
         if self._is_running():
@@ -68,10 +93,30 @@ def _on_toggle_clicked(self):
     def _on_start(self):
         if self._is_running():
             return
-        endpoint = self.endpoint_edit.text().strip()
-        bind_endpoint = self._resolve_endpoint_for_worker(endpoint)
+        mode = self._current_mode()
+        user_entry = self.endpoint_edit.text().strip()
+        endpoint = user_entry or self._endpoint_for_mode(mode)
+        if mode == "local":
+            endpoint = self._ensure_local_endpoint(endpoint, mode="local")
+        elif mode == "private":
+            endpoint = self._ensure_private_endpoint(endpoint)
+        elif mode == "tunnel":
+            ssh_target, ssh_port = self._parse_ssh_target(user_entry)
+            if not ssh_target:
+                self.status_label.setText("Enter SSH target for reverse tunnel.")
+                return
+            endpoint = self._ensure_local_endpoint(
+                self._endpoint_tunnel or f"tcp://127.0.0.1:{DEFAULT_TCP_PORT}",
+                mode="tunnel",
+                update_field=False,
+            )
+            self._store_endpoint_for_mode("tunnel", user_entry)
+        else:
+            self.status_label.setText("Unknown mode.")
+            return
+
         self._thread = QThread()
-        self._worker = ZMQImageListener(bind_endpoint)
+        self._worker = ZMQImageListener(endpoint)
         self._worker.moveToThread(self._thread)
 
         self._thread.started.connect(self._worker.start)
@@ -83,7 +128,12 @@ def _on_start(self):
         self.btn_run.setEnabled(True)
         self._thread.start()
 
+        if mode == "tunnel":
+            if not self._start_reverse_tunnel(ssh_target, ssh_port, endpoint):
+                self._on_stop()
+
     def _on_stop(self):
+        self._stop_tunnel()
         if self._worker is not None:
             self._worker.stop()
         if self._thread is not None:
@@ -94,30 +144,132 @@ def _on_stop(self):
         self.btn_run.setText("Start")
         self.btn_run.setEnabled(True)
 
-    def _on_public_toggled(self, checked: int):
-        is_public = bool(checked)
-        suggested = default_endpoint(public=is_public)
-        current = self.endpoint_edit.text().strip()
-        if current == self._last_auto_endpoint:
-            self.endpoint_edit.setText(suggested)
-        self._last_auto_endpoint = suggested
-        if self._is_running():
-            self._restart_listener()
-
     def _copy_endpoint(self):
         endpoint = self.endpoint_edit.text().strip()
         QGuiApplication.clipboard().setText(endpoint)
 
-    def _resolve_endpoint_for_worker(self, endpoint: str) -> str:
-        if not self.public_access.isChecked():
-            return endpoint
-        if endpoint.startswith("tcp://"):
-            return bind_endpoint_for_public(endpoint)
-        # If the field holds a non-TCP endpoint, fall back to a sensible TCP default.
-        fallback = default_endpoint(public=True)
-        self.endpoint_edit.setText(fallback)
-        self._last_auto_endpoint = fallback
-        return bind_endpoint_for_public(fallback)
+    def _on_mode_changed(self, *_):
+        # Stop if running to allow reconfiguration
+        if self._is_running():
+            self._on_stop()
+        # Persist current text to the previous mode slot
+        prev_mode = getattr(self, "_last_mode", None)
+        if prev_mode:
+            self._store_endpoint_for_mode(prev_mode, self.endpoint_edit.text().strip())
+        mode = self._current_mode()
+        self.endpoint_edit.setPlaceholderText("user@remote or ssh-alias[#port]" if mode == "tunnel" else "")
+        self.endpoint_edit.setText(self._endpoint_for_mode(mode))
+        self._last_mode = mode
+        self.status_label.setText("Idle")
+
+    def _current_mode(self) -> str:
+        if self.mode_private.isChecked():
+            return "private"
+        if self.mode_tunnel.isChecked():
+            return "tunnel"
+        return "local"
+
+    def _endpoint_for_mode(self, mode: str) -> str:
+        if mode == "private":
+            return self._endpoint_private
+        if mode == "tunnel":
+            # Show blank to prompt SSH target entry; fall back to stored value if present
+            return "" if not self._endpoint_tunnel else self._endpoint_tunnel
+        return self._endpoint_local
+
+    def _store_endpoint_for_mode(self, mode: str, endpoint: str) -> None:
+        if mode == "private":
+            self._endpoint_private = endpoint or self._endpoint_private
+        elif mode == "tunnel":
+            self._endpoint_tunnel = endpoint or self._endpoint_tunnel
+        else:
+            self._endpoint_local = endpoint or self._endpoint_local
+
+    def _ensure_local_endpoint(self, endpoint: str, *, mode: str = "local", update_field: bool = True) -> str:
+        port = self._extract_port(endpoint) or DEFAULT_TCP_PORT
+        value = f"tcp://127.0.0.1:{port}"
+        if update_field:
+            self.endpoint_edit.setText(value)
+            self._store_endpoint_for_mode(mode, value)
+        return value
+
+    def _ensure_private_endpoint(self, endpoint: str) -> str:
+        # Prefer a shareable IP, then bind on all interfaces.
+        if not endpoint.startswith("tcp://"):
+            endpoint = self._endpoint_private or default_endpoint(public=True)
+        if endpoint.count(":") < 2:
+            # Missing port: append default
+            endpoint = endpoint.rstrip("/") + f":{DEFAULT_TCP_PORT}"
+        if not endpoint.startswith("tcp://"):
+            endpoint = "tcp://" + endpoint
+        self._endpoint_private = endpoint
+        self.endpoint_edit.setText(endpoint)
+        return bind_endpoint_for_public(endpoint)
+
+    def _extract_port(self, endpoint: str) -> Optional[int]:
+        if not endpoint:
+            return None
+        try:
+            _, port_str = endpoint.rsplit(":", 1)
+            return int(port_str)
+        except Exception:
+            return None
+
+    def _parse_ssh_target(self, raw: str) -> Tuple[str, int]:
+        raw = (raw or "").strip()
+        if not raw:
+            return "", 22
+        if "#" in raw:
+            target, port_part = raw.rsplit("#", 1)
+            try:
+                port = int(port_part)
+            except Exception:
+                port = 22
+            return target or "", port
+        return raw, 22
+
+    def _start_reverse_tunnel(self, target: str, ssh_port: int, endpoint: str) -> bool:
+        self._stop_tunnel()
+        cmd = shutil.which("autossh") or shutil.which("ssh")
+        if cmd is None:
+            self.status_label.setText("autossh/ssh not found in PATH.")
+            return False
+        local_port = self._extract_port(endpoint) or DEFAULT_TCP_PORT
+        remote_port = local_port
+        args = [cmd]
+        if cmd.endswith("autossh"):
+            args += ["-M", "0"]
+        args += ["-N", "-R", f"{remote_port}:localhost:{local_port}"]
+        if ssh_port != 22:
+            args += ["-p", str(ssh_port)]
+        args.append(target)
+        try:
+            self._tunnel_proc = subprocess.Popen(
+                args,
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+                start_new_session=False,
+            )
+            self.status_label.setText(f"Reverse tunnel via {target} (port {ssh_port})")
+            return True
+        except Exception as exc:  # noqa: BLE001
+            self._tunnel_proc = None
+            self.status_label.setText(f"Failed to start tunnel: {exc}")
+            return False
+
+    def _stop_tunnel(self, *_):
+        if self._tunnel_proc is None:
+            return
+        try:
+            self._tunnel_proc.terminate()
+            self._tunnel_proc.wait(timeout=2)
+        except Exception:
+            try:
+                self._tunnel_proc.kill()
+            except Exception:
+                pass
+        finally:
+            self._tunnel_proc = None
 
     def _is_running(self) -> bool:
         return self._thread is not None and self._thread.isRunning()
@@ -134,7 +286,7 @@ def _on_received(self, arr: np.ndarray, meta: dict):
         viewer_kwargs = {}
 
         # Affine: accept any square >= 2x2 (2x2, 3x3, 4x4, ...)
-        if "affine" in meta:
+        if "affine" in meta and not self.ignore_affine.isChecked():
             try:
                 A = np.asarray(meta["affine"], dtype=float)
                 if A.ndim == 2 and A.shape[0] == A.shape[1] and A.shape[0] >= 2:

src/napari_stream/_utils.py

@@ -1,5 +1,4 @@
 from __future__ import annotations
-import os
 import socket
 
 DEFAULT_TCP_PORT = 5556
@@ -8,10 +7,8 @@
 def default_endpoint(public: bool = False) -> str:
     if public:
         return f"tcp://{_preferred_ip()}:{DEFAULT_TCP_PORT}"
-    if os.name == "nt":  # Windows: prefer TCP
-        return f"tcp://127.0.0.1:{DEFAULT_TCP_PORT}"
-    # Unix: fast local IPC
-    return "ipc:///tmp/napari_stream.sock"
+    # TCP everywhere for simplicity
+    return f"tcp://127.0.0.1:{DEFAULT_TCP_PORT}"
 
 
 def _preferred_ip() -> str:
@@ -23,4 +20,4 @@ def _preferred_ip() -> str:
                 return addr
     except Exception:
         pass
-    return "127.0.0.1"
+    return "127.0.0.1"