sshpeek 0.1.0__py3-none-any.whl

sshpeek/pool.py

@@ -0,0 +1,175 @@
+"""SSH connection pool.
+
+One multiplexed connection per host, created lazily and reconnected on
+demand. Authentication, host keys, proxies etc. all come from the normal
+OpenSSH machinery: asyncssh reads ~/.ssh/config, ~/.ssh/known_hosts and
+talks to the ssh-agent (SSH_AUTH_SOCK), so anything you can reach with
+plain `ssh <host>` works here too.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from dataclasses import dataclass, field
+
+import asyncssh
+
+log = logging.getLogger("sshpeek.pool")
+
+
+@dataclass
+class Forward:
+    """A local port forward: 127.0.0.1:<local> -> <host> ssh -> remote_host:remote_port."""
+
+    id: str
+    host: str
+    remote_host: str
+    remote_port: int
+    listener: asyncssh.SSHListener | None = None
+    desired_local: int = 0  # try to rebind the same local port after a reconnect
+    internal: bool = False  # backing an HTTP service; hidden from the tunnels UI
+    declared: bool = False  # comes from sshpeek.yaml; re-ensured periodically
+
+    @property
+    def local_port(self) -> int | None:
+        try:
+            return self.listener.get_port() if self.listener else None
+        except Exception:
+            return None
+
+
+@dataclass
+class HostState:
+    name: str
+    conn: asyncssh.SSHClientConnection
+    sftp: asyncssh.SFTPClient | None = None
+    forwards: dict[str, Forward] = field(default_factory=dict)
+    closed: bool = False
+    watcher: asyncio.Task | None = None
+
+
+class SSHPool:
+    def __init__(self) -> None:
+        self._states: dict[str, HostState] = {}
+        self._locks: dict[str, asyncio.Lock] = {}
+
+    def _lock(self, host: str) -> asyncio.Lock:
+        return self._locks.setdefault(host, asyncio.Lock())
+
+    async def get(self, host: str) -> HostState:
+        """Return a live HostState, connecting or reconnecting if needed."""
+        async with self._lock(host):
+            st = self._states.get(host)
+            if st is not None and not st.closed:
+                return st
+            return await self._connect(host, previous=st)
+
+    async def _connect(self, host: str, previous: HostState | None) -> HostState:
+        log.info("connecting to %s ...", host)
+        conn = await asyncio.wait_for(
+            asyncssh.connect(
+                host,
+                keepalive_interval=30,
+                keepalive_count_max=3,
+            ),
+            timeout=20,
+        )
+        st = HostState(name=host, conn=conn)
+        st.sftp = await conn.start_sftp_client()
+
+        async def _watch() -> None:
+            await conn.wait_closed()
+            st.closed = True
+            log.warning("connection to %s closed", host)
+
+        st.watcher = asyncio.create_task(_watch())
+        self._states[host] = st
+        log.info("connected to %s", host)
+
+        # Re-establish forwards that existed on the previous connection.
+        if previous is not None:
+            for f in list(previous.forwards.values()):
+                try:
+                    await self._open_forward(st, f)
+                    log.info(
+                        "restored forward %s on 127.0.0.1:%s", f.id, f.local_port
+                    )
+                except Exception as e:  # noqa: BLE001 - best effort restore
+                    log.warning("could not restore forward %s: %s", f.id, e)
+        return st
+
+    async def _open_forward(self, st: HostState, f: Forward) -> Forward:
+        want = f.desired_local or 0
+        try:
+            f.listener = await st.conn.forward_local_port(
+                "127.0.0.1", want, f.remote_host, f.remote_port
+            )
+        except OSError:
+            # Preferred local port taken; let the OS pick one.
+            f.listener = await st.conn.forward_local_port(
+                "127.0.0.1", 0, f.remote_host, f.remote_port
+            )
+        f.desired_local = f.listener.get_port()
+        st.forwards[f.id] = f
+        return f
+
+    async def add_forward(
+        self,
+        host: str,
+        remote_port: int,
+        remote_host: str = "localhost",
+        desired_local: int = 0,
+        fid: str | None = None,
+        internal: bool = False,
+        declared: bool = False,
+    ) -> Forward:
+        st = await self.get(host)
+        fid = fid or f"{host}:{remote_host}:{remote_port}"
+        existing = st.forwards.get(fid)
+        if existing is not None and existing.listener is not None:
+            return existing
+        f = Forward(
+            id=fid,
+            host=host,
+            remote_host=remote_host,
+            remote_port=remote_port,
+            desired_local=desired_local,
+            internal=internal,
+            declared=declared,
+        )
+        return await self._open_forward(st, f)
+
+    def peek_state(self, host: str) -> HostState | None:
+        """Current state without connecting."""
+        return self._states.get(host)
+
+    async def remove_forward(self, host: str, fid: str) -> bool:
+        st = self._states.get(host)
+        if st is None:
+            return False
+        f = st.forwards.pop(fid, None)
+        if f is None:
+            return False
+        if f.listener is not None:
+            f.listener.close()
+        return True
+
+    def all_forwards(self) -> list[tuple[Forward, bool]]:
+        """All registered forwards as (forward, up) pairs."""
+        out: list[tuple[Forward, bool]] = []
+        for st in self._states.values():
+            for f in st.forwards.values():
+                out.append((f, not st.closed))
+        return out
+
+    def hosts(self) -> dict[str, bool]:
+        """Hosts we have talked to, mapped to whether the connection is up."""
+        return {name: not st.closed for name, st in self._states.items()}
+
+    async def close(self) -> None:
+        for st in self._states.values():
+            try:
+                st.conn.close()
+            except Exception:  # noqa: BLE001
+                pass

sshpeek/proxy.py

@@ -0,0 +1,199 @@
+"""Reverse proxy for declared HTTP services.
+
+Each service in sshpeek.yaml gets a stable URL:
+
+    http://<service>.<host>.localhost:<sshpeek port>/
+
+Browsers resolve *.localhost to 127.0.0.1 natively, so no /etc/hosts edits.
+Under the hood every service is backed by an (internal, ephemeral) SSH local
+forward from the pool; the proxy just relays to 127.0.0.1:<forwarded port>.
+Proxying to a real local port keeps things simple and gives WebSockets and
+streaming for free -- no hand-rolled HTTP over SSH channels.
+
+Path rewriting is deliberately absent: because every service lives on its own
+(sub)domain at path /, apps that generate absolute paths (Jupyter, Grafana,
+...) work unmodified, and cookies/origins stay isolated per service.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+
+import httpx
+from starlette.background import BackgroundTask
+from starlette.requests import Request
+from starlette.responses import PlainTextResponse, StreamingResponse
+from starlette.websockets import WebSocket
+from websockets.asyncio.client import connect as ws_connect
+
+from .config import Config, HostSpec, HttpService
+from .pool import SSHPool
+
+log = logging.getLogger("sshpeek.proxy")
+
+# Hop-by-hop headers are between us and each peer, never forwarded.
+HOP = {
+    "connection", "keep-alive", "proxy-authenticate", "proxy-authorization",
+    "te", "trailer", "transfer-encoding", "upgrade", "host",
+}
+
+
+def service_fid(host: str, service: str) -> str:
+    return f"http:{host}:{service}"
+
+
+class ProxyRouter:
+    """ASGI wrapper: routes *.localhost hosts to the proxy, everything else
+    to the inner app."""
+
+    def __init__(self, app, cfg: Config, pool: SSHPool) -> None:
+        self.app = app
+        self.cfg = cfg
+        self.pool = pool
+        self.client = httpx.AsyncClient(
+            timeout=httpx.Timeout(None, connect=10), follow_redirects=False
+        )
+
+    # -- routing -----------------------------------------------------------
+
+    def _match(self, scope) -> tuple[HostSpec, HttpService] | None:
+        headers = {k.decode().lower(): v.decode() for k, v in scope.get("headers") or []}
+        hostname = (headers.get("host") or "").split(":")[0].lower()
+        if not hostname.endswith(".localhost"):
+            return None
+        labels = hostname[: -len(".localhost")]
+        if "." not in labels:
+            return None
+        service, halias = labels.split(".", 1)
+        hs = self.cfg.hosts.get(halias)
+        if hs is None:
+            return None
+        for s in hs.http:
+            if s.name == service:
+                return hs, s
+        return None
+
+    async def _local_port(self, hs: HostSpec, s: HttpService) -> int:
+        f = await self.pool.add_forward(
+            hs.name, s.remote_port, s.remote_host,
+            fid=service_fid(hs.name, s.name), internal=True, declared=True,
+        )
+        if f.local_port is None:
+            raise OSError("forward has no bound port")
+        return f.local_port
+
+    async def __call__(self, scope, receive, send):
+        if scope["type"] in ("http", "websocket"):
+            m = self._match(scope)
+            if m is not None:
+                if scope["type"] == "http":
+                    return await self._http(scope, receive, send, *m)
+                return await self._ws(scope, receive, send, *m)
+        await self.app(scope, receive, send)
+
+    # -- http --------------------------------------------------------------
+
+    async def _http(self, scope, receive, send, hs: HostSpec, s: HttpService):
+        request = Request(scope, receive)
+        try:
+            port = await self._local_port(hs, s)
+        except Exception as e:  # noqa: BLE001
+            resp = PlainTextResponse(
+                f"sshpeek: cannot reach {s.name} on {hs.name}: {e}", status_code=502
+            )
+            return await resp(scope, receive, send)
+
+        url = f"http://127.0.0.1:{port}{scope['path']}"
+        if scope.get("query_string"):
+            url += "?" + scope["query_string"].decode()
+        headers = [(k, v) for k, v in request.headers.items() if k.lower() not in HOP]
+        req = self.client.build_request(
+            request.method, url, headers=headers, content=request.stream()
+        )
+        try:
+            r = await self.client.send(req, stream=True)
+        except httpx.HTTPError as e:
+            resp = PlainTextResponse(
+                f"sshpeek: {s.name} on {hs.name} not responding: {e}", status_code=502
+            )
+            return await resp(scope, receive, send)
+
+        resp = StreamingResponse(
+            r.aiter_raw(), status_code=r.status_code, background=BackgroundTask(r.aclose)
+        )
+        # Bypass Starlette's dict headers to preserve duplicates (Set-Cookie)
+        # and pass raw bytes through untouched (Content-Encoding intact).
+        resp.raw_headers = [
+            (k.encode(), v.encode())
+            for k, v in r.headers.multi_items()
+            if k.lower() not in HOP
+        ]
+        await resp(scope, receive, send)
+
+    # -- websocket ---------------------------------------------------------
+
+    async def _ws(self, scope, receive, send, hs: HostSpec, s: HttpService):
+        ws = WebSocket(scope, receive, send)
+        try:
+            port = await self._local_port(hs, s)
+        except Exception as e:  # noqa: BLE001
+            log.warning("ws: cannot reach %s on %s: %s", s.name, hs.name, e)
+            await ws.close(code=1011)
+            return
+
+        url = f"ws://127.0.0.1:{port}{scope['path']}"
+        if scope.get("query_string"):
+            url += "?" + scope["query_string"].decode()
+        fwd_headers = {}
+        for k, v in scope.get("headers") or []:
+            if k.decode().lower() in ("cookie", "authorization"):
+                fwd_headers[k.decode()] = v.decode()
+
+        try:
+            upstream = await ws_connect(
+                url,
+                subprotocols=scope.get("subprotocols") or None,
+                additional_headers=fwd_headers,
+                max_size=None,
+            )
+        except Exception as e:  # noqa: BLE001
+            log.warning("ws upstream connect failed for %s: %s", url, e)
+            await ws.close(code=1011)
+            return
+
+        await ws.accept(subprotocol=upstream.subprotocol)
+
+        async def client_to_upstream():
+            try:
+                while True:
+                    msg = await ws.receive()
+                    if msg["type"] == "websocket.disconnect":
+                        return
+                    if msg.get("text") is not None:
+                        await upstream.send(msg["text"])
+                    elif msg.get("bytes") is not None:
+                        await upstream.send(msg["bytes"])
+            except Exception:  # noqa: BLE001
+                return
+
+        async def upstream_to_client():
+            try:
+                async for m in upstream:
+                    if isinstance(m, str):
+                        await ws.send_text(m)
+                    else:
+                        await ws.send_bytes(m)
+            except Exception:  # noqa: BLE001
+                return
+
+        t1 = asyncio.create_task(client_to_upstream())
+        t2 = asyncio.create_task(upstream_to_client())
+        _, pending = await asyncio.wait({t1, t2}, return_when=asyncio.FIRST_COMPLETED)
+        for p in pending:
+            p.cancel()
+        for closer in (upstream.close(), ws.close()):
+            try:
+                await closer
+            except Exception:  # noqa: BLE001
+                pass

sshpeek/s3.py

@@ -0,0 +1,110 @@
+"""S3 sources.
+
+Buckets declared under `s3:` in sshpeek.yaml show up as browsable sources
+next to the SSH hosts. Object keys are presented as paths: "directories"
+are the usual `/`-delimited common prefixes. Credentials come from the
+normal boto3 chain (env vars, ~/.aws profiles, SSO, instance roles).
+
+boto3 is an optional dependency (`pip install sshpeek[s3]`) and blocking,
+so it is imported lazily and every call runs in a worker thread.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+
+from .config import S3Spec
+
+log = logging.getLogger("sshpeek.s3")
+
+
+class S3Error(Exception):
+    """Normalized error for anything that goes wrong talking to S3."""
+
+
+class S3Source:
+    def __init__(self, spec: S3Spec) -> None:
+        self.spec = spec
+        self.connected = False  # true after the first successful call
+        self._client = None
+
+    # -- blocking helpers (always called via asyncio.to_thread) -------------
+
+    def _client_or_raise(self):
+        if self._client is None:
+            try:
+                import boto3
+            except ImportError as e:
+                raise S3Error(
+                    "boto3 is not installed -- pip install 'sshpeek[s3]'"
+                ) from e
+            session = boto3.session.Session(
+                profile_name=self.spec.profile, region_name=self.spec.region
+            )
+            self._client = session.client("s3", endpoint_url=self.spec.endpoint)
+        return self._client
+
+    def _key(self, path: str) -> str:
+        """Map a browser path like /runs/2024 onto a full object key."""
+        return self.spec.prefix + path.strip("/")
+
+    def _listdir(self, path: str) -> list[dict]:
+        c = self._client_or_raise()
+        key = self._key(path)
+        pfx = key + "/" if key else ""
+        dirs: list[dict] = []
+        files: list[dict] = []
+        for page in c.get_paginator("list_objects_v2").paginate(
+            Bucket=self.spec.bucket, Prefix=pfx, Delimiter="/"
+        ):
+            for p in page.get("CommonPrefixes", []):
+                name = p["Prefix"][len(pfx):].rstrip("/")
+                dirs.append({"name": name, "dir": True, "size": None, "mtime": None})
+            for o in page.get("Contents", []):
+                if o["Key"] == pfx:  # the "directory marker" object itself
+                    continue
+                files.append(
+                    {
+                        "name": o["Key"][len(pfx):],
+                        "dir": False,
+                        "size": o["Size"],
+                        "mtime": o["LastModified"].timestamp(),
+                    }
+                )
+        self.connected = True
+        out = sorted(dirs, key=lambda e: e["name"].lower())
+        out += sorted(files, key=lambda e: e["name"].lower())
+        return out
+
+    def _stat(self, path: str) -> tuple[float, int]:
+        c = self._client_or_raise()
+        head = c.head_object(Bucket=self.spec.bucket, Key=self._key(path))
+        self.connected = True
+        return head["LastModified"].timestamp(), head["ContentLength"]
+
+    def _open(self, path: str):
+        c = self._client_or_raise()
+        obj = c.get_object(Bucket=self.spec.bucket, Key=self._key(path))
+        self.connected = True
+        return obj["ContentLength"], obj["Body"]
+
+    # -- async surface used by the endpoints ---------------------------------
+
+    async def listdir(self, path: str) -> list[dict]:
+        return await self._call(self._listdir, path)
+
+    async def stat(self, path: str) -> tuple[float, int]:
+        return await self._call(self._stat, path)
+
+    async def open(self, path: str):
+        """Return (size, blocking StreamingBody); read it with to_thread."""
+        return await self._call(self._open, path)
+
+    async def _call(self, fn, *args):
+        try:
+            return await asyncio.to_thread(fn, *args)
+        except S3Error:
+            raise
+        except Exception as e:  # noqa: BLE001 - boto raises a small zoo
+            raise S3Error(f"{self.spec.bucket}: {e}") from e