ccssh 1.0.0__py3-none-any.whl

ccssh/mux.py

@@ -0,0 +1,301 @@
+"""
+ Channel multiplexer: share one SSH transport between exec/sftp/shell
+ channels and port forwarders.
+
+ SSH has a single transport stream, so exactly one thread may read it.
+ Without a multiplexer, two concurrent readers (e.g. an exec loop and a
+ forwarding pump) steal each other's packets and both stall. The mux owns
+ the only reader thread; everyone else consumes per-channel queues:
+
+ * channel messages carrying our channel id go to that channel's queue;
+ * ``CHANNEL_SUCCESS``/``CHANNEL_FAILURE`` (99/100, which carry NO channel
+   id on the wire) go to the channel with the oldest outstanding channel
+   request (see :meth:`expect_reply` - register *before* sending);
+ * server-initiated ``CHANNEL_OPEN`` and stray 81/82 go to the catch-all
+   queue (port forwarders);
+ * ``DISCONNECT``/``KEXINIT`` are broadcast to every queue (in-flight
+   operations keep their existing "peer rekey/closed" error paths);
+ * ``IGNORE``/``DEBUG`` are skipped, ``GLOBAL_REQUEST`` answered centrally.
+
+ Rekey still needs idle channels (same rule as before): :meth:`busy`
+ reports whether any channel is registered. The reader also emits idle
+ SSH ``IGNORE`` keepalives so quiet forward-only connections survive.
+"""
+from __future__ import annotations
+import queue
+import struct
+import threading
+import time
+
+from ccssh.transport import CCSSHError, CCSSHTimeoutError, enc_ignore
+
+_DEAD = object() # sentinel: connection died, wake every blocked getter
+
+
+class ChannelMux:
+    """Demultiplex one transport into per-channel queues (sync only)."""
+
+    def __init__(self, transport, *, idle_keepalive: float = 30.0):
+        self._t = transport
+        self._idle_keepalive = idle_keepalive
+        self._lock = threading.Lock()
+        self._queues: dict[int, queue.Queue] = {}
+        self._catch: queue.Queue = queue.Queue()
+        self._expect: list[int] = [] # channels awaiting 99/100, oldest first
+        self._stop = threading.Event()
+        self._dead = threading.Event()
+        self._dead_err: Exception | None = None
+        self._paused = threading.Event()
+        self._parked = threading.Event()
+        # Locks: _lock guards maps/queues; request_lock serialises
+        # expect_reply()+send pairs so 99/100 replies cannot be misrouted
+        # when two threads set up channels concurrently.
+        self.request_lock = threading.Lock()
+        self._chan_next = 1
+        self._last_io = time.monotonic()
+        self._th = threading.Thread(target=self._reader, daemon=True)
+
+    # lifecycle
+    def start(self) -> "ChannelMux":
+        self._th.start()
+        return self
+
+    def stop(self) -> None:
+        self._stop.set()
+        self._paused.clear() # let a parked reader observe the stop
+        self._die(CCSSHError("multiplexer stopped"))
+        th = self._th
+        if th.is_alive():
+            th.join(timeout=2.0)
+
+    def pause(self) -> None:
+        """Park the reader outside any transport read (for rekey).
+
+        Rekey swaps the transport's crypto state mid-stream; a concurrent
+        reader would decrypt one half of a packet with the old keys and
+        the other half with the new ones. Blocks (bounded) until parked.
+        No-op when the reader isn't running yet.
+        """
+        self._paused.set()
+        if not self._th.is_alive():
+            # Reader not running (e.g. pause-before-start): just latch the
+            # flag; the first loop iteration will park on it.
+            return
+        try:
+            margin = self.timeout
+        except Exception:
+            margin = 10.0
+        deadline = time.monotonic() + max(2.0, margin + 1.0)
+        while not self._parked.is_set():
+            if self._dead.is_set() or self._stop.is_set():
+                break
+            if time.monotonic() > deadline:
+                raise CCSSHError("timed out parking the mux reader")
+            time.sleep(0.01)
+
+    def resume(self) -> None:
+        self._paused.clear()
+        # The reader may be inside its park-wait: give it a beat to
+        # observe the cleared flag and leave the parked state, otherwise
+        # a stale parked bit would confuse the next pause().
+        deadline = time.monotonic() + 2.0
+        while self._parked.is_set() and time.monotonic() < deadline:
+            time.sleep(0.01)
+        self._parked.clear()
+
+    def busy(self) -> bool:
+        """
+         Any registered channel queues (open or pending channels).
+
+         A forwarder with zero open channels does not count: its pump only
+         consumes mux queues and idle-sends nothing that a rekey could
+         corrupt (rekey pauses the reader; sends are atomic per packet).
+        """
+        with self._lock:
+            return bool(self._queues)
+
+    @property
+    def dead(self) -> bool:
+        return self._dead.is_set()
+
+    # channels
+    def register(self, chan_id: int) -> queue.Queue:
+        """Create (or reuse) the inbound queue for our channel id."""
+        with self._lock:
+            q = self._queues.get(chan_id)
+            if q is None:
+                q = queue.Queue()
+                self._queues[chan_id] = q
+            return q
+
+    def alloc_chan(self) -> int:
+        """
+         Allocate a fresh channel id from the single shared namespace.
+
+         All mux users (exec/shell/sftp/tunnels/forwards) must allocate
+         here: disjoint local counters would eventually collide on the
+         shared queues and misroute packets (channel-confusion).
+        """
+        with self._lock:
+            n = self._chan_next
+            self._chan_next = n + 1
+            return n
+
+    def _queue_for(self, chan_id: int) -> queue.Queue | None:
+        """Queue lookup without creating (for drain loops)."""
+        with self._lock:
+            return self._queues.get(chan_id)
+
+    def unregister(self, chan_id: int) -> None:
+        with self._lock:
+            self._queues.pop(chan_id, None)
+            self._expect = [c for c in self._expect if c != chan_id]
+
+    def expect_reply(self, chan_id: int) -> None:
+        """
+         Record an outstanding channel request (call BEFORE sending it,
+         so a fast reply can never overtake the registration).
+        """
+        with self._lock:
+            self._expect.append(chan_id)
+
+    def catch_all(self) -> queue.Queue:
+        return self._catch
+
+    def send(self, payload: bytes) -> None:
+        self._t.send_packet(payload)
+
+    @property
+    def timeout(self) -> float:
+        try:
+            return float(getattr(self._t, "timeout", 10.0))
+        except Exception:
+            return 10.0
+
+    # consumption
+    def _raise_if_dead(self) -> None:
+        if self._dead.is_set():
+            raise CCSSHError(str(self._dead_err or "connection closed by peer"))
+
+    def recv_chan(self, chan_id: int, timeout: float | None = None) -> bytes:
+        """Next packet for a channel (None timeout blocks until it arrives
+        or the connection dies; a numeric timeout raises
+        CCSSHTimeoutError like a direct transport read)."""
+        with self._lock:
+            q = self._queues.get(chan_id)
+        if q is None:
+            raise CCSSHError(f"channel {chan_id} is not registered with the mux")
+        while True:
+            try:
+                p = q.get(timeout=timeout)
+            except queue.Empty:
+                self._raise_if_dead()
+                raise CCSSHTimeoutError("channel read timeout")
+            if p is _DEAD:
+                # Re-queue for any other getter on this channel, then fail.
+                try:
+                    q.put(_DEAD)
+                except Exception:
+                    pass
+                self._raise_if_dead()
+                raise CCSSHError("connection closed by peer")
+            return p
+
+    # the only transport reader
+    def _die(self, err: Exception) -> None:
+        with self._lock:
+            if self._dead.is_set():
+                return
+            self._dead_err = err
+            self._dead.set()
+            targets = list(self._queues.values()) + [self._catch]
+        for q in targets:
+            try:
+                q.put(_DEAD)
+            except Exception:
+                pass
+
+    def _broadcast(self, p: bytes) -> None:
+        with self._lock:
+            targets = list(self._queues.values()) + [self._catch]
+        for q in targets:
+            try:
+                q.put(p)
+            except Exception:
+                pass
+
+    def _route(self, p: bytes) -> None:
+        # NB: IGNORE/DEBUG/GLOBAL_REQUEST never reach here: the reader uses
+        # transport.recv_packet(), which skips and answers them centrally.
+        if not p:
+            return
+        m = p[0]
+        if m in (1, 20): # DISCONNECT, KEXINIT: everyone must see these
+            self._broadcast(p)
+            if m == 1:
+                self._die(CCSSHError("connection closed by peer"))
+            return
+        if m in (99, 100):  # SUCCESS/FAILURE carry no channel id
+            with self._lock:
+                chan = self._expect.pop(0) if self._expect else None
+                q = self._queues.get(chan) if chan is not None else None
+            if q is not None:
+                try:
+                    q.put(p)
+                except Exception:
+                    pass
+            return
+        if m == 90:
+            # Server-initiated open (forwarded-tcpip): p[1:5] is the type
+            # string length, NOT a channel id - always to catch-all.
+            try:
+                self._catch.put(p)
+            except Exception:
+                pass
+            return
+        if len(p) >= 5 and m in (91, 92, 93, 94, 95, 96, 97, 98):
+            try:
+                (rcpt,) = struct.unpack(">I", p[1:5])
+            except Exception:
+                return
+            with self._lock:
+                q = self._queues.get(rcpt)
+            if q is not None:
+                try:
+                    q.put(p)
+                except Exception:
+                    pass
+            # anything else for an unknown/stale channel is dropped
+            return
+        try:
+            self._catch.put(p)
+        except Exception:
+            pass
+
+    def _reader(self) -> None:
+        while not self._stop.is_set():
+            if self._paused.is_set():
+                self._parked.set()
+                self._paused.wait(0.5)
+                continue
+            self._parked.clear()
+            try:
+                p = self._t.recv_packet(keepalive=False)
+            except CCSSHTimeoutError:
+                if self._idle_keepalive > 0 and not self._stop.is_set() \
+                        and time.monotonic() - self._last_io >= self._idle_keepalive:
+                    try:
+                        self.send(enc_ignore())
+                    except Exception:
+                        self._die(CCSSHError("connection closed by peer"))
+                        return
+                    self._last_io = time.monotonic()
+                continue
+            except Exception as e:
+                self._die(e)
+                return
+            self._last_io = time.monotonic()
+            try:
+                self._route(p)
+            except Exception:
+                pass

ccssh/sftp.py

@@ -0,0 +1,270 @@
+"""
+ SFTP version 3 client (draft-ietf-secsh-filexfer-02): stat, list, get, put,
+ remove, mkdir over a subsystem channel.
+
+ SFTP packets may straddle SSH channel messages, so responses accumulate in
+ a length-prefixed buffer. Each request goes out as a single CHANNEL_DATA
+ message within the server's maxpacket (servers drop oversized ones).
+"""
+from __future__ import annotations
+
+from ccssh import wire as W
+
+# Message types.
+F_INIT, F_VERSION = 1, 2
+F_OPEN, F_CLOSE, F_READ, F_WRITE = 3, 4, 5, 6
+F_OPENDIR, F_READDIR = 11, 12
+F_REMOVE, F_MKDIR = 13, 14
+F_REALPATH, F_STAT = 16, 17
+F_HANDLE, F_DATA, F_NAME, F_ATTRS, F_STATUS = 102, 103, 104, 105, 101
+
+# Status codes.
+S_OK, S_EOF = 0, 1
+S_NO_SUCH_FILE, S_FAILURE = 2, 4
+S_OP_UNSUPPORTED = 8
+
+# OPEN pflags.
+P_READ, P_WRITE, P_APPEND, P_CREAT, P_TRUNC, P_EXCL = 0x01, 0x02, 0x04, 0x08, 0x10, 0x20
+
+# Attribute flags.
+A_SIZE, A_PERM, A_ACMOD = 0x01, 0x04, 0x08
+A_EXTENDED = 0x80000000
+
+CHUNK = 32768
+
+
+class SFTPError(ValueError):
+    pass
+
+
+def _req(t: int, rid: int, body: bytes) -> bytes:
+    return W.enc_u32(1 + 4 + len(body)) + bytes([t]) + W.enc_u32(rid) + body
+
+
+def _dec_attrs_reader(r: W.Reader) -> dict:
+    fl = r.u32()
+    out: dict = {}
+    if fl & A_SIZE:
+        out["size"] = r.u64()
+    if fl & 0x02:
+        out["uid"], out["gid"] = r.u32(), r.u32()
+    if fl & A_PERM:
+        out["permissions"] = r.u32()
+    if fl & A_ACMOD:
+        out["atime"], out["mtime"] = r.u32(), r.u32()
+    if fl & A_EXTENDED:
+        n = r.u32()
+        if n > 64:
+            raise SFTPError("too many extended attrs")
+        for _ in range(n):
+            r.string()
+            r.string()
+    return out
+
+
+def dec_attrs(buf: bytes) -> dict:
+    return _dec_attrs_reader(W.Reader(buf))
+
+
+class SFTPClient:
+    """An open SFTP channel. Use CCSSHClient.open_sftp()."""
+
+    """
+     Upper bound on one SFTP reply (length prefix). Our own requests are
+     small (reads are <= maxpacket); anything bigger is a malicious or
+     broken server - fail instead of buffering it into memory.
+    """
+    MAX_REPLY = 1024 * 1024
+
+    def __init__(self, send_chan_data, recv_packet_fn, server_chan: int, on_close=None, maxpacket: int = 32768):
+        self._send = send_chan_data # fn(server_chan, bytes)
+        self._recv = recv_packet_fn # fn() -> channel-data bytes (one message)
+        self._srv = server_chan
+        self._on_close = on_close
+        """
+         Servers drop CHANNEL_DATA beyond their maxpacket ("rcvd big packet"),
+         so every request goes out as one message within maxpacket.
+        """
+        self._maxpacket = max(1024, maxpacket)
+        self._rid = 1
+        self._buf = b""
+        self._closed = False
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        if self._on_close:
+            self._on_close()
+
+    def _call(self, ftype: int, body: bytes) -> tuple[int, bytes]:
+        rid = self._rid & 0xFFFFFFFF
+        if rid == 0:  # 0 is never used on the wire: skip it explicitly
+            rid = 1
+            self._rid = 2
+        else:
+            self._rid += 1
+        req = _req(ftype, rid, body)
+        if len(req) > self._maxpacket:
+            raise SFTPError("sftp request exceeds server maxpacket")
+        self._send(self._srv, req)
+        while True:
+            data = self._recv()
+            self._buf += data
+            if len(self._buf) > 2 * self.MAX_REPLY + 4:
+                # Pipelined/garbage bytes beyond any single reply: fail
+                # closed instead of buffering a runaway server into OOM.
+                try:
+                    self.close()
+                except Exception:
+                    pass
+                raise SFTPError("sftp: reply buffer exceeded")
+            if len(self._buf) < 4:
+                continue
+            ln = int.from_bytes(self._buf[:4], "big")
+            if ln > self.MAX_REPLY:
+                raise SFTPError(f"sftp: reply too big ({ln} bytes)")
+            if len(self._buf) < 4 + ln:
+                continue
+            pkt, self._buf = self._buf[4:4 + ln], self._buf[4 + ln:]
+            r = W.Reader(pkt)
+            rt, rr = r.byte(), r.u32()
+            if rr != rid:
+                # Desynchronised stream: fail closed, do not continue on a
+                # half-consumed buffer.
+                try:
+                    self.close()
+                except Exception:
+                    pass
+                raise SFTPError(f"sftp: id mismatch {rr}!={rid}")
+            return rt, pkt[5:]
+
+    def _check(self, rtype: int, rest: bytes, want: int) -> bytes:
+        if rtype == F_STATUS:
+            r = W.Reader(rest)
+            code = r.u32()
+            msg = r.string().decode(errors="replace")
+            if code != S_OK:
+                raise SFTPError(f"sftp status {code}: {msg}")
+            return b""
+        if rtype != want:
+            raise SFTPError(f"sftp: expected {want}, got {rtype}")
+        return rest
+
+    def realpath(self, path: str) -> str:
+        t, rest = self._call(F_REALPATH, W.enc_string(path.encode()))
+        r = W.Reader(self._check(t, rest, F_NAME))
+        r.u32() # count, always 1 here
+        return r.string().decode()
+
+    def stat(self, path: str) -> dict:
+        t, rest = self._call(F_STAT, W.enc_string(path.encode()))
+        return dec_attrs(self._check(t, rest, F_ATTRS))
+
+    def listdir(self, path: str, *, max_entries: int = 50000) -> list[str]:
+        t, rest = self._call(F_OPENDIR, W.enc_string(path.encode()))
+        handle = W.Reader(self._check(t, rest, F_HANDLE)).string()
+        out: list[str] = []
+        out_bytes = 0
+        try:
+            while True:
+                t2, r2 = self._call(F_READDIR, W.enc_string(handle))
+                if t2 == F_STATUS:
+                    r = W.Reader(r2)
+                    if r.u32() == S_EOF:
+                        break
+                    raise SFTPError("readdir failed")
+                r = W.Reader(self._check(t2, r2, F_NAME))
+                n = r.u32()
+                if n > 1024:
+                    raise SFTPError("readdir count too large")
+                for _ in range(n):
+                    name = r.string().decode()
+                    r.string() # longname
+                    _dec_attrs_reader(r)
+                    out_bytes += len(name.encode())
+                    if len(out) >= max_entries or out_bytes > 4 * 1024 * 1024:
+                        raise SFTPError("directory listing too large")
+                    out.append(name)
+        finally:
+            try:
+                self._call(F_CLOSE, W.enc_string(handle))
+            except Exception:
+                pass # never mask the original error with a close failure
+        return out
+
+    def _open(self, path: str, pflags: int) -> bytes:
+        t, rest = self._call(F_OPEN, W.enc_string(path.encode()) + W.enc_u32(pflags) + W.enc_u32(0))
+        return W.Reader(self._check(t, rest, F_HANDLE)).string()
+
+    def _close(self, handle: bytes) -> None:
+        t, rest = self._call(F_CLOSE, W.enc_string(handle))
+        if t == F_STATUS:
+            r = W.Reader(rest)
+            if r.u32() != S_OK:
+                raise SFTPError("close failed")
+
+    def get(self, path: str, *, max_bytes: int | None = 1 << 30) -> bytes:
+        """
+         Download a whole file. A server error mid-transfer raises
+         (never silent truncation); max_bytes caps a runaway server
+         (None disables the cap).
+        """
+        h = self._open(path, P_READ)
+        chunks: list[bytes] = []
+        off = 0
+        try:
+            while True:
+                t, rest = self._call(F_READ, W.enc_string(h) + W.enc_u64(off) + W.enc_u32(CHUNK))
+                if t == F_STATUS:
+                    r = W.Reader(rest)
+                    if r.u32() == S_EOF:
+                        break
+                    raise SFTPError(f"read failed: {r.string().decode(errors='replace')}")
+                data = W.Reader(self._check(t, rest, F_DATA)).string()
+                if not data:
+                    break
+                chunks.append(data)
+                off += len(data)
+                if max_bytes is not None and off > max_bytes:
+                    raise SFTPError(f"file exceeds max_bytes={max_bytes}")
+        finally:
+            try:
+                self._close(h)
+            except Exception:
+                pass # never mask the transfer error with a close failure
+        return b"".join(chunks)
+
+    def put(self, path: str, data: bytes, *, mode: int = 0o644) -> None:
+        attrs = W.enc_u32(A_PERM) + W.enc_u32(mode)
+        t, rest = self._call(F_OPEN, W.enc_string(path.encode()) + W.enc_u32(P_WRITE | P_CREAT | P_TRUNC) + attrs)
+        h = W.Reader(self._check(t, rest, F_HANDLE)).string()
+        if not data:
+            # Empty file: OPEN|TRUNC already truncated it; no zero-WRITE.
+            self._close(h)
+            return
+        wchunk = max(1024, self._maxpacket - 128)
+        try:
+            off = 0
+            while off < len(data):
+                chunk = data[off:off + wchunk]
+                t2, r2 = self._call(F_WRITE, W.enc_string(h) + W.enc_u64(off) + W.enc_string(chunk))
+                if t2 == F_STATUS:
+                    r = W.Reader(r2)
+                    if r.u32() != S_OK:
+                        raise SFTPError("write failed")
+                off += len(chunk)
+        finally:
+            try:
+                self._close(h)
+            except Exception:
+                pass # never mask the transfer error with a close failure
+
+    def remove(self, path: str) -> None:
+        t, rest = self._call(F_REMOVE, W.enc_string(path.encode()))
+        self._check(t, rest, F_STATUS)
+
+    def mkdir(self, path: str, *, mode: int = 0o755) -> None:
+        attrs = W.enc_u32(A_PERM) + W.enc_u32(mode)
+        t, rest = self._call(F_MKDIR, W.enc_string(path.encode()) + attrs)
+        self._check(t, rest, F_STATUS)