Added gRPC gnmi protocol to UTCP

plugins/communication_protocols/gnmi/pyproject.toml

@@ -0,0 +1,32 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "utcp-gnmi"
+version = "1.0.0"
+authors = [
+  { name = "UTCP Contributors" },
+]
+description = "UTCP gNMI communication protocol plugin over gRPC"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "pydantic>=2.0",
+    "protobuf>=4.21",
+    "grpcio>=1.60",
+    "utcp>=1.0",
+    "aiohttp>=3.8"
+]
+license = "MPL-2.0"
+
+[project.optional-dependencies]
+dev = [
+    "build",
+    "pytest",
+    "pytest-asyncio",
+    "pytest-cov",
+]
+
+[project.entry-points."utcp.plugins"]
+gnmi = "utcp_gnmi:register"

plugins/communication_protocols/gnmi/src/utcp_gnmi/__init__.py

@@ -0,0 +1,12 @@
+from utcp.plugins.discovery import register_communication_protocol, register_call_template
+from utcp_gnmi.gnmi_communication_protocol import GnmiCommunicationProtocol
+from utcp_gnmi.gnmi_call_template import GnmiCallTemplateSerializer
+
+def register():
+    register_communication_protocol("gnmi", GnmiCommunicationProtocol())
+    register_call_template("gnmi", GnmiCallTemplateSerializer())
+
+__all__ = [
+    "GnmiCommunicationProtocol",
+    "GnmiCallTemplateSerializer",
+]

plugins/communication_protocols/gnmi/src/utcp_gnmi/gnmi_call_template.py

@@ -0,0 +1,27 @@
+from typing import Optional, Dict, List, Literal
+from pydantic import Field
+
+from utcp.data.call_template import CallTemplate
+from utcp.interfaces.serializer import Serializer
+from utcp.exceptions import UtcpSerializerValidationError
+import traceback
+
+class GnmiCallTemplate(CallTemplate):
+    call_template_type: Literal["gnmi"] = "gnmi"
+    target: str
+    use_tls: bool = True
+    metadata: Optional[Dict[str, str]] = None
+    metadata_fields: Optional[List[str]] = None
+    operation: Literal["capabilities", "get", "set", "subscribe"] = "get"
+    stub_module: str = "gnmi_pb2_grpc"
+    message_module: str = "gnmi_pb2"
+
+class GnmiCallTemplateSerializer(Serializer[GnmiCallTemplate]):
+    def to_dict(self, obj: GnmiCallTemplate) -> dict:
+        return obj.model_dump()
+
+    def validate_dict(self, obj: dict) -> GnmiCallTemplate:
+        try:
+            return GnmiCallTemplate.model_validate(obj)
+        except Exception as e:
+            raise UtcpSerializerValidationError("Invalid GnmiCallTemplate: " + traceback.format_exc()) from e

plugins/communication_protocols/gnmi/src/utcp_gnmi/gnmi_communication_protocol.py

@@ -0,0 +1,245 @@
+import importlib
+from typing import Dict, Any, List, Optional, AsyncGenerator
+
+from utcp.interfaces.communication_protocol import CommunicationProtocol
+from utcp.data.call_template import CallTemplate
+from utcp.data.tool import Tool, JsonSchema
+from utcp.data.utcp_manual import UtcpManual
+from utcp.data.register_manual_response import RegisterManualResult
+from utcp_gnmi.gnmi_call_template import GnmiCallTemplate
+from utcp.data.auth_implementations.api_key_auth import ApiKeyAuth
+from utcp.data.auth_implementations.basic_auth import BasicAuth
+from utcp.data.auth_implementations.oauth2_auth import OAuth2Auth
+
+class GnmiCommunicationProtocol(CommunicationProtocol):
+    async def register_manual(self, caller, manual_call_template: CallTemplate) -> RegisterManualResult:
+        if not isinstance(manual_call_template, GnmiCallTemplate):
+            raise ValueError("GnmiCommunicationProtocol can only be used with GnmiCallTemplate")
+
+        target = manual_call_template.target
+        if manual_call_template.use_tls:
+            pass
+        else:
+            if not (target.startswith("localhost") or target.startswith("127.0.0.1")):
+                return RegisterManualResult(
+                    success=False,
+                    manual_call_template=manual_call_template,
+                    manual=UtcpManual(manual_version="0.0.0", tools=[]),
+                    errors=["Insecure channel only allowed for localhost or 127.0.0.1"]
+                )
+
+        tools: List[Tool] = []
+        ops = ["capabilities", "get", "set", "subscribe"]
+        for op in ops:
+            tct = GnmiCallTemplate(
+                name=manual_call_template.name,
+                call_template_type="gnmi",
+                auth=manual_call_template.auth,
+                target=manual_call_template.target,
+                use_tls=manual_call_template.use_tls,
+                metadata=manual_call_template.metadata,
+                metadata_fields=manual_call_template.metadata_fields,
+                operation=op,
+                stub_module=manual_call_template.stub_module,
+                message_module=manual_call_template.message_module,
+            )
+            inputs = JsonSchema(type="object", properties={})
+            outputs = JsonSchema(type="object", properties={})
+            tool = Tool(
+                name=op,
+                description="",
+                inputs=inputs,
+                outputs=outputs,
+                tags=["gnmi", op],
+                tool_call_template=tct,
+            )
+            tools.append(tool)
+
+        manual = UtcpManual(manual_version="1.0.0", tools=tools)
+        return RegisterManualResult(
+            success=True,
+            manual_call_template=manual_call_template,
+            manual=manual,
+            errors=[],
+        )
+
+    async def deregister_manual(self, caller, manual_call_template: CallTemplate) -> None:
+        return None
+
+    async def call_tool(self, caller, tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate) -> Any:
+        if not isinstance(tool_call_template, GnmiCallTemplate):
+            raise ValueError("GnmiCommunicationProtocol can only be used with GnmiCallTemplate")
+
+        op = tool_call_template.operation
+        target = tool_call_template.target
+
+        metadata: List[tuple[str, str]] = []
+        if tool_call_template.metadata:
+            metadata.extend([(k, v) for k, v in tool_call_template.metadata.items()])
+        if tool_call_template.metadata_fields:
+            for k in tool_call_template.metadata_fields:
+                if k in tool_args:
+                    metadata.append((k, str(tool_args[k])))
+        if tool_call_template.auth:
+            if isinstance(tool_call_template.auth, ApiKeyAuth):
+                if tool_call_template.auth.api_key:
+                    metadata.append((tool_call_template.auth.var_name or "authorization", tool_call_template.auth.api_key))
+            elif isinstance(tool_call_template.auth, BasicAuth):
+                import base64
+                token = base64.b64encode(f"{tool_call_template.auth.username}:{tool_call_template.auth.password}".encode()).decode()
+                metadata.append(("authorization", f"Basic {token}"))
+            elif isinstance(tool_call_template.auth, OAuth2Auth):
+                token = await self._handle_oauth2(tool_call_template.auth)
+                metadata.append(("authorization", f"Bearer {token}"))
+
+        grpc = importlib.import_module("grpc")
+        aio = importlib.import_module("grpc.aio")
+        json_format = importlib.import_module("google.protobuf.json_format")
+        stub_mod = importlib.import_module(tool_call_template.stub_module)
+        msg_mod = importlib.import_module(tool_call_template.message_module)
+
+        if tool_call_template.use_tls:
+            creds = grpc.ssl_channel_credentials()
+            channel = aio.secure_channel(target, creds)
+        else:
+            channel = aio.insecure_channel(target)
+
+        stub = None
+        for attr in dir(stub_mod):
+            if attr.endswith("Stub"):
+                stub_cls = getattr(stub_mod, attr)
+                stub = stub_cls(channel)
+                break
+        if stub is None:
+            raise ValueError("gNMI stub not found in stub_module")
+
+        if op == "capabilities":
+            req = getattr(msg_mod, "CapabilityRequest")()
+            resp = await stub.Capabilities(req, metadata=metadata)
+        elif op == "get":
+            req = getattr(msg_mod, "GetRequest")()
+            paths = tool_args.get("paths", [])
+            for p in paths:
+                path_msg = getattr(msg_mod, "Path")()
+                for elem in [e for e in p.strip("/").split("/") if e]:
+                    pe = getattr(msg_mod, "PathElem")(name=elem)
+                    path_msg.elem.append(pe)
+                req.path.append(path_msg)
+            resp = await stub.Get(req, metadata=metadata)
+        elif op == "set":
+            req = getattr(msg_mod, "SetRequest")()
+            updates = tool_args.get("updates", [])
+            for upd in updates:
+                path_msg = getattr(msg_mod, "Path")()
+                for elem in [e for e in str(upd.get("path", "")).strip("/").split("/") if e]:
+                    pe = getattr(msg_mod, "PathElem")(name=elem)
+                    path_msg.elem.append(pe)
+                val = getattr(msg_mod, "TypedValue")(json_val=str(upd.get("value", "")))
+                update_msg = getattr(msg_mod, "Update")(path=path_msg, val=val)
+                req.update.append(update_msg)
+            resp = await stub.Set(req, metadata=metadata)
+        elif op == "subscribe":
+            req = getattr(msg_mod, "SubscribeRequest")()
+            sub_list = getattr(msg_mod, "SubscriptionList")()
+            mode = tool_args.get("mode", "stream").upper()
+            sub_list.mode = getattr(msg_mod, "SubscriptionList".upper(), None) or 0
+            paths = tool_args.get("paths", [])
+            for p in paths:
+                path_msg = getattr(msg_mod, "Path")()
+                for elem in [e for e in p.strip("/").split("/") if e]:
+                    pe = getattr(msg_mod, "PathElem")(name=elem)
+                    path_msg.elem.append(pe)
+                sub = getattr(msg_mod, "Subscription")(path=path_msg)
+                sub_list.subscription.append(sub)
+            req.subscribe.CopyFrom(sub_list)
+            raise ValueError("Subscribe is streaming; use call_tool_streaming")
+        else:
+            raise ValueError("Unsupported gNMI operation")
+
+        return json_format.MessageToDict(resp)
+
+    async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate) -> AsyncGenerator[Any, None]:
+        if not isinstance(tool_call_template, GnmiCallTemplate):
+            raise ValueError("GnmiCommunicationProtocol can only be used with GnmiCallTemplate")
+        if tool_call_template.operation != "subscribe":
+            result = await self.call_tool(caller, tool_name, tool_args, tool_call_template)
+            yield result
+            return
+        grpc = importlib.import_module("grpc")
+        aio = importlib.import_module("grpc.aio")
+        json_format = importlib.import_module("google.protobuf.json_format")
+        stub_mod = importlib.import_module(tool_call_template.stub_module)
+        msg_mod = importlib.import_module(tool_call_template.message_module)
+        target = tool_call_template.target
+        if tool_call_template.use_tls:
+            creds = grpc.ssl_channel_credentials()
+            channel = aio.secure_channel(target, creds)
+        else:
+            channel = aio.insecure_channel(target)
+        stub = None
+        for attr in dir(stub_mod):
+            if attr.endswith("Stub"):
+                stub_cls = getattr(stub_mod, attr)
+                stub = stub_cls(channel)
+                break
+        if stub is None:
+            raise ValueError("gNMI stub not found in stub_module")
+        metadata: List[tuple[str, str]] = []
+        if tool_call_template.metadata:
+            metadata.extend([(k, v) for k, v in tool_call_template.metadata.items()])
+        if tool_call_template.metadata_fields:
+            for k in tool_call_template.metadata_fields:
+                if k in tool_args:
+                    metadata.append((k, str(tool_args[k])))
+        if tool_call_template.auth:
+            if isinstance(tool_call_template.auth, ApiKeyAuth):
+                if tool_call_template.auth.api_key:
+                    metadata.append((tool_call_template.auth.var_name or "authorization", tool_call_template.auth.api_key))
+            elif isinstance(tool_call_template.auth, BasicAuth):
+                import base64
+                token = base64.b64encode(f"{tool_call_template.auth.username}:{tool_call_template.auth.password}".encode()).decode()
+                metadata.append(("authorization", f"Basic {token}"))
+            elif isinstance(tool_call_template.auth, OAuth2Auth):
+                token = await self._handle_oauth2(tool_call_template.auth)
+                metadata.append(("authorization", f"Bearer {token}"))
+        req = getattr(msg_mod, "SubscribeRequest")()
+        sub_list = getattr(msg_mod, "SubscriptionList")()
+        paths = tool_args.get("paths", [])
+        for p in paths:
+            path_msg = getattr(msg_mod, "Path")()
+            for elem in [e for e in p.strip("/").split("/") if e]:
+                pe = getattr(msg_mod, "PathElem")(name=elem)
+                path_msg.elem.append(pe)
+            sub = getattr(msg_mod, "Subscription")(path=path_msg)
+            sub_list.subscription.append(sub)
+        req.subscribe.CopyFrom(sub_list)
+        call = stub.Subscribe(req, metadata=metadata)
+        async for resp in call:
+            yield json_format.MessageToDict(resp)
+
+    async def _handle_oauth2(self, auth_details: OAuth2Auth) -> str:
+        import aiohttp
+        client_id = auth_details.client_id
+        async with aiohttp.ClientSession() as session:
+            try:
+                body_data = {
+                    "grant_type": "client_credentials",
+                    "client_id": auth_details.client_id,
+                    "client_secret": auth_details.client_secret,
+                    "scope": auth_details.scope,
+                }
+                async with session.post(auth_details.token_url, data=body_data) as response:
+                    response.raise_for_status()
+                    token_response = await response.json()
+                    return token_response["access_token"]
+            except Exception:
+                from aiohttp import BasicAuth as AiohttpBasicAuth
+                header_auth = AiohttpBasicAuth(auth_details.client_id, auth_details.client_secret)
+                header_data = {
+                    "grant_type": "client_credentials",
+                    "scope": auth_details.scope,
+                }
+                async with session.post(auth_details.token_url, data=header_data, auth=header_auth) as response:
+                    response.raise_for_status()
+                    token_response = await response.json()
+                    return token_response["access_token"]

plugins/communication_protocols/gnmi/tests/test_gnmi_plugin.py

@@ -0,0 +1,50 @@
+import sys
+from pathlib import Path
+import pytest
+
+plugin_src = Path(__file__).parent.parent / "src"
+sys.path.insert(0, str(plugin_src))
+
+core_src = Path(__file__).parent.parent.parent.parent.parent / "core" / "src"
+sys.path.insert(0, str(core_src))
+
+from utcp.utcp_client import UtcpClient
+from utcp_gnmi import register
+
+@pytest.mark.asyncio
+async def test_register_manual_and_tools():
+    register()
+    client = await UtcpClient.create(config={
+        "manual_call_templates": [
+            {
+                "name": "routerA",
+                "call_template_type": "gnmi",
+                "target": "localhost:50051",
+                "use_tls": False,
+                "operation": "get"
+            }
+        ]
+    })
+    tools = await client.config.tool_repository.get_tools()
+    names = [t.name for t in tools]
+    assert any(n.startswith("routerA.") for n in names)
+    assert any(n.endswith("subscribe") for n in names)
+
+def test_serializer_roundtrip():
+    from utcp_gnmi.gnmi_call_template import GnmiCallTemplateSerializer
+    serializer = GnmiCallTemplateSerializer()
+    data = {
+        "name": "routerB",
+        "call_template_type": "gnmi",
+        "target": "localhost:50051",
+        "use_tls": False,
+        "metadata": {"authorization": "Bearer token"},
+        "metadata_fields": ["tenant-id"],
+        "operation": "set",
+        "stub_module": "gnmi_pb2_grpc",
+        "message_module": "gnmi_pb2"
+    }
+    obj = serializer.validate_dict(data)
+    out = serializer.to_dict(obj)
+    assert out["call_template_type"] == "gnmi"
+    assert out["operation"] == "set"