@josephyan/qingflow-mcp 0.1.0-beta.2

package/src/qingflow_mcp/tools/auth_tools.py

@@ -0,0 +1,496 @@
+from __future__ import annotations
+
+import base64
+from typing import Any
+
+from mcp.server.fastmcp import FastMCP
+from Crypto.PublicKey import RSA
+from Crypto.Cipher import PKCS1_v1_5
+
+from ..backend_client import BackendRequestContext, BackendResponse
+from ..config import (
+    DEFAULT_PROFILE,
+    get_default_base_url,
+    get_default_qf_version,
+    normalize_base_url,
+)
+from ..errors import QingflowApiError, raise_tool_error
+from ..session_store import SessionStore
+from .base import ToolBase
+
+
+class AuthTools(ToolBase):
+    def __init__(self, sessions: SessionStore, backend) -> None:
+        super().__init__(sessions, backend)
+
+    def register(self, mcp: FastMCP) -> None:
+        @mcp.tool()
+        def auth_login(
+            profile: str = DEFAULT_PROFILE,
+            base_url: str | None = None,
+            qf_version: str | None = None,
+            email: str = "",
+            password: str = "",
+            persist: bool = True,
+        ) -> dict[str, Any]:
+            return self.auth_login(profile=profile, base_url=base_url, qf_version=qf_version, email=email, password=password, persist=persist)
+
+        @mcp.tool()
+        def auth_use_token(
+            profile: str = DEFAULT_PROFILE,
+            base_url: str | None = None,
+            qf_version: str | None = None,
+            token: str = "",
+            ws_id: int | None = None,
+            persist: bool = False,
+        ) -> dict[str, Any]:
+            return self.auth_use_token(profile=profile, base_url=base_url, qf_version=qf_version, token=token, ws_id=ws_id, persist=persist)
+
+        @mcp.tool()
+        def auth_whoami(profile: str = DEFAULT_PROFILE) -> dict[str, Any]:
+            return self.auth_whoami(profile=profile)
+
+        @mcp.tool()
+        def auth_logout(profile: str = DEFAULT_PROFILE, forget_persisted: bool = False) -> dict[str, Any]:
+            return self.auth_logout(profile=profile, forget_persisted=forget_persisted)
+
+    def auth_login(
+        self,
+        *,
+        profile: str = DEFAULT_PROFILE,
+        base_url: str | None = None,
+        qf_version: str | None = None,
+        email: str,
+        password: str,
+        persist: bool,
+    ) -> dict[str, Any]:
+        normalized_base_url = self._normalize_base_url(base_url)
+        normalized_qf_version, qf_version_source = self._resolve_qf_version_input(qf_version)
+        if not email or not password:
+            raise_tool_error(QingflowApiError.config_error("email and password are required"))
+
+        # Try to fetch public key and encrypt password
+        public_key_str = self._fetch_public_key(normalized_base_url, qf_version=normalized_qf_version)
+        encrypted_password = self._encrypt_password(password, public_key_str) if public_key_str else password
+
+        try:
+            # Try 'email' first (aPaas/Public Cloud style)
+            login_response = self.backend.public_request_with_meta(
+                "POST",
+                normalized_base_url,
+                "/user/login",
+                json_body={"email": email, "password": encrypted_password},
+                qf_version=normalized_qf_version,
+            )
+        except QingflowApiError as error:
+            # If failed, try 'account' (QMC/Private Cloud style)
+            try:
+                login_response = self.backend.public_request_with_meta(
+                    "POST",
+                    normalized_base_url,
+                    "/user/login",
+                    json_body={"account": email, "password": encrypted_password},
+                    qf_version=normalized_qf_version,
+                )
+            except QingflowApiError:
+                # If both failed, raise the original error
+                self._handle_error(profile, error)
+                raise AssertionError("unreachable")
+        login_result = login_response.data
+        if not isinstance(login_result, dict):
+            raise_tool_error(QingflowApiError(category="auth", message="Login did not return a valid result"))
+
+        token = login_result.get("token")
+        login_token = login_result.get("loginToken")
+        
+        if not token and login_token:
+            raise_tool_error(
+                QingflowApiError.not_supported(
+                    "Current environment requires an additional login challenge. Qingflow MCP v1 only supports direct token login."
+                )
+            )
+        
+        if not token:
+            raise_tool_error(QingflowApiError(category="auth", message="Login did not return a valid Qingflow token"))
+        detected_qf_version = login_response.qf_response_version
+        resolved_qf_version, resolved_qf_version_source = self._resolve_backend_qf_version(
+            detected_qf_version,
+            fallback_qf_version=normalized_qf_version,
+            fallback_source=qf_version_source,
+        )
+
+        user_info = login_result.get("userVO") or login_result.get("userInfo") or {}
+        if not isinstance(user_info, dict):
+            user_info = {}
+        verified_user_info, verified_qf_version = self._try_fetch_user_info(
+            normalized_base_url,
+            token,
+            qf_version=resolved_qf_version,
+            qf_version_source=resolved_qf_version_source,
+        )
+        if verified_qf_version:
+            resolved_qf_version, resolved_qf_version_source = self._resolve_backend_qf_version(
+                verified_qf_version,
+                fallback_qf_version=resolved_qf_version,
+                fallback_source=resolved_qf_version_source,
+            )
+        if isinstance(verified_user_info, dict):
+            user_info = verified_user_info
+        last_ws_info = user_info.get("lastWsInfo") or {}
+        session_profile = self.sessions.save_session(
+            profile=profile,
+            base_url=normalized_base_url,
+            qf_version=resolved_qf_version,
+            qf_version_source=resolved_qf_version_source,
+            token=token,
+            login_token=login_token,
+            uid=int(user_info.get("uid")),
+            email=user_info.get("email"),
+            nick_name=user_info.get("nickName"),
+            persist=persist,
+        )
+        return {
+            "profile": session_profile.profile,
+            "base_url": session_profile.base_url,
+            "qf_version": session_profile.qf_version,
+            "qf_version_source": session_profile.qf_version_source,
+            "uid": session_profile.uid,
+            "email": session_profile.email,
+            "nick_name": session_profile.nick_name,
+            "selected_ws_id": session_profile.selected_ws_id,
+            "selected_ws_name": session_profile.selected_ws_name,
+            "suggested_ws_id": last_ws_info.get("wsId"),
+            "suggested_ws_name": last_ws_info.get("wsName") or last_ws_info.get("workspaceName"),
+            "persisted": session_profile.persisted,
+            "request_route": self._request_route_payload(
+                BackendRequestContext(
+                    base_url=session_profile.base_url,
+                    token=token,
+                    ws_id=session_profile.selected_ws_id,
+                    qf_version=session_profile.qf_version,
+                    qf_version_source=session_profile.qf_version_source,
+                )
+            ),
+        }
+
+    def auth_use_token(
+        self,
+        *,
+        profile: str = DEFAULT_PROFILE,
+        base_url: str | None = None,
+        qf_version: str | None = None,
+        token: str,
+        ws_id: int | None = None,
+        persist: bool = False,
+    ) -> dict[str, Any]:
+        normalized_base_url = self._normalize_base_url(base_url)
+        normalized_qf_version, qf_version_source = self._resolve_qf_version_input(qf_version)
+        if not token:
+            raise_tool_error(QingflowApiError.config_error("token is required"))
+        if ws_id is not None and ws_id <= 0:
+            raise_tool_error(QingflowApiError.config_error("ws_id must be positive"))
+        try:
+            user_info, detected_qf_version = self._fetch_user_info(
+                normalized_base_url,
+                token,
+                ws_id,
+                qf_version=normalized_qf_version,
+                qf_version_source=qf_version_source,
+            )
+            resolved_qf_version, resolved_qf_version_source = self._resolve_backend_qf_version(
+                detected_qf_version,
+                fallback_qf_version=normalized_qf_version,
+                fallback_source=qf_version_source,
+            )
+            uid = user_info.get("uid")
+            if uid is None:
+                raise_tool_error(QingflowApiError(category="auth", message="Token validation did not return valid user info"))
+            last_ws_info = user_info.get("lastWsInfo") or {}
+            session_profile = self.sessions.save_session(
+                profile=profile,
+                base_url=normalized_base_url,
+                qf_version=resolved_qf_version,
+                qf_version_source=resolved_qf_version_source,
+                token=token,
+                login_token=None,
+                uid=int(uid),
+                email=user_info.get("email"),
+                nick_name=user_info.get("nickName") or user_info.get("displayName") or user_info.get("name"),
+                persist=persist,
+            )
+            selected_ws_name = None
+            if ws_id is not None:
+                workspace = self._fetch_workspace(
+                    normalized_base_url,
+                    token,
+                    ws_id,
+                    qf_version=resolved_qf_version,
+                    qf_version_source=resolved_qf_version_source,
+                )
+                selected_ws_name = workspace.get("workspaceName") or workspace.get("wsName") or workspace.get("remark")
+                session_profile = self.sessions.select_workspace(profile, ws_id=ws_id, ws_name=selected_ws_name)
+            return {
+                "profile": session_profile.profile,
+                "base_url": session_profile.base_url,
+                "qf_version": session_profile.qf_version,
+                "qf_version_source": session_profile.qf_version_source,
+                "uid": session_profile.uid,
+                "email": session_profile.email,
+                "nick_name": session_profile.nick_name,
+                "selected_ws_id": session_profile.selected_ws_id,
+                "selected_ws_name": session_profile.selected_ws_name,
+                "suggested_ws_id": last_ws_info.get("wsId"),
+                "suggested_ws_name": last_ws_info.get("wsName") or last_ws_info.get("workspaceName"),
+                "persisted": session_profile.persisted,
+                "request_route": self._request_route_payload(
+                    BackendRequestContext(
+                        base_url=session_profile.base_url,
+                        token=token,
+                        ws_id=session_profile.selected_ws_id,
+                        qf_version=session_profile.qf_version,
+                        qf_version_source=session_profile.qf_version_source,
+                    )
+                ),
+            }
+        except QingflowApiError as error:
+            self._handle_error(profile, error)
+            raise AssertionError("unreachable")
+
+    def auth_whoami(self, *, profile: str = DEFAULT_PROFILE) -> dict[str, Any]:
+        try:
+            session_profile, _, context = self._require_context(profile, require_workspace=False)
+        except QingflowApiError as error:
+            self._handle_error(profile, error)
+            raise AssertionError("unreachable")
+        return {
+            "profile": session_profile.profile,
+            "base_url": session_profile.base_url,
+            "qf_version": session_profile.qf_version,
+            "qf_version_source": session_profile.qf_version_source,
+            "uid": session_profile.uid,
+            "email": session_profile.email,
+            "nick_name": session_profile.nick_name,
+            "selected_ws_id": session_profile.selected_ws_id,
+            "selected_ws_name": session_profile.selected_ws_name,
+            "persisted": session_profile.persisted,
+            "request_route": self._request_route_payload(context),
+        }
+
+    def auth_logout(self, *, profile: str = DEFAULT_PROFILE, forget_persisted: bool = False) -> dict[str, Any]:
+        if not self.sessions.has_profile(profile):
+            raise_tool_error(QingflowApiError.auth_required(profile))
+        self.sessions.logout(profile, forget_persisted=forget_persisted)
+        return {
+            "profile": profile,
+            "logged_out": True,
+            "forgot_persisted": forget_persisted,
+        }
+
+    def _normalize_base_url(self, base_url: str | None) -> str:
+        normalized_base_url = normalize_base_url(base_url) or get_default_base_url()
+        if not normalized_base_url:
+            raise_tool_error(
+                QingflowApiError.config_error(
+                    "base_url is required or configure default_base_url / QINGFLOW_MCP_DEFAULT_BASE_URL"
+                )
+            )
+        return normalized_base_url
+
+    def _normalize_qf_version(self, qf_version: str | None) -> str | None:
+        if qf_version is not None:
+            normalized = str(qf_version).strip()
+            return normalized or None
+        return get_default_qf_version()
+
+    def _resolve_qf_version_input(self, qf_version: str | None) -> tuple[str | None, str]:
+        if qf_version is not None:
+            normalized = self._normalize_qf_version(qf_version)
+            return normalized, "explicit" if normalized else "unset"
+        normalized = self._normalize_qf_version(None)
+        if normalized:
+            return normalized, "default_config"
+        return None, "unset"
+
+    def _resolve_backend_qf_version(
+        self,
+        backend_qf_version: str | None,
+        *,
+        fallback_qf_version: str | None,
+        fallback_source: str,
+    ) -> tuple[str | None, str]:
+        if backend_qf_version:
+            return backend_qf_version, "backend_response"
+        return fallback_qf_version, fallback_source
+
+    def _fetch_user_info(
+        self,
+        base_url: str,
+        token: str,
+        ws_id: int | None,
+        *,
+        qf_version: str | None,
+        qf_version_source: str | None,
+    ) -> tuple[dict[str, Any], str | None]:
+        request_context = BackendRequestContext(
+            base_url=base_url,
+            token=token,
+            ws_id=ws_id,
+            qf_version=qf_version,
+            qf_version_source=qf_version_source,
+        )
+        try:
+            user_response = self.backend.request_with_meta("GET", request_context, "/user")
+            user_info = user_response.data
+            if isinstance(user_info, dict):
+                return user_info, user_response.qf_response_version
+        except QingflowApiError as original_error:
+            if ws_id is not None:
+                raise original_error
+            first_workspace, workspace_qf_version = self._fetch_first_workspace(
+                base_url,
+                token,
+                qf_version=qf_version,
+                qf_version_source=qf_version_source,
+            )
+            if not first_workspace:
+                raise original_error
+            first_ws_id = first_workspace.get("wsId")
+            if not first_ws_id:
+                raise original_error
+            effective_qf_version = workspace_qf_version or qf_version
+            effective_qf_version_source = "backend_response" if workspace_qf_version else qf_version_source
+            fallback_context = BackendRequestContext(
+                base_url=base_url,
+                token=token,
+                ws_id=int(first_ws_id),
+                qf_version=effective_qf_version,
+                qf_version_source=effective_qf_version_source,
+            )
+            user_response = self.backend.request_with_meta("GET", fallback_context, "/user")
+            user_info = user_response.data
+            if isinstance(user_info, dict):
+                return user_info, user_response.qf_response_version or effective_qf_version
+            raise original_error
+        raise_tool_error(QingflowApiError(category="auth", message="Token validation did not return valid user info"))
+
+    def _try_fetch_user_info(
+        self,
+        base_url: str,
+        token: str,
+        *,
+        qf_version: str | None,
+        qf_version_source: str | None,
+    ) -> tuple[dict[str, Any] | None, str | None]:
+        try:
+            return self._fetch_user_info(
+                base_url,
+                token,
+                None,
+                qf_version=qf_version,
+                qf_version_source=qf_version_source,
+            )
+        except QingflowApiError:
+            return None, None
+
+    def _fetch_first_workspace(
+        self,
+        base_url: str,
+        token: str,
+        *,
+        qf_version: str | None,
+        qf_version_source: str | None,
+    ) -> tuple[dict[str, Any] | None, str | None]:
+        page_response = self.backend.request_with_meta(
+            "POST",
+            BackendRequestContext(
+                base_url=base_url,
+                token=token,
+                ws_id=None,
+                qf_version=qf_version,
+                qf_version_source=qf_version_source,
+            ),
+            "/user/workspaceList/pageQuery",
+            json_body={"pageNum": 1, "pageSize": 1},
+        )
+        page = page_response.data
+        if not isinstance(page, dict):
+            return None, page_response.qf_response_version
+        workspaces = page.get("list") or []
+        if not workspaces:
+            return None, page_response.qf_response_version
+        first_workspace = workspaces[0]
+        return (first_workspace if isinstance(first_workspace, dict) else None), page_response.qf_response_version
+
+    def _fetch_workspace(
+        self,
+        base_url: str,
+        token: str,
+        ws_id: int,
+        *,
+        qf_version: str | None,
+        qf_version_source: str | None,
+    ) -> dict[str, Any]:
+        workspace = self.backend.request(
+            "GET",
+            BackendRequestContext(
+                base_url=base_url,
+                token=token,
+                ws_id=None,
+                qf_version=qf_version,
+                qf_version_source=qf_version_source,
+            ),
+            f"/user/workspace/{ws_id}",
+        )
+        if not isinstance(workspace, dict):
+            raise_tool_error(QingflowApiError(category="workspace", message=f"Workspace {ws_id} is not accessible"))
+        return workspace
+
+    def _request_route_payload(self, context: BackendRequestContext) -> dict[str, Any]:
+        describe_route = getattr(self.backend, "describe_route", None)
+        if callable(describe_route):
+            payload = describe_route(context)
+            if isinstance(payload, dict):
+                return payload
+        return {
+            "base_url": context.base_url,
+            "qf_version": context.qf_version,
+            "qf_version_source": context.qf_version_source or ("context" if context.qf_version else "unknown"),
+        }
+
+    def _fetch_public_key(self, base_url: str, *, qf_version: str | None) -> str | None:
+        # Endpoints to try (order matters, lowercase 'pubkey' is for Public Cloud)
+        endpoints = ["/user/pubkey", "/api/user/pubkey", "/user/publicKey", "/api/user/publicKey"]
+        for endpoint in endpoints:
+            try:
+                # We use unwrap=False to handle various response formats
+                result = self.backend.public_request("GET", base_url, endpoint, unwrap=False, qf_version=qf_version)
+                if isinstance(result, dict):
+                    # Try various common response structures
+                    data = result.get("data") or result.get("result") or result
+                    if isinstance(data, dict):
+                        # Try case-insensitive keys
+                        for key in ["pubkey", "publicKey", "pubKey"]:
+                            if key in data:
+                                return str(data[key])
+                    if isinstance(data, str) and not data.startswith("{"):
+                        return data
+            except Exception:
+                continue
+        
+        return None
+
+    def _encrypt_password(self, password: str, public_key_str: str) -> str:
+        try:
+            if not public_key_str.startswith("-----BEGIN"):
+                key_content = public_key_str.strip()
+                formatted_key = f"-----BEGIN PUBLIC KEY-----\n{key_content}\n-----END PUBLIC KEY-----"
+            else:
+                formatted_key = public_key_str
+            
+            key = RSA.import_key(formatted_key)
+            cipher = PKCS1_v1_5.new(key)
+            encrypted = cipher.encrypt(password.encode("utf-8"))
+            return base64.b64encode(encrypted).decode("utf-8")
+        except Exception as e:
+            # If encryption fails, fallback to plain text (might be a legacy or custom environment)
+            return password

package/src/qingflow_mcp/tools/base.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+from typing import Callable, TypeVar
+
+from ..backend_client import BackendRequestContext, BackendClient
+from ..errors import QingflowApiError, raise_tool_error
+from ..json_types import JSONObject
+from ..session_store import BackendSession, SessionProfile, SessionStore
+
+
+T = TypeVar("T")
+
+
+class ToolBase:
+    def __init__(self, sessions: SessionStore, backend: BackendClient) -> None:
+        self.sessions = sessions
+        self.backend = backend
+
+    def _require_context(self, profile: str, *, require_workspace: bool) -> tuple[SessionProfile, BackendSession, BackendRequestContext]:
+        session_profile = self.sessions.get_profile(profile)
+        if session_profile is None:
+            raise QingflowApiError.auth_required(profile)
+        backend_session = self.sessions.get_backend_session(profile)
+        if backend_session is None:
+            raise QingflowApiError.auth_required(profile)
+        if require_workspace and session_profile.selected_ws_id is None:
+            raise QingflowApiError.workspace_not_selected(profile)
+        context = BackendRequestContext(
+            base_url=backend_session.base_url,
+            token=backend_session.token,
+            ws_id=session_profile.selected_ws_id if require_workspace else None,
+            qf_version=backend_session.qf_version,
+            qf_version_source=backend_session.qf_version_source,
+        )
+        return session_profile, backend_session, context
+
+    def _run(self, profile: str, func: Callable[[SessionProfile, BackendRequestContext], T], *, require_workspace: bool = True) -> T:
+        try:
+            session_profile, _, context = self._require_context(profile, require_workspace=require_workspace)
+            return func(session_profile, context)
+        except QingflowApiError as error:
+            self._handle_error(profile, error)
+        raise AssertionError("unreachable")
+
+    def _handle_error(self, profile: str, error: QingflowApiError) -> None:
+        if error.looks_like_invalid_token():
+            self.sessions.invalidate(profile)
+            error = QingflowApiError(
+                category="auth",
+                message=f"Qingflow session for profile '{profile}' has expired. Run auth_login again.",
+                backend_code=error.backend_code,
+                request_id=error.request_id,
+                http_status=error.http_status,
+            )
+        raise_tool_error(error)
+
+    def _require_dict(self, payload: JSONObject | None, field_name: str = "payload") -> JSONObject:
+        if not isinstance(payload, dict) or not payload:
+            raise_tool_error(QingflowApiError.config_error(f"{field_name} must be a non-empty object"))
+        return payload
+
+    def _high_risk_tool_description(self, *, operation: str, target: str) -> str:
+        return (
+            f"High-risk {operation} operation for {target}. Read the current state first, "
+            "confirm the exact target IDs and intended diff with a human, and avoid running "
+            "against production without explicit approval."
+        )
+
+    def _attach_human_review_notice(self, response: JSONObject, *, operation: str, target: str) -> JSONObject:
+        payload = dict(response)
+        payload["requires_human_review"] = True
+        payload["risk_notice"] = {
+            "operation": operation,
+            "target": target,
+            "severity": "high",
+            "guidance": (
+                "Read the current state first, confirm the exact target IDs and intended diff with a human, "
+                "and require explicit approval before running against production."
+            ),
+        }
+        return payload