python-plugin 0.1.0__py3-none-any.whl

pyplugin/__init__.py

@@ -0,0 +1,54 @@
+"""pyplugin — wire-compatible Python port of HashiCorp's go-plugin."""
+from __future__ import annotations
+
+from .broker import GRPCBroker
+from .client import Client, ClientConfig
+from .errors import (
+    AppProtocolMismatch,
+    CoreProtocolMismatch,
+    HandshakeError,
+    MagicCookieMismatch,
+    ProcessExitedError,
+    PyPluginError,
+    StartTimeout,
+    TLSError,
+    UnsupportedProtocol,
+)
+from .handshake import (
+    CORE_PROTOCOL_VERSION,
+    HandshakeConfig,
+    HandshakeLine,
+    NETWORK_TCP,
+    NETWORK_UNIX,
+    PROTOCOL_GRPC,
+)
+from .plugin import Plugin, PluginSet, VersionedPlugins
+from .reattach import ReattachConfig
+from .server import ServeConfig, serve
+
+__all__ = [
+    "CORE_PROTOCOL_VERSION",
+    "Client",
+    "ClientConfig",
+    "GRPCBroker",
+    "HandshakeConfig",
+    "HandshakeLine",
+    "NETWORK_TCP",
+    "NETWORK_UNIX",
+    "PROTOCOL_GRPC",
+    "Plugin",
+    "PluginSet",
+    "ReattachConfig",
+    "ServeConfig",
+    "VersionedPlugins",
+    "serve",
+    "PyPluginError",
+    "HandshakeError",
+    "CoreProtocolMismatch",
+    "AppProtocolMismatch",
+    "UnsupportedProtocol",
+    "MagicCookieMismatch",
+    "ProcessExitedError",
+    "StartTimeout",
+    "TLSError",
+]

pyplugin/_generated/grpc_broker_grpc.py

@@ -0,0 +1,40 @@
+# Generated by the Protocol Buffers compiler. DO NOT EDIT!
+# source: grpc_broker.proto
+# plugin: grpclib.plugin.main
+import abc
+import typing
+
+import grpclib.const
+import grpclib.client
+if typing.TYPE_CHECKING:
+    import grpclib.server
+
+from . import grpc_broker_pb2
+
+
+class GRPCBrokerBase(abc.ABC):
+
+    @abc.abstractmethod
+    async def StartStream(self, stream: 'grpclib.server.Stream[grpc_broker_pb2.ConnInfo, grpc_broker_pb2.ConnInfo]') -> None:
+        pass
+
+    def __mapping__(self) -> typing.Dict[str, grpclib.const.Handler]:
+        return {
+            '/plugin.GRPCBroker/StartStream': grpclib.const.Handler(
+                self.StartStream,
+                grpclib.const.Cardinality.STREAM_STREAM,
+                grpc_broker_pb2.ConnInfo,
+                grpc_broker_pb2.ConnInfo,
+            ),
+        }
+
+
+class GRPCBrokerStub:
+
+    def __init__(self, channel: grpclib.client.Channel) -> None:
+        self.StartStream = grpclib.client.StreamStreamMethod(
+            channel,
+            '/plugin.GRPCBroker/StartStream',
+            grpc_broker_pb2.ConnInfo,
+            grpc_broker_pb2.ConnInfo,
+        )

pyplugin/_generated/grpc_broker_pb2.py

@@ -0,0 +1,41 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: grpc_broker.proto
+# Protobuf Python Version: 6.31.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    6,
+    31,
+    1,
+    '',
+    'grpc_broker.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x11grpc_broker.proto\x12\x06plugin\"\x9b\x01\n\x08\x43onnInfo\x12\x12\n\nservice_id\x18\x01 \x01(\r\x12\x0f\n\x07network\x18\x02 \x01(\t\x12\x0f\n\x07\x61\x64\x64ress\x18\x03 \x01(\t\x12%\n\x05knock\x18\x04 \x01(\x0b\x32\x16.plugin.ConnInfo.Knock\x1a\x32\n\x05Knock\x12\r\n\x05knock\x18\x01 \x01(\x08\x12\x0b\n\x03\x61\x63k\x18\x02 \x01(\x08\x12\r\n\x05\x65rror\x18\x03 \x01(\t2C\n\nGRPCBroker\x12\x35\n\x0bStartStream\x12\x10.plugin.ConnInfo\x1a\x10.plugin.ConnInfo(\x01\x30\x01\x42\x30Z.github.com/hashicorp/go-plugin/internal/pluginb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'grpc_broker_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  _globals['DESCRIPTOR']._loaded_options = None
+  _globals['DESCRIPTOR']._serialized_options = b'Z.github.com/hashicorp/go-plugin/internal/plugin'
+  _globals['_CONNINFO']._serialized_start=30
+  _globals['_CONNINFO']._serialized_end=185
+  _globals['_CONNINFO_KNOCK']._serialized_start=135
+  _globals['_CONNINFO_KNOCK']._serialized_end=185
+  _globals['_GRPCBROKER']._serialized_start=187
+  _globals['_GRPCBROKER']._serialized_end=254
+# @@protoc_insertion_point(module_scope)

pyplugin/_generated/grpc_controller_grpc.py

@@ -0,0 +1,40 @@
+# Generated by the Protocol Buffers compiler. DO NOT EDIT!
+# source: grpc_controller.proto
+# plugin: grpclib.plugin.main
+import abc
+import typing
+
+import grpclib.const
+import grpclib.client
+if typing.TYPE_CHECKING:
+    import grpclib.server
+
+from . import grpc_controller_pb2
+
+
+class GRPCControllerBase(abc.ABC):
+
+    @abc.abstractmethod
+    async def Shutdown(self, stream: 'grpclib.server.Stream[grpc_controller_pb2.Empty, grpc_controller_pb2.Empty]') -> None:
+        pass
+
+    def __mapping__(self) -> typing.Dict[str, grpclib.const.Handler]:
+        return {
+            '/plugin.GRPCController/Shutdown': grpclib.const.Handler(
+                self.Shutdown,
+                grpclib.const.Cardinality.UNARY_UNARY,
+                grpc_controller_pb2.Empty,
+                grpc_controller_pb2.Empty,
+            ),
+        }
+
+
+class GRPCControllerStub:
+
+    def __init__(self, channel: grpclib.client.Channel) -> None:
+        self.Shutdown = grpclib.client.UnaryUnaryMethod(
+            channel,
+            '/plugin.GRPCController/Shutdown',
+            grpc_controller_pb2.Empty,
+            grpc_controller_pb2.Empty,
+        )

pyplugin/_generated/grpc_controller_pb2.py

@@ -0,0 +1,39 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: grpc_controller.proto
+# Protobuf Python Version: 6.31.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    6,
+    31,
+    1,
+    '',
+    'grpc_controller.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x15grpc_controller.proto\x12\x06plugin\"\x07\n\x05\x45mpty2:\n\x0eGRPCController\x12(\n\x08Shutdown\x12\r.plugin.Empty\x1a\r.plugin.EmptyB0Z.github.com/hashicorp/go-plugin/internal/pluginb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'grpc_controller_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  _globals['DESCRIPTOR']._loaded_options = None
+  _globals['DESCRIPTOR']._serialized_options = b'Z.github.com/hashicorp/go-plugin/internal/plugin'
+  _globals['_EMPTY']._serialized_start=33
+  _globals['_EMPTY']._serialized_end=40
+  _globals['_GRPCCONTROLLER']._serialized_start=42
+  _globals['_GRPCCONTROLLER']._serialized_end=100
+# @@protoc_insertion_point(module_scope)

pyplugin/_generated/grpc_stdio_grpc.py

@@ -0,0 +1,41 @@
+# Generated by the Protocol Buffers compiler. DO NOT EDIT!
+# source: grpc_stdio.proto
+# plugin: grpclib.plugin.main
+import abc
+import typing
+
+import grpclib.const
+import grpclib.client
+if typing.TYPE_CHECKING:
+    import grpclib.server
+
+import google.protobuf.empty_pb2
+from . import grpc_stdio_pb2
+
+
+class GRPCStdioBase(abc.ABC):
+
+    @abc.abstractmethod
+    async def StreamStdio(self, stream: 'grpclib.server.Stream[google.protobuf.empty_pb2.Empty, grpc_stdio_pb2.StdioData]') -> None:
+        pass
+
+    def __mapping__(self) -> typing.Dict[str, grpclib.const.Handler]:
+        return {
+            '/plugin.GRPCStdio/StreamStdio': grpclib.const.Handler(
+                self.StreamStdio,
+                grpclib.const.Cardinality.UNARY_STREAM,
+                google.protobuf.empty_pb2.Empty,
+                grpc_stdio_pb2.StdioData,
+            ),
+        }
+
+
+class GRPCStdioStub:
+
+    def __init__(self, channel: grpclib.client.Channel) -> None:
+        self.StreamStdio = grpclib.client.UnaryStreamMethod(
+            channel,
+            '/plugin.GRPCStdio/StreamStdio',
+            google.protobuf.empty_pb2.Empty,
+            grpc_stdio_pb2.StdioData,
+        )

pyplugin/_generated/grpc_stdio_pb2.py

@@ -0,0 +1,42 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: grpc_stdio.proto
+# Protobuf Python Version: 6.31.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    6,
+    31,
+    1,
+    '',
+    'grpc_stdio.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x10grpc_stdio.proto\x12\x06plugin\x1a\x1bgoogle/protobuf/empty.proto\"u\n\tStdioData\x12*\n\x07\x63hannel\x18\x01 \x01(\x0e\x32\x19.plugin.StdioData.Channel\x12\x0c\n\x04\x64\x61ta\x18\x02 \x01(\x0c\".\n\x07\x43hannel\x12\x0b\n\x07INVALID\x10\x00\x12\n\n\x06STDOUT\x10\x01\x12\n\n\x06STDERR\x10\x02\x32G\n\tGRPCStdio\x12:\n\x0bStreamStdio\x12\x16.google.protobuf.Empty\x1a\x11.plugin.StdioData0\x01\x42\x30Z.github.com/hashicorp/go-plugin/internal/pluginb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'grpc_stdio_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  _globals['DESCRIPTOR']._loaded_options = None
+  _globals['DESCRIPTOR']._serialized_options = b'Z.github.com/hashicorp/go-plugin/internal/plugin'
+  _globals['_STDIODATA']._serialized_start=57
+  _globals['_STDIODATA']._serialized_end=174
+  _globals['_STDIODATA_CHANNEL']._serialized_start=128
+  _globals['_STDIODATA_CHANNEL']._serialized_end=174
+  _globals['_GRPCSTDIO']._serialized_start=176
+  _globals['_GRPCSTDIO']._serialized_end=247
+# @@protoc_insertion_point(module_scope)

pyplugin/broker.py

@@ -0,0 +1,225 @@
+"""GRPCBroker — bidi sub-channel multiplexer (grpclib async).
+
+Same semantics as go-plugin's non-multiplexed broker: ``accept_and_serve(id, …)``
+opens a fresh listener and sends a ``ConnInfo`` over the broker stream;
+``dial(id)`` waits for that ``ConnInfo`` and opens a channel to the address.
+A demux task pumps inbound stream messages into per-id ``asyncio.Queue``s.
+
+Multiplexing the broker over the main socket (``PLUGIN_MULTIPLEX_GRPC``) is
+not implemented — we always advertise ``false`` if the env var is set.
+
+Sub-channels use mTLS reusing the same cert material as the main channel:
+each side has its leaf cert and key, plus the peer's cert pinned as trust root.
+We hold those PEMs on the broker so we can build the correct SSL context
+(server-side for ``accept_and_serve``, client-side for ``dial``).
+"""
+from __future__ import annotations
+
+import asyncio
+import itertools
+import ssl
+from dataclasses import dataclass
+from typing import Awaitable, Callable, Optional
+
+from grpclib.client import Channel
+from grpclib.config import Configuration
+from grpclib.server import Server, Stream
+
+from . import mtls
+from ._generated import grpc_broker_grpc, grpc_broker_pb2
+from .transport import open_listener
+
+
+_DIAL_TIMEOUT = 5.0
+
+
+@dataclass(frozen=True)
+class TLSMaterial:
+    """PEM bytes used to derive per-direction SSL contexts."""
+    cert_pem: bytes
+    key_pem: bytes
+    peer_cert_pem: bytes
+
+
+class _ServerStreamServicer(grpc_broker_grpc.GRPCBrokerBase):
+    """Plugin-side bridge: pumps a single bidi stream into per-side queues."""
+
+    def __init__(self) -> None:
+        self.incoming: asyncio.Queue[grpc_broker_pb2.ConnInfo | None] = asyncio.Queue()
+        self.outgoing: asyncio.Queue[grpc_broker_pb2.ConnInfo | None] = asyncio.Queue()
+        self.connected = asyncio.Event()
+
+    async def StartStream(self, stream: Stream) -> None:  # noqa: N802
+        self.connected.set()
+
+        async def reader() -> None:
+            try:
+                async for msg in stream:
+                    await self.incoming.put(msg)
+            finally:
+                await self.incoming.put(None)
+
+        async def writer() -> None:
+            while True:
+                msg = await self.outgoing.get()
+                if msg is None:
+                    return
+                await stream.send_message(msg)
+
+        await asyncio.gather(reader(), writer(), return_exceptions=True)
+
+
+class GRPCBroker:
+    """Public broker facade. Lives on both host and plugin sides."""
+
+    def __init__(
+        self,
+        *,
+        send: Callable[[grpc_broker_pb2.ConnInfo], Awaitable[None]],
+        close: Callable[[], None],
+        tls: Optional[TLSMaterial] = None,
+    ) -> None:
+        self._send = send
+        self._close = close
+        self._tls = tls
+        self._lock = asyncio.Lock()
+        self._pending: dict[int, asyncio.Queue[grpc_broker_pb2.ConnInfo]] = {}
+        self._next_id = itertools.count(1)
+        self._closed = False
+        self._servers: list[Server] = []
+
+    def next_id(self) -> int:
+        return next(self._next_id)
+
+    def _q_for(self, sid: int) -> asyncio.Queue[grpc_broker_pb2.ConnInfo]:
+        q = self._pending.get(sid)
+        if q is None:
+            q = asyncio.Queue(maxsize=1)
+            self._pending[sid] = q
+        return q
+
+    async def deliver(self, msg: grpc_broker_pb2.ConnInfo) -> None:
+        """Run-loop callback — push an incoming message onto the per-id queue."""
+        q = self._q_for(msg.service_id)
+        try:
+            q.put_nowait(msg)
+        except asyncio.QueueFull:
+            pass  # drop duplicate Knock/ack — not used for non-mux model
+
+    def _server_ssl(self) -> Optional[ssl.SSLContext]:
+        if self._tls is None:
+            return None
+        return mtls.server_ssl_context(
+            cert_pem=self._tls.cert_pem,
+            key_pem=self._tls.key_pem,
+            peer_cert_pem=self._tls.peer_cert_pem,
+        )
+
+    def _client_ssl(self) -> Optional[ssl.SSLContext]:
+        if self._tls is None:
+            return None
+        return mtls.client_ssl_context(
+            cert_pem=self._tls.cert_pem,
+            key_pem=self._tls.key_pem,
+            peer_cert_pem=self._tls.peer_cert_pem,
+        )
+
+    async def accept_and_serve(self, service_id: int, servicers: list) -> Server:
+        """Open a fresh listener, run a grpclib ``Server`` on it for ``service_id``,
+        and notify the peer via the broker stream."""
+        listener = open_listener()
+        server = Server(servicers)
+        srv_ssl = self._server_ssl()
+        if listener.network == "unix":
+            await server.start(path=listener.address, ssl=srv_ssl)
+        else:
+            host, port = listener.address.split(":")
+            await server.start(host=host, port=int(port), ssl=srv_ssl)
+        self._servers.append(server)
+        await self._send(grpc_broker_pb2.ConnInfo(
+            service_id=service_id, network=listener.network, address=listener.address,
+        ))
+        return server
+
+    async def dial(self, service_id: int, *, timeout: float = _DIAL_TIMEOUT) -> Channel:
+        """Wait for the peer's ``ConnInfo`` for ``service_id`` and open a channel."""
+        q = self._q_for(service_id)
+        try:
+            ci = await asyncio.wait_for(q.get(), timeout=timeout)
+        except asyncio.TimeoutError as e:
+            raise TimeoutError(f"timeout waiting for broker conn info for id {service_id}") from e
+
+        cli_ssl = self._client_ssl()
+        cfg = Configuration(ssl_target_name_override="localhost") if cli_ssl else None
+        if ci.network == "unix":
+            return Channel(path=ci.address, ssl=cli_ssl, config=cfg)
+        host, port = ci.address.split(":")
+        return Channel(host=host, port=int(port), ssl=cli_ssl, config=cfg)
+
+    async def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        for srv in self._servers:
+            srv.close()
+            try:
+                await srv.wait_closed()
+            except Exception:  # noqa: BLE001
+                pass
+        self._close()
+
+
+def make_server_side_broker(tls: TLSMaterial | None) -> tuple[_ServerStreamServicer, GRPCBroker, "asyncio.Task[None]"]:
+    """Build (servicer, broker, demux-task) for the plugin side."""
+    servicer = _ServerStreamServicer()
+
+    async def send(msg: grpc_broker_pb2.ConnInfo) -> None:
+        await servicer.outgoing.put(msg)
+
+    def close() -> None:
+        servicer.outgoing.put_nowait(None)
+
+    broker = GRPCBroker(send=send, close=close, tls=tls)
+
+    async def demux() -> None:
+        while True:
+            msg = await servicer.incoming.get()
+            if msg is None:
+                return
+            await broker.deliver(msg)
+
+    return servicer, broker, asyncio.ensure_future(demux())
+
+
+def make_client_side_broker(channel: Channel, tls: TLSMaterial | None) -> tuple[GRPCBroker, "asyncio.Task[None]"]:
+    """Build (broker, run-task) for the host side, dialing StartStream."""
+    stub = grpc_broker_grpc.GRPCBrokerStub(channel)
+    outgoing: asyncio.Queue[grpc_broker_pb2.ConnInfo | None] = asyncio.Queue()
+    stream_holder: dict = {}
+
+    async def send(msg: grpc_broker_pb2.ConnInfo) -> None:
+        await outgoing.put(msg)
+
+    def close() -> None:
+        outgoing.put_nowait(None)
+
+    broker = GRPCBroker(send=send, close=close, tls=tls)
+
+    async def runner() -> None:
+        async with stub.StartStream.open() as stream:
+            stream_holder["stream"] = stream
+
+            async def writer() -> None:
+                while True:
+                    msg = await outgoing.get()
+                    if msg is None:
+                        return
+                    await stream.send_message(msg)
+
+            async def reader() -> None:
+                async for msg in stream:
+                    await broker.deliver(msg)
+
+            await asyncio.gather(writer(), reader(), return_exceptions=True)
+
+    return broker, asyncio.ensure_future(runner())