sproxy 0.1.0__py3-none-any.whl

sproxy/pg.py

@@ -0,0 +1,637 @@
+"""A protocol-terminating PostgreSQL gateway for guarded sessions.
+
+PostgreSQL SCRAM authentication proves knowledge of the password *inside the
+client*, so the HTTP trick of substituting a placeholder in passing traffic
+cannot work. Instead, every ``[postgres.NAME]`` profile gets a local listener
+that terminates the wire protocol: the agent's client authenticates to sproxy
+with a per-session token, and sproxy alone opens the real connection -
+TLS-verified and SCRAM-authenticated - to the host pinned in policy. The real
+password exists only in this process; anything that bypasses the gateway is
+holding a token that no actual database accepts.
+
+Nothing here assumes the agent and the gateway share a host: the client hop is
+plain TCP with token auth, so the listener can later move behind a network
+boundary (e.g. a separate pod) unchanged.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import hashlib
+import hmac
+import secrets
+import ssl
+import sys
+import threading
+import time
+from datetime import datetime, timezone
+from functools import partial
+
+from sproxy.audit import AuditLog
+from sproxy.policy import PostgresProfile
+
+PROTOCOL_3_0 = 196608  # protocol version 3.0
+SSL_REQUEST = 80877103
+GSSENC_REQUEST = 80877104
+CANCEL_REQUEST = 80877102
+
+_MAX_STARTUP = 10_000
+_MAX_MESSAGE = 1 << 30  # PostgreSQL's own wire limit
+
+
+class PgError(RuntimeError):
+    """A connection cannot proceed; ``code`` is a PostgreSQL SQLSTATE."""
+
+    def __init__(self, message: str, code: str = "08000"):
+        super().__init__(message)
+        self.code = code
+
+
+# --- wire format ------------------------------------------------------------
+
+def message(kind: bytes, payload: bytes) -> bytes:
+    return kind + (len(payload) + 4).to_bytes(4, "big") + payload
+
+
+async def read_message(reader: asyncio.StreamReader) -> tuple[bytes, bytes]:
+    kind = await reader.readexactly(1)
+    length = int.from_bytes(await reader.readexactly(4), "big")
+    if not 4 <= length <= _MAX_MESSAGE:
+        raise PgError(f"invalid message length {length}", code="08P01")
+    return kind, await reader.readexactly(length - 4)
+
+
+def startup_message(params: dict[str, str]) -> bytes:
+    body = PROTOCOL_3_0.to_bytes(4, "big")
+    for key, value in params.items():
+        body += key.encode("utf-8") + b"\x00" + value.encode("utf-8") + b"\x00"
+    body += b"\x00"
+    return (len(body) + 4).to_bytes(4, "big") + body
+
+
+async def read_startup(reader: asyncio.StreamReader) -> bytes:
+    """Read one length-prefixed startup-phase packet, sans its length word."""
+    length = int.from_bytes(await reader.readexactly(4), "big")
+    if not 8 <= length <= _MAX_STARTUP:
+        raise PgError(f"invalid startup packet length {length}", code="08P01")
+    return await reader.readexactly(length - 4)
+
+
+def parse_startup_params(body: bytes) -> dict[str, str]:
+    params = {}
+    fields = body.split(b"\x00")
+    for key, value in zip(fields[::2], fields[1::2]):
+        if key:
+            params[key.decode("utf-8", "replace")] = value.decode("utf-8", "replace")
+    return params
+
+
+def error_response(text: str, code: str = "08000") -> bytes:
+    fields = b""
+    for kind, value in (("S", "FATAL"), ("V", "FATAL"), ("C", code), ("M", text)):
+        fields += kind.encode("ascii") + value.encode("utf-8") + b"\x00"
+    return message(b"E", fields + b"\x00")
+
+
+def error_text(payload: bytes) -> str:
+    for chunk in payload.split(b"\x00"):
+        if chunk[:1] == b"M":
+            return chunk[1:].decode("utf-8", "replace")
+    return "unknown error"
+
+
+# --- SCRAM-SHA-256 client (RFC 7677) ----------------------------------------
+
+def _saslname(value: str) -> str:
+    return value.replace("=", "=3D").replace(",", "=2C")
+
+
+def _scram_attrs(text: str) -> dict[str, str]:
+    attrs: dict[str, str] = {}
+    for part in text.split(","):
+        key, sep, value = part.partition("=")
+        if sep:
+            attrs.setdefault(key, value)
+    return attrs
+
+
+class ScramClient:
+    """Client side of SCRAM-SHA-256, without channel binding.
+
+    libpq sends an empty SASL username because PostgreSQL takes the role from
+    the startup message; ``username`` and ``nonce`` are parameterised so the
+    RFC 7677 test vectors can drive this implementation directly.
+    """
+
+    def __init__(self, password: str, *, username: str = "", nonce: str | None = None):
+        self._password = password
+        self._nonce = nonce or secrets.token_urlsafe(18)
+        self._first_bare = f"n={_saslname(username)},r={self._nonce}"
+        self._salted: bytes | None = None
+        self._auth_message: bytes | None = None
+
+    def client_first(self) -> str:
+        return "n,," + self._first_bare
+
+    def client_final(self, server_first: str) -> str:
+        attrs = _scram_attrs(server_first)
+        try:
+            combined = attrs["r"]
+            salt = base64.b64decode(attrs["s"])
+            iterations = int(attrs["i"])
+        except (KeyError, ValueError) as exc:
+            raise PgError(f"malformed SCRAM server-first message: {server_first!r}") from exc
+        if not combined.startswith(self._nonce) or len(combined) <= len(self._nonce):
+            raise PgError("SCRAM server nonce does not extend the client nonce")
+        self._salted = hashlib.pbkdf2_hmac(
+            "sha256", self._password.encode("utf-8"), salt, iterations
+        )
+        client_key = hmac.digest(self._salted, b"Client Key", "sha256")
+        without_proof = f"c=biws,r={combined}"  # biws = base64("n,,")
+        self._auth_message = ",".join(
+            (self._first_bare, server_first, without_proof)
+        ).encode("utf-8")
+        signature = hmac.digest(
+            hashlib.sha256(client_key).digest(), self._auth_message, "sha256"
+        )
+        proof = bytes(a ^ b for a, b in zip(client_key, signature))
+        return f"{without_proof},p={base64.b64encode(proof).decode('ascii')}"
+
+    def verify_server_final(self, server_final: str) -> None:
+        attrs = _scram_attrs(server_final)
+        if "e" in attrs:
+            raise PgError(f"SCRAM authentication failed: {attrs['e']}", code="28P01")
+        if self._salted is None or self._auth_message is None or "v" not in attrs:
+            raise PgError("SCRAM exchange out of order")
+        server_key = hmac.digest(self._salted, b"Server Key", "sha256")
+        expected = hmac.digest(server_key, self._auth_message, "sha256")
+        try:
+            got = base64.b64decode(attrs["v"])
+        except ValueError as exc:
+            raise PgError("malformed SCRAM server signature") from exc
+        if not hmac.compare_digest(expected, got):
+            raise PgError("SCRAM server signature mismatch - not the real server?")
+
+
+# --- upstream side ----------------------------------------------------------
+
+async def _open_upstream(
+    profile: PostgresProfile,
+) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
+    try:
+        reader, writer = await asyncio.open_connection(profile.host, profile.port)
+    except OSError as exc:
+        raise PgError(f"cannot reach {profile.host}:{profile.port}: {exc}") from exc
+    if profile.sslmode == "disable":
+        return reader, writer
+    try:
+        writer.write((8).to_bytes(4, "big") + SSL_REQUEST.to_bytes(4, "big"))
+        await writer.drain()
+        answer = await reader.readexactly(1)
+        if answer != b"S":
+            raise PgError(f"{profile.host} refused TLS, which this profile requires")
+        context = ssl.create_default_context(cafile=profile.sslrootcert or None)
+        await writer.start_tls(context, server_hostname=profile.host)
+    except PgError:
+        writer.close()
+        raise
+    except (ssl.SSLError, OSError, asyncio.IncompleteReadError) as exc:
+        writer.close()
+        raise PgError(f"TLS to {profile.host} failed: {exc}") from exc
+    return reader, writer
+
+
+async def _authenticate_upstream(
+    reader: asyncio.StreamReader,
+    writer: asyncio.StreamWriter,
+    *,
+    user: str,
+    password: str,
+) -> None:
+    scram: ScramClient | None = None
+    while True:
+        kind, payload = await read_message(reader)
+        if kind == b"E":
+            raise PgError(
+                f"upstream authentication failed: {error_text(payload)}", code="28P01"
+            )
+        if kind != b"R":
+            raise PgError(f"unexpected message {kind!r} during upstream authentication")
+        code = int.from_bytes(payload[:4], "big")
+        if code == 0:  # AuthenticationOk
+            return
+        if code == 3:  # cleartext
+            writer.write(message(b"p", password.encode("utf-8") + b"\x00"))
+        elif code == 5:  # md5
+            inner = hashlib.md5(password.encode("utf-8") + user.encode("utf-8")).hexdigest()
+            digest = "md5" + hashlib.md5(inner.encode("ascii") + payload[4:8]).hexdigest()
+            writer.write(message(b"p", digest.encode("ascii") + b"\x00"))
+        elif code == 10:  # SASL mechanism negotiation
+            if b"SCRAM-SHA-256" not in payload[4:].split(b"\x00"):
+                raise PgError("upstream offers no supported SASL mechanism (need SCRAM-SHA-256)")
+            scram = ScramClient(password)
+            first = scram.client_first().encode("utf-8")
+            writer.write(
+                message(b"p", b"SCRAM-SHA-256\x00" + len(first).to_bytes(4, "big") + first)
+            )
+        elif code == 11:  # SASL continue
+            if scram is None:
+                raise PgError("SCRAM continue before mechanism negotiation")
+            final = scram.client_final(payload[4:].decode("utf-8"))
+            writer.write(message(b"p", final.encode("utf-8")))
+        elif code == 12:  # SASL final
+            if scram is None:
+                raise PgError("SCRAM final before mechanism negotiation")
+            scram.verify_server_final(payload[4:].decode("utf-8"))
+        else:
+            raise PgError(f"unsupported upstream authentication method {code}")
+        await writer.drain()
+
+
+async def _forward_cancel(profile: PostgresProfile, packet: bytes) -> None:
+    """Relay a CancelRequest; its key is the relayed upstream one, so it fits."""
+    try:
+        _, writer = await _open_upstream(profile)
+    except PgError:
+        return
+    try:
+        writer.write((len(packet) + 4).to_bytes(4, "big") + packet)
+        await writer.drain()
+    except (ConnectionError, OSError):
+        pass
+    finally:
+        writer.close()
+
+
+async def _splice(
+    client_reader: asyncio.StreamReader,
+    client_writer: asyncio.StreamWriter,
+    up_reader: asyncio.StreamReader,
+    up_writer: asyncio.StreamWriter,
+) -> int:
+    """Relay until either side hangs up; returns the number of queries seen."""
+    queries = 0
+
+    async def client_to_upstream() -> None:
+        nonlocal queries
+        while True:
+            kind = await client_reader.read(1)
+            if not kind:
+                break
+            length_bytes = await client_reader.readexactly(4)
+            length = int.from_bytes(length_bytes, "big")
+            if not 4 <= length <= _MAX_MESSAGE:
+                break
+            payload = await client_reader.readexactly(length - 4)
+            if kind in (b"Q", b"P"):  # simple query / extended-protocol Parse
+                queries += 1
+            up_writer.write(kind + length_bytes + payload)
+            await up_writer.drain()
+
+    async def upstream_to_client() -> None:
+        while True:
+            chunk = await up_reader.read(65536)
+            if not chunk:
+                break
+            client_writer.write(chunk)
+            await client_writer.drain()
+
+    tasks = [
+        asyncio.create_task(client_to_upstream()),
+        asyncio.create_task(upstream_to_client()),
+    ]
+    _, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+    for task in pending:
+        task.cancel()
+    await asyncio.gather(*tasks, return_exceptions=True)
+    return queries
+
+
+# --- the gateway ------------------------------------------------------------
+
+class Endpoint:
+    """One profile plus its resolved upstream password."""
+
+    def __init__(self, profile: PostgresProfile, password: str):
+        self.profile = profile
+        self.password = password
+        self.bound_port: int | None = None
+
+
+class PostgresGateway:
+    """Serve every profile on a local port and audit each connection.
+
+    The gateway is a separate process from the mitmdump addon and each audit
+    hash chain tolerates only one writer, so it chains its own file.
+
+    The live profile set can be replaced while the gateway runs (see
+    :meth:`apply`), because a policy file may be edited mid-session. Listeners
+    are keyed by profile name and looked up per connection, so re-pointing a
+    profile at another database changes only what the *next* connection does.
+    """
+
+    def __init__(
+        self,
+        endpoints: list[Endpoint],
+        token: str,
+        *,
+        audit_log: str = "",
+        listen_host: str = "127.0.0.1",
+    ):
+        self._pending = list(endpoints)
+        self._token = token
+        self._listen_host = listen_host
+        self._audit_path = audit_log
+        self._audit: AuditLog | None = None
+        # Both keyed by profile name: the endpoint being served, and the
+        # listener the session has been told about.
+        self._live: dict[str, Endpoint] = {}
+        self._listeners: dict[str, asyncio.Server] = {}
+
+    @property
+    def endpoints(self) -> list[Endpoint]:
+        """The profiles currently being served, in policy order."""
+        return list(self._live.values())
+
+    async def start(self) -> None:
+        if self._audit_path:
+            self._audit = AuditLog(self._audit_path)
+        await self.apply(self._pending)
+
+    async def apply(
+        self, endpoints: list[Endpoint], *, audit_log: str | None = None
+    ) -> dict[str, list[str]]:
+        """Make ``endpoints`` the live set without disturbing issued endpoints.
+
+        A profile that is still configured keeps its local port - the session
+        was already handed that port and the session token, and a policy edit
+        must not invalidate credentials the agent already holds. Established
+        connections are never interrupted; a changed profile takes effect on the
+        next connection. Removed profiles stop accepting new connections.
+
+        New listeners are bound before any old one is closed, so a port
+        collision aborts the change with the previous set still serving.
+
+        Returns the per-kind profile names that changed, for the caller to report.
+        """
+        current = dict(self._live)
+        incoming = {e.profile.name: e for e in endpoints}
+        changes: dict[str, list[str]] = {
+            "added": [], "removed": [], "updated": [], "rebound": []
+        }
+        fresh: dict[str, asyncio.Server] = {}
+        try:
+            for name, endpoint in incoming.items():
+                previous = current.get(name)
+                if previous is not None and previous.bound_port is not None and (
+                    endpoint.profile.listen in (0, previous.bound_port)
+                ):
+                    # Same local endpoint: keep the listener already advertised.
+                    endpoint.bound_port = previous.bound_port
+                    if (
+                        endpoint.profile != previous.profile
+                        or endpoint.password != previous.password
+                    ):
+                        changes["updated"].append(name)
+                    continue
+                fresh[name] = await self._listen(name, endpoint)
+                endpoint.bound_port = fresh[name].sockets[0].getsockname()[1]
+                changes["rebound" if previous is not None else "added"].append(name)
+        except PgError:
+            for server in fresh.values():
+                server.close()
+            raise
+
+        if audit_log is not None and audit_log != self._audit_path:
+            self._audit_path = audit_log
+            self._audit = AuditLog(audit_log) if audit_log else None
+
+        for name, server in list(self._listeners.items()):
+            if name in fresh or name not in incoming:
+                # close() stops accepting; connections in progress are left to
+                # finish, so saving a config file never kills a running query.
+                # Deliberately no wait_closed(): from Python 3.12 it blocks
+                # until every relayed session has ended.
+                server.close()
+                del self._listeners[name]
+                if name not in incoming:
+                    changes["removed"].append(name)
+        self._listeners.update(fresh)
+        self._live = incoming
+        return changes
+
+    async def _listen(self, name: str, endpoint: Endpoint) -> asyncio.Server:
+        """Bind this profile's local listener, preferring its advertised port."""
+        wanted = endpoint.profile.listen or endpoint.bound_port or 0
+        try:
+            return await self._start_server(name, wanted)
+        except PgError:
+            if endpoint.profile.listen or not wanted:
+                raise
+            # ``bound_port`` was only a hint - the port this profile had before
+            # it was removed from the policy. Re-adding it should not fail just
+            # because something else took that port meanwhile.
+            return await self._start_server(name, 0)
+
+    async def _start_server(self, name: str, port: int) -> asyncio.Server:
+        try:
+            return await asyncio.start_server(
+                partial(self._serve, name), self._listen_host, port
+            )
+        except OSError as exc:
+            raise PgError(
+                f"postgres profile {name!r} cannot listen on "
+                f"{self._listen_host}:{port}: {exc}"
+            ) from exc
+
+    async def stop(self) -> None:
+        for server in self._listeners.values():
+            server.close()
+        for server in self._listeners.values():
+            await server.wait_closed()
+        self._listeners.clear()
+
+    def child_env(self) -> dict[str, str]:
+        """Endpoint advertisements for the guarded session.
+
+        Values are libpq conninfo strings, so ``psql "$SPROXY_POSTGRES_NAME"`` just
+        works and extra parameters can be appended to override the defaults.
+        Nothing password-like beyond the session token is ever placed here.
+        """
+        env = {"SPROXY_POSTGRES_TOKEN": self._token}
+        for endpoint in self.endpoints:
+            env[endpoint.profile.env_var] = self._conninfo(endpoint)
+        return env
+
+    def _conninfo(self, endpoint: Endpoint) -> str:
+        profile = endpoint.profile
+        # sslmode=disable covers only the local hop to the gateway; the
+        # upstream connection is TLS-verified by sproxy per the profile.
+        parts = [
+            f"host={self._listen_host}",
+            f"port={endpoint.bound_port}",
+            f"password={self._token}",
+            "sslmode=disable",
+        ]
+        if profile.user:
+            parts.append(f"user={profile.user}")
+        if profile.database:
+            parts.append(f"dbname={profile.database}")
+        return " ".join(parts)
+
+    async def _serve(
+        self,
+        name: str,
+        reader: asyncio.StreamReader,
+        writer: asyncio.StreamWriter,
+    ) -> None:
+        started = time.monotonic()
+        # Look the profile up per connection so a reloaded policy applies here
+        # without rebinding the port the session was told to use.
+        endpoint = self._live.get(name)
+        if endpoint is None:  # removed by a reload as this connection arrived
+            await _fail(writer, PgError(f"postgres profile {name!r} is no longer configured"))
+            writer.close()
+            return
+        profile = endpoint.profile
+        record = {
+            "kind": "postgres",
+            "profile": profile.name,
+            "host": profile.host,
+            "port": profile.port,
+            "tls": profile.sslmode != "disable",
+            "user": "",
+            "database": "",
+            "decision": "error",
+            "reason": "",
+            "queries": 0,
+        }
+        up_writer: asyncio.StreamWriter | None = None
+        try:
+            packet = await read_startup(reader)
+            code = int.from_bytes(packet[:4], "big")
+            while code in (SSL_REQUEST, GSSENC_REQUEST):
+                # The local hop is plain TCP for now; libpq falls back cleanly.
+                writer.write(b"N")
+                await writer.drain()
+                packet = await read_startup(reader)
+                code = int.from_bytes(packet[:4], "big")
+            if code == CANCEL_REQUEST:
+                record["decision"] = "cancel"
+                await _forward_cancel(profile, packet)
+                return
+            if code != PROTOCOL_3_0:
+                raise PgError(f"unsupported protocol version {code}", code="08P01")
+
+            params = parse_startup_params(packet[4:])
+            # The profile's user/database are advertised defaults, not
+            # overrides: the client's explicit choice is forwarded, and the
+            # database itself enforces what this profile's credential may do.
+            user = params.get("user") or profile.user
+            if not user:
+                raise PgError("no user in startup message", code="28000")
+            database = params.get("database") or profile.database or user
+            record["user"] = user
+            record["database"] = database
+
+            writer.write(message(b"R", (3).to_bytes(4, "big")))  # cleartext request
+            await writer.drain()
+            kind, payload = await read_message(reader)
+            supplied = payload.rstrip(b"\x00") if kind == b"p" else b""
+            if not hmac.compare_digest(supplied, self._token.encode("utf-8")):
+                record["decision"] = "denied"
+                record["reason"] = "bad session token"
+                raise PgError(
+                    "invalid session token (connect with the advertised "
+                    f"${profile.env_var} conninfo or $SPROXY_POSTGRES_TOKEN as the password)",
+                    code="28P01",
+                )
+
+            upstream_params = dict(params)
+            upstream_params["user"] = user
+            upstream_params["database"] = database
+            up_reader, up_writer = await _open_upstream(profile)
+            up_writer.write(startup_message(upstream_params))
+            await up_writer.drain()
+            await _authenticate_upstream(
+                up_reader, up_writer, user=user, password=endpoint.password
+            )
+
+            writer.write(message(b"R", (0).to_bytes(4, "big")))  # AuthenticationOk
+            await writer.drain()
+            record["decision"] = "allow"
+            record["queries"] = await _splice(reader, writer, up_reader, up_writer)
+        except PgError as exc:
+            record["reason"] = record["reason"] or str(exc)
+            await _fail(writer, exc)
+        except (asyncio.IncompleteReadError, ConnectionError, OSError) as exc:
+            record["reason"] = record["reason"] or f"connection error: {exc}"
+        finally:
+            record["duration_ms"] = int((time.monotonic() - started) * 1000)
+            self._log(record)
+            for w in (up_writer, writer):
+                if w is not None:
+                    w.close()
+
+    def _log(self, record: dict) -> None:
+        if self._audit is None:
+            return
+        try:
+            self._audit.append({"ts": datetime.now(timezone.utc).isoformat(), **record})
+        except Exception as exc:  # never let audit failure break traffic
+            print(f"sproxy: postgres audit write failed: {exc}", file=sys.stderr, flush=True)
+
+
+async def _fail(writer: asyncio.StreamWriter, exc: PgError) -> None:
+    try:
+        writer.write(error_response(f"sproxy: {exc}", code=exc.code))
+        await writer.drain()
+    except (ConnectionError, OSError):
+        pass
+
+
+class GatewayThread:
+    """Run a :class:`PostgresGateway` on its own event loop beside the runner."""
+
+    def __init__(self, gateway: PostgresGateway):
+        self.gateway = gateway
+        self._loop = asyncio.new_event_loop()
+        self._thread = threading.Thread(
+            target=self._run, name="sproxy-postgres", daemon=True
+        )
+
+    def _run(self) -> None:
+        asyncio.set_event_loop(self._loop)
+        try:
+            self._loop.run_forever()
+        finally:
+            self._loop.close()
+
+    def start(self, timeout: float = 10.0) -> None:
+        self._thread.start()
+        asyncio.run_coroutine_threadsafe(self.gateway.start(), self._loop).result(timeout)
+
+    def apply(
+        self,
+        endpoints: list[Endpoint],
+        *,
+        audit_log: str | None = None,
+        timeout: float = 10.0,
+    ) -> dict[str, list[str]]:
+        """Replace the live profile set from another thread (a policy reload)."""
+        return asyncio.run_coroutine_threadsafe(
+            self.gateway.apply(endpoints, audit_log=audit_log), self._loop
+        ).result(timeout)
+
+    def child_env(self) -> dict[str, str]:
+        return self.gateway.child_env()
+
+    def stop(self, timeout: float = 5.0) -> None:
+        if not self._thread.is_alive():
+            return
+        try:
+            asyncio.run_coroutine_threadsafe(self.gateway.stop(), self._loop).result(timeout)
+        finally:
+            self._loop.call_soon_threadsafe(self._loop.stop)
+            self._thread.join(timeout)