trustgraph-base 2.4.4__tar.gz → 2.4.5__tar.gz

{trustgraph_base-2.4.4 → trustgraph_base-2.4.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-base
-Version: 2.4.4
+Version: 2.4.5
 Summary: TrustGraph provides a means to run a pipeline of flexible AI processing components in a flexible means to achieve a processing pipeline.
 Author-email: "trustgraph.ai" <security@trustgraph.ai>
 Project-URL: Homepage, https://github.com/trustgraph-ai/trustgraph

{trustgraph_base-2.4.4 → trustgraph_base-2.4.5}/trustgraph/api/async_socket_client.py

@@ -49,21 +49,67 @@ class AsyncSocketClient:
             return f"ws://{url}"
 
     def _build_ws_url(self):
-        ws_url = f"{self.url.rstrip('/')}/api/v1/socket"
-        if self.token:
-            ws_url = f"{ws_url}?token={self.token}"
-        return ws_url
+        # /api/v1/socket uses the first-frame auth protocol — the
+        # token is sent as the first frame after connecting rather
+        # than in the URL.  This avoids browser issues with 401 on
+        # the WebSocket handshake and lets long-lived sockets
+        # refresh credentials mid-session.
+        return f"{self.url.rstrip('/')}/api/v1/socket"
 
     async def connect(self):
-        """Establish the persistent websocket connection."""
+        """Establish the persistent websocket connection and run the
+        first-frame auth handshake."""
         if self._connected:
             return
 
+        if not self.token:
+            raise ProtocolException(
+                "AsyncSocketClient requires a token for first-frame "
+                "auth against /api/v1/socket"
+            )
+
         ws_url = self._build_ws_url()
         self._connect_cm = websockets.connect(
             ws_url, ping_interval=20, ping_timeout=self.timeout
         )
         self._socket = await self._connect_cm.__aenter__()
+
+        # First-frame auth: send {"type":"auth","token":"..."} and
+        # wait for auth-ok / auth-failed.  Run before starting the
+        # reader task so the response isn't consumed by the reader's
+        # id-based routing.
+        await self._socket.send(json.dumps({
+            "type": "auth", "token": self.token,
+        }))
+        try:
+            raw = await asyncio.wait_for(
+                self._socket.recv(), timeout=self.timeout,
+            )
+        except asyncio.TimeoutError:
+            await self._socket.close()
+            raise ProtocolException("Timeout waiting for auth response")
+
+        try:
+            resp = json.loads(raw)
+        except Exception:
+            await self._socket.close()
+            raise ProtocolException(
+                f"Unexpected non-JSON auth response: {raw!r}"
+            )
+
+        if resp.get("type") == "auth-ok":
+            self.workspace = resp.get("workspace", self.workspace)
+        elif resp.get("type") == "auth-failed":
+            await self._socket.close()
+            raise ProtocolException(
+                f"auth failure: {resp.get('error', 'unknown')}"
+            )
+        else:
+            await self._socket.close()
+            raise ProtocolException(
+                f"Unexpected auth response: {resp!r}"
+            )
+
         self._connected = True
         self._reader_task = asyncio.create_task(self._reader())
 

{trustgraph_base-2.4.4 → trustgraph_base-2.4.5}/trustgraph/api/socket_client.py

@@ -112,10 +112,10 @@ class SocketClient:
             return f"ws://{url}"
 
     def _build_ws_url(self):
-        ws_url = f"{self.url.rstrip('/')}/api/v1/socket"
-        if self.token:
-            ws_url = f"{ws_url}?token={self.token}"
-        return ws_url
+        # /api/v1/socket uses the first-frame auth protocol — the
+        # token is sent as the first frame after connecting rather
+        # than in the URL.
+        return f"{self.url.rstrip('/')}/api/v1/socket"
 
     def _get_loop(self):
         """Get or create the event loop, reusing across calls."""
@@ -132,15 +132,58 @@ class SocketClient:
         return self._loop
 
     async def _ensure_connected(self):
-        """Lazily establish the persistent websocket connection."""
+        """Lazily establish the persistent websocket connection and
+        run the first-frame auth handshake."""
         if self._connected:
             return
 
+        if not self.token:
+            raise ProtocolException(
+                "SocketClient requires a token for first-frame auth "
+                "against /api/v1/socket"
+            )
+
         ws_url = self._build_ws_url()
         self._connect_cm = websockets.connect(
             ws_url, ping_interval=20, ping_timeout=self.timeout
         )
         self._socket = await self._connect_cm.__aenter__()
+
+        # First-frame auth — run before starting the reader so the
+        # auth-ok / auth-failed response isn't consumed by the reader
+        # loop's id-based routing.
+        await self._socket.send(json.dumps({
+            "type": "auth", "token": self.token,
+        }))
+        try:
+            raw = await asyncio.wait_for(
+                self._socket.recv(), timeout=self.timeout,
+            )
+        except asyncio.TimeoutError:
+            await self._socket.close()
+            raise ProtocolException("Timeout waiting for auth response")
+
+        try:
+            resp = json.loads(raw)
+        except Exception:
+            await self._socket.close()
+            raise ProtocolException(
+                f"Unexpected non-JSON auth response: {raw!r}"
+            )
+
+        if resp.get("type") == "auth-ok":
+            self.workspace = resp.get("workspace", self.workspace)
+        elif resp.get("type") == "auth-failed":
+            await self._socket.close()
+            raise ProtocolException(
+                f"auth failure: {resp.get('error', 'unknown')}"
+            )
+        else:
+            await self._socket.close()
+            raise ProtocolException(
+                f"Unexpected auth response: {resp!r}"
+            )
+
         self._connected = True
         self._reader_task = asyncio.create_task(self._reader())
 

trustgraph_base-2.4.5/trustgraph/base/iam_client.py

@@ -0,0 +1,279 @@
+
+from . request_response_spec import RequestResponse, RequestResponseSpec
+from .. schema import (
+    IamRequest, IamResponse,
+    UserInput, WorkspaceInput, ApiKeyInput,
+)
+
+IAM_TIMEOUT = 10
+
+
+class IamClient(RequestResponse):
+    """Client for the IAM service request/response pub/sub protocol.
+
+    Mirrors ``ConfigClient``: a thin wrapper around ``RequestResponse``
+    that knows the IAM request / response schemas.  Only the subset of
+    operations actually implemented by the server today has helper
+    methods here; callers that need an unimplemented operation can
+    build ``IamRequest`` and call ``request()`` directly.
+    """
+
+    async def _request(self, timeout=IAM_TIMEOUT, **kwargs):
+        resp = await self.request(
+            IamRequest(**kwargs),
+            timeout=timeout,
+        )
+        if resp.error:
+            raise RuntimeError(
+                f"{resp.error.type}: {resp.error.message}"
+            )
+        return resp
+
+    async def bootstrap(self, timeout=IAM_TIMEOUT):
+        """Initial-run IAM self-seed.  Returns a tuple of
+        ``(admin_user_id, admin_api_key_plaintext)``.  Both are empty
+        strings on repeat calls — the operation is a no-op once the
+        IAM tables are populated."""
+        resp = await self._request(
+            operation="bootstrap", timeout=timeout,
+        )
+        return resp.bootstrap_admin_user_id, resp.bootstrap_admin_api_key
+
+    async def resolve_api_key(self, api_key, timeout=IAM_TIMEOUT):
+        """Resolve a plaintext API key to its identity triple.
+
+        Returns ``(user_id, workspace, roles)`` or raises
+        ``RuntimeError`` with error type ``auth-failed`` if the key is
+        unknown / expired / revoked."""
+        resp = await self._request(
+            operation="resolve-api-key",
+            api_key=api_key,
+            timeout=timeout,
+        )
+        return (
+            resp.resolved_user_id,
+            resp.resolved_workspace,
+            list(resp.resolved_roles),
+        )
+
+    async def create_user(self, workspace, user, actor="",
+                          timeout=IAM_TIMEOUT):
+        """Create a user.  ``user`` is a ``UserInput``."""
+        resp = await self._request(
+            operation="create-user",
+            workspace=workspace,
+            actor=actor,
+            user=user,
+            timeout=timeout,
+        )
+        return resp.user
+
+    async def list_users(self, workspace, actor="", timeout=IAM_TIMEOUT):
+        resp = await self._request(
+            operation="list-users",
+            workspace=workspace,
+            actor=actor,
+            timeout=timeout,
+        )
+        return list(resp.users)
+
+    async def create_api_key(self, workspace, key, actor="",
+                             timeout=IAM_TIMEOUT):
+        """Create an API key.  ``key`` is an ``ApiKeyInput``.  Returns
+        ``(plaintext, record)`` — plaintext is returned once and the
+        caller is responsible for surfacing it to the operator."""
+        resp = await self._request(
+            operation="create-api-key",
+            workspace=workspace,
+            actor=actor,
+            key=key,
+            timeout=timeout,
+        )
+        return resp.api_key_plaintext, resp.api_key
+
+    async def list_api_keys(self, workspace, user_id, actor="",
+                            timeout=IAM_TIMEOUT):
+        resp = await self._request(
+            operation="list-api-keys",
+            workspace=workspace,
+            actor=actor,
+            user_id=user_id,
+            timeout=timeout,
+        )
+        return list(resp.api_keys)
+
+    async def revoke_api_key(self, workspace, key_id, actor="",
+                             timeout=IAM_TIMEOUT):
+        await self._request(
+            operation="revoke-api-key",
+            workspace=workspace,
+            actor=actor,
+            key_id=key_id,
+            timeout=timeout,
+        )
+
+    async def login(self, username, password, workspace="",
+                    timeout=IAM_TIMEOUT):
+        """Validate credentials and return ``(jwt, expires_iso)``.
+        ``workspace`` is optional; defaults at the server to the
+        OSS default workspace."""
+        resp = await self._request(
+            operation="login",
+            workspace=workspace,
+            username=username,
+            password=password,
+            timeout=timeout,
+        )
+        return resp.jwt, resp.jwt_expires
+
+    async def get_signing_key_public(self, timeout=IAM_TIMEOUT):
+        """Return the active JWT signing public key in PEM.  The
+        gateway calls this at startup and caches the result."""
+        resp = await self._request(
+            operation="get-signing-key-public",
+            timeout=timeout,
+        )
+        return resp.signing_key_public
+
+    async def change_password(self, user_id, current_password,
+                              new_password, timeout=IAM_TIMEOUT):
+        await self._request(
+            operation="change-password",
+            user_id=user_id,
+            password=current_password,
+            new_password=new_password,
+            timeout=timeout,
+        )
+
+    async def reset_password(self, workspace, user_id, actor="",
+                             timeout=IAM_TIMEOUT):
+        """Admin-driven password reset.  Returns the plaintext
+        temporary password (returned once)."""
+        resp = await self._request(
+            operation="reset-password",
+            workspace=workspace,
+            actor=actor,
+            user_id=user_id,
+            timeout=timeout,
+        )
+        return resp.temporary_password
+
+    async def get_user(self, workspace, user_id, actor="",
+                       timeout=IAM_TIMEOUT):
+        resp = await self._request(
+            operation="get-user",
+            workspace=workspace,
+            actor=actor,
+            user_id=user_id,
+            timeout=timeout,
+        )
+        return resp.user
+
+    async def update_user(self, workspace, user_id, user, actor="",
+                          timeout=IAM_TIMEOUT):
+        resp = await self._request(
+            operation="update-user",
+            workspace=workspace,
+            actor=actor,
+            user_id=user_id,
+            user=user,
+            timeout=timeout,
+        )
+        return resp.user
+
+    async def disable_user(self, workspace, user_id, actor="",
+                           timeout=IAM_TIMEOUT):
+        await self._request(
+            operation="disable-user",
+            workspace=workspace,
+            actor=actor,
+            user_id=user_id,
+            timeout=timeout,
+        )
+
+    async def enable_user(self, workspace, user_id, actor="",
+                          timeout=IAM_TIMEOUT):
+        await self._request(
+            operation="enable-user",
+            workspace=workspace,
+            actor=actor,
+            user_id=user_id,
+            timeout=timeout,
+        )
+
+    async def delete_user(self, workspace, user_id, actor="",
+                          timeout=IAM_TIMEOUT):
+        await self._request(
+            operation="delete-user",
+            workspace=workspace,
+            actor=actor,
+            user_id=user_id,
+            timeout=timeout,
+        )
+
+    async def create_workspace(self, workspace_record, actor="",
+                               timeout=IAM_TIMEOUT):
+        resp = await self._request(
+            operation="create-workspace",
+            actor=actor,
+            workspace_record=workspace_record,
+            timeout=timeout,
+        )
+        return resp.workspace
+
+    async def list_workspaces(self, actor="", timeout=IAM_TIMEOUT):
+        resp = await self._request(
+            operation="list-workspaces",
+            actor=actor,
+            timeout=timeout,
+        )
+        return list(resp.workspaces)
+
+    async def get_workspace(self, workspace_id, actor="",
+                            timeout=IAM_TIMEOUT):
+        from ..schema import WorkspaceInput
+        resp = await self._request(
+            operation="get-workspace",
+            actor=actor,
+            workspace_record=WorkspaceInput(id=workspace_id),
+            timeout=timeout,
+        )
+        return resp.workspace
+
+    async def update_workspace(self, workspace_record, actor="",
+                               timeout=IAM_TIMEOUT):
+        resp = await self._request(
+            operation="update-workspace",
+            actor=actor,
+            workspace_record=workspace_record,
+            timeout=timeout,
+        )
+        return resp.workspace
+
+    async def disable_workspace(self, workspace_id, actor="",
+                                timeout=IAM_TIMEOUT):
+        from ..schema import WorkspaceInput
+        await self._request(
+            operation="disable-workspace",
+            actor=actor,
+            workspace_record=WorkspaceInput(id=workspace_id),
+            timeout=timeout,
+        )
+
+    async def rotate_signing_key(self, actor="", timeout=IAM_TIMEOUT):
+        await self._request(
+            operation="rotate-signing-key",
+            actor=actor,
+            timeout=timeout,
+        )
+
+
+class IamClientSpec(RequestResponseSpec):
+    def __init__(self, request_name, response_name):
+        super().__init__(
+            request_name=request_name,
+            request_schema=IamRequest,
+            response_name=response_name,
+            response_schema=IamResponse,
+            impl=IamClient,
+        )

trustgraph_base-2.4.5/trustgraph/base_version.py

@@ -0,0 +1 @@
+__version__ = "2.4.5"

{trustgraph_base-2.4.4 → trustgraph_base-2.4.5}/trustgraph/messaging/__init__.py

@@ -15,6 +15,7 @@ from .translators.library import LibraryRequestTranslator, LibraryResponseTransl
 from .translators.document_loading import DocumentTranslator, TextDocumentTranslator
 from .translators.config import ConfigRequestTranslator, ConfigResponseTranslator
 from .translators.flow import FlowRequestTranslator, FlowResponseTranslator
+from .translators.iam import IamRequestTranslator, IamResponseTranslator
 from .translators.prompt import PromptRequestTranslator, PromptResponseTranslator
 from .translators.tool import ToolRequestTranslator, ToolResponseTranslator
 from .translators.embeddings_query import (
@@ -85,11 +86,17 @@ TranslatorRegistry.register_service(
 )
 
 TranslatorRegistry.register_service(
-    "flow", 
-    FlowRequestTranslator(), 
+    "flow",
+    FlowRequestTranslator(),
     FlowResponseTranslator()
 )
 
+TranslatorRegistry.register_service(
+    "iam",
+    IamRequestTranslator(),
+    IamResponseTranslator()
+)
+
 TranslatorRegistry.register_service(
     "prompt", 
     PromptRequestTranslator(), 

trustgraph_base-2.4.5/trustgraph/messaging/translators/iam.py

@@ -0,0 +1,194 @@
+from typing import Dict, Any, Tuple
+
+from ...schema import IamRequest, IamResponse
+from ...schema import (
+    UserInput, UserRecord,
+    WorkspaceInput, WorkspaceRecord,
+    ApiKeyInput, ApiKeyRecord,
+)
+from .base import MessageTranslator
+
+
+def _user_input_from_dict(d):
+    if d is None:
+        return None
+    return UserInput(
+        username=d.get("username", ""),
+        name=d.get("name", ""),
+        email=d.get("email", ""),
+        password=d.get("password", ""),
+        roles=list(d.get("roles", [])),
+        enabled=d.get("enabled", True),
+        must_change_password=d.get("must_change_password", False),
+    )
+
+
+def _workspace_input_from_dict(d):
+    if d is None:
+        return None
+    return WorkspaceInput(
+        id=d.get("id", ""),
+        name=d.get("name", ""),
+        enabled=d.get("enabled", True),
+    )
+
+
+def _api_key_input_from_dict(d):
+    if d is None:
+        return None
+    return ApiKeyInput(
+        user_id=d.get("user_id", ""),
+        name=d.get("name", ""),
+        expires=d.get("expires", ""),
+    )
+
+
+def _user_record_to_dict(r):
+    if r is None:
+        return None
+    return {
+        "id": r.id,
+        "workspace": r.workspace,
+        "username": r.username,
+        "name": r.name,
+        "email": r.email,
+        "roles": list(r.roles),
+        "enabled": r.enabled,
+        "must_change_password": r.must_change_password,
+        "created": r.created,
+    }
+
+
+def _workspace_record_to_dict(r):
+    if r is None:
+        return None
+    return {
+        "id": r.id,
+        "name": r.name,
+        "enabled": r.enabled,
+        "created": r.created,
+    }
+
+
+def _api_key_record_to_dict(r):
+    if r is None:
+        return None
+    return {
+        "id": r.id,
+        "user_id": r.user_id,
+        "name": r.name,
+        "prefix": r.prefix,
+        "expires": r.expires,
+        "created": r.created,
+        "last_used": r.last_used,
+    }
+
+
+class IamRequestTranslator(MessageTranslator):
+
+    def decode(self, data: Dict[str, Any]) -> IamRequest:
+        return IamRequest(
+            operation=data.get("operation", ""),
+            workspace=data.get("workspace", ""),
+            actor=data.get("actor", ""),
+            user_id=data.get("user_id", ""),
+            username=data.get("username", ""),
+            key_id=data.get("key_id", ""),
+            api_key=data.get("api_key", ""),
+            password=data.get("password", ""),
+            new_password=data.get("new_password", ""),
+            user=_user_input_from_dict(data.get("user")),
+            workspace_record=_workspace_input_from_dict(
+                data.get("workspace_record")
+            ),
+            key=_api_key_input_from_dict(data.get("key")),
+        )
+
+    def encode(self, obj: IamRequest) -> Dict[str, Any]:
+        result = {"operation": obj.operation}
+        for fname in (
+                "workspace", "actor", "user_id", "username", "key_id",
+                "api_key", "password", "new_password",
+        ):
+            v = getattr(obj, fname, "")
+            if v:
+                result[fname] = v
+        if obj.user is not None:
+            result["user"] = {
+                "username": obj.user.username,
+                "name": obj.user.name,
+                "email": obj.user.email,
+                "password": obj.user.password,
+                "roles": list(obj.user.roles),
+                "enabled": obj.user.enabled,
+                "must_change_password": obj.user.must_change_password,
+            }
+        if obj.workspace_record is not None:
+            result["workspace_record"] = {
+                "id": obj.workspace_record.id,
+                "name": obj.workspace_record.name,
+                "enabled": obj.workspace_record.enabled,
+            }
+        if obj.key is not None:
+            result["key"] = {
+                "user_id": obj.key.user_id,
+                "name": obj.key.name,
+                "expires": obj.key.expires,
+            }
+        return result
+
+
+class IamResponseTranslator(MessageTranslator):
+
+    def decode(self, data: Dict[str, Any]) -> IamResponse:
+        raise NotImplementedError(
+            "IamResponse is a server-produced message; no HTTP→schema "
+            "path is needed"
+        )
+
+    def encode(self, obj: IamResponse) -> Dict[str, Any]:
+        result: Dict[str, Any] = {}
+
+        if obj.user is not None:
+            result["user"] = _user_record_to_dict(obj.user)
+        if obj.users:
+            result["users"] = [_user_record_to_dict(u) for u in obj.users]
+        if obj.workspace is not None:
+            result["workspace"] = _workspace_record_to_dict(obj.workspace)
+        if obj.workspaces:
+            result["workspaces"] = [
+                _workspace_record_to_dict(w) for w in obj.workspaces
+            ]
+        if obj.api_key_plaintext:
+            result["api_key_plaintext"] = obj.api_key_plaintext
+        if obj.api_key is not None:
+            result["api_key"] = _api_key_record_to_dict(obj.api_key)
+        if obj.api_keys:
+            result["api_keys"] = [
+                _api_key_record_to_dict(k) for k in obj.api_keys
+            ]
+        if obj.jwt:
+            result["jwt"] = obj.jwt
+        if obj.jwt_expires:
+            result["jwt_expires"] = obj.jwt_expires
+        if obj.signing_key_public:
+            result["signing_key_public"] = obj.signing_key_public
+        if obj.resolved_user_id:
+            result["resolved_user_id"] = obj.resolved_user_id
+        if obj.resolved_workspace:
+            result["resolved_workspace"] = obj.resolved_workspace
+        if obj.resolved_roles:
+            result["resolved_roles"] = list(obj.resolved_roles)
+        if obj.temporary_password:
+            result["temporary_password"] = obj.temporary_password
+        if obj.bootstrap_admin_user_id:
+            result["bootstrap_admin_user_id"] = obj.bootstrap_admin_user_id
+        if obj.bootstrap_admin_api_key:
+            result["bootstrap_admin_api_key"] = obj.bootstrap_admin_api_key
+
+        return result
+
+    def encode_with_completion(
+            self, obj: IamResponse,
+    ) -> Tuple[Dict[str, Any], bool]:
+        return self.encode(obj), True

{trustgraph_base-2.4.4 → trustgraph_base-2.4.5}/trustgraph/schema/services/__init__.py

@@ -5,6 +5,7 @@ from .agent import *
 from .flow import *
 from .prompt import *
 from .config import *
+from .iam import *
 from .library import *
 from .lookup import *
 from .nlp_query import *