workflow-exec-engine 0.0.2__py3-none-any.whl

workflow_engine/client/auth_manager.py

@@ -0,0 +1,127 @@
+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
+# All Rights Reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+"""Auth manager - builds interceptors from AgentCard, self-contained."""
+
+from typing import Dict, Any, List, Optional
+from loguru import logger
+
+try:
+    from a2a.client.auth import AuthInterceptor
+    from a2a.client.interceptors import ClientCallInterceptor, BeforeArgs, AfterArgs
+    _A2A_AUTH_AVAILABLE = True
+except ImportError:
+    _A2A_AUTH_AVAILABLE = False
+    AuthInterceptor = None
+    ClientCallInterceptor = None
+    BeforeArgs = None
+    AfterArgs = None
+
+
+from workflow_engine.client.credential_service import AgentAuthManager, CustomAuthInterceptor
+from workflow_engine.client.extension_interceptor import ExtensionInterceptor
+from workflow_engine.client.auth_provider import AuthProvider
+
+
+class AuthManager:
+    """Builds auth/extension interceptors from AgentCard securitySchemes."""
+
+    def __init__(self, agent_cards: List[Any], credentials_config: Optional[str | Dict] = None):
+        self._interceptors: Dict[str, List[Any]] = {}
+        self._auth_manager: Optional[AgentAuthManager] = None
+
+        if not _A2A_AUTH_AVAILABLE:
+            logger.info("[AuthManager] a2a-sdk auth not available, authentication disabled")
+            return
+        logger.info(f"[AuthManager] Initializing with {len(agent_cards)} agent card(s), config={credentials_config is not None}")
+
+        if credentials_config:
+            if isinstance(credentials_config, str):
+                self._auth_manager = AgentAuthManager(config_path=credentials_config)
+            elif isinstance(credentials_config, dict):
+                self._auth_manager = AgentAuthManager(config=credentials_config)
+        else:
+            self._auth_manager = AgentAuthManager()
+
+        self._build_interceptors(agent_cards)
+
+    def _build_interceptors(self, agent_cards: List[Any]):
+        if not self._auth_manager:
+            logger.info("[AuthManager] No auth manager, skipping interceptor build")
+            return
+        for card in agent_cards:
+            if not hasattr(card, "name"):
+                continue
+            interceptors = []
+            cred_svc = None
+            if card.security_schemes and card.security_requirements:
+                cred_svc = self._auth_manager.get_service(card.name)
+            else:
+                logger.info(f"[AuthManager] Agent {card.name}: no security schemes, skipping auth")
+                if not (getattr(card, "capabilities", None) and card.capabilities.extensions):
+                    continue
+            if cred_svc is not None:
+                logger.info(f"[AuthManager] Agent {card.name}: credentials found")
+                agent_cfg = self._auth_manager.get_config(card.name) or {}
+                if any(isinstance(v, dict) and (v.get("auth_header") or v.get("accept_header"))
+                       for v in agent_cfg.values()):
+                    interceptors.append(CustomAuthInterceptor(cred_svc, agent_cfg))
+                else:
+                    interceptors.append(AuthInterceptor(cred_svc))
+                logger.info(f"[AuthManager] Agent {card.name}: configured with {type(interceptors[0]).__name__}")
+            if getattr(card, "capabilities", None) and card.capabilities.extensions:
+                ext_uris = [ext.uri for ext in card.capabilities.extensions if ext.uri]
+                if ext_uris:
+                    interceptors.append(ExtensionInterceptor(ext_uris))
+            if interceptors:
+                self._interceptors[card.name] = interceptors
+
+    def get_interceptors(self, agent_name: str) -> List[Any]:
+        return self._interceptors.get(agent_name, [])
+
+    def set_httpx_client(self, client):
+        if self._auth_manager:
+            self._auth_manager.set_httpx_client(client)
+class AuthProviderInterceptor(ClientCallInterceptor if _A2A_AUTH_AVAILABLE else object):
+    """Wraps a custom AuthProvider as an a2a-sdk ClientCallInterceptor.
+
+    Calls ``auth_provider.apply_auth(agent_name, agent_card, headers)`` on
+    every ``before`` to inject auth headers. Mirrors the Java SDK's
+    ``AuthProvider.applyAuth`` being called in buildClientCallContext.
+    """
+    def __init__(self, auth_provider: AuthProvider, agent_name: str):
+        self._auth_provider = auth_provider
+        self._agent_name = agent_name
+
+    async def before(self, args: "BeforeArgs") -> None:
+        agent_card = args.agent_card
+        headers: Dict[str, str] = {}
+        try:
+            self._auth_provider.apply_auth(self._agent_name, agent_card, headers)
+        except Exception as e:
+            logger.warning(f"[AuthProvider] apply_auth raised: {e}")
+        if args.context is None:
+            from a2a.client.client import ClientCallContext
+            args.context = ClientCallContext()
+        if args.context.service_parameters is None:
+            args.context.service_parameters = {}
+        args.context.service_parameters.update(headers)
+        if headers:
+            logger.info(f"[AuthProvider] Injected {len(headers)} header(s) for {self._agent_name}")
+
+    async def after(self, args: "AfterArgs") -> None:
+        pass

workflow_engine/client/auth_provider.py

@@ -0,0 +1,47 @@
+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
+# All Rights Reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+"""Custom authentication provider for injecting auth headers.
+
+Implement this when the agent's authentication is not covered by the
+credentials file or the AgentCard's security schemes (e.g. corporate SSO,
+non-standard auth). Mirrors the Java SDK's ``AuthProvider`` interface.
+
+Register via ``WorkflowEngineClient(..., auth_provider=my_provider)``.
+The provider is called for every message send, regardless of whether the
+AgentCard declares security schemes. If both a credentials config and a
+custom AuthProvider are configured, both run (custom provider first,
+credentials-based auth second).
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+
+class AuthProvider(ABC):
+    """Apply authentication headers for sending a message to an agent."""
+
+    @abstractmethod
+    def apply_auth(self, agent_name: str, agent_card: Any, headers: Dict[str, str]) -> None:
+        """Add auth headers to the mutable ``headers`` dict.
+
+        Args:
+            agent_name: target agent name (matches AgentCard.name)
+            agent_card: the agent's card (security_schemes may be empty)
+            headers: mutable header map to add auth headers to
+        """
+        ...

workflow_engine/client/credential_crypto.py

@@ -0,0 +1,102 @@
+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
+# All Rights Reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+"""AES-GCM credential encryption/decryption utility.
+
+Supports encrypted values in credential config files using the ``enc:``
+prefix. The encryption key is read from the ``A2AT_CRED_KEY`` environment
+variable (32-byte hex string). Mirrors the Java SDK's ``CredentialCrypto``.
+
+Usage in credentials JSON::
+
+    {"value": "enc:<base64-iv>:<base64-ciphertext>"}
+
+Plaintext values (no ``enc:`` prefix) are returned as-is for backward compat.
+"""
+
+import os
+import base64
+import secrets
+from typing import Optional
+
+from loguru import logger
+
+_ENV_KEY = "A2AT_CRED_KEY"
+_PREFIX = "enc:"
+_IV_LENGTH = 12  # 96-bit IV for GCM
+_TAG_LENGTH = 16  # 128-bit auth tag
+
+
+def _resolve_key() -> Optional[str]:
+    """Resolve the encryption key from OS environment."""
+    key = os.environ.get(_ENV_KEY)
+    if key:
+        return key
+    return None
+
+
+def decrypt_if_needed(value: Optional[str]) -> Optional[str]:
+    """Decrypt a credential value if it has the ``enc:`` prefix.
+
+    Values without the prefix are returned as-is (plaintext fallback).
+    """
+    if not value or not value.startswith(_PREFIX):
+        return value
+    key_hex = _resolve_key()
+    if not key_hex or not key_hex.strip():
+        logger.warning(
+            f"[CredentialCrypto] Encrypted value found but {_ENV_KEY} not set, using as-is"
+        )
+        return value
+    try:
+        encoded = value[len(_PREFIX):]
+        parts = encoded.split(":", 1)
+        if len(parts) != 2:
+            logger.error("[CredentialCrypto] Invalid encrypted format, expected enc:<iv>:<ciphertext>")
+            return value
+        iv = base64.b64decode(parts[0])
+        ciphertext_and_tag = base64.b64decode(parts[1])
+        key_bytes = bytes.fromhex(key_hex)
+        ciphertext = ciphertext_and_tag[:-_TAG_LENGTH]
+        tag = ciphertext_and_tag[-_TAG_LENGTH:]
+        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+        plaintext = AESGCM(key_bytes).decrypt(iv, ciphertext + tag, None)
+        return plaintext.decode("utf-8")
+    except Exception as e:
+        logger.error(f"[CredentialCrypto] Decryption failed: {e}")
+        return value
+
+
+def encrypt(plaintext: str) -> str:
+    """Encrypt a plaintext value using AES-GCM with the key from A2AT_CRED_KEY.
+
+    Returns encrypted string in format ``enc:<base64-iv>:<base64-ciphertext>``.
+    Raises RuntimeError if the key env var is not set.
+    """
+    key_hex = _resolve_key()
+    if not key_hex or not key_hex.strip():
+        raise RuntimeError(f"{_ENV_KEY} environment variable not set")
+    key_bytes = bytes.fromhex(key_hex)
+    iv = secrets.token_bytes(_IV_LENGTH)
+    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+    ciphertext_and_tag = AESGCM(key_bytes).encrypt(iv, plaintext.encode("utf-8"), None)
+    return (
+        _PREFIX
+        + base64.b64encode(iv).decode("ascii")
+        + ":"
+        + base64.b64encode(ciphertext_and_tag).decode("ascii")
+    )

workflow_engine/client/credential_service.py

@@ -0,0 +1,229 @@
+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
+# All Rights Reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+"""Agent credential service - self-contained, no orchestration center dependency.
+
+Obtains Bearer tokens via login endpoints for agents requiring authentication.
+Reads credentials from a user-provided config (JSON file or dict).
+"""
+
+import json
+import time
+from pathlib import Path
+from typing import Dict, Optional
+import httpx
+from loguru import logger
+
+from workflow_engine.client.credential_crypto import decrypt_if_needed
+
+try:
+    from a2a.client.auth import CredentialService
+    from a2a.client.auth import InMemoryContextCredentialStore
+    from a2a.client.interceptors import ClientCallInterceptor, BeforeArgs, AfterArgs
+    _A2A_AVAILABLE = True
+except ImportError:
+    _A2A_AVAILABLE = False
+    CredentialService = object
+
+
+class AgentCredentialService(CredentialService if _A2A_AVAILABLE else object):
+    """Obtains tokens via login endpoint, caches with TTL."""
+
+    def __init__(self, agent_name: str, scheme_configs: Dict[str, dict],
+                 httpx_client: Optional[httpx.AsyncClient] = None):
+        self._agent_name = agent_name
+        self._schemes = scheme_configs
+        self._httpx_client = httpx_client
+        self._tokens: Dict[str, tuple] = {}
+        self._lock = None
+
+    def _ensure_lock(self):
+        if self._lock is None:
+            import asyncio
+            self._lock = asyncio.Lock()
+
+    def set_httpx_client(self, client: httpx.AsyncClient):
+        self._httpx_client = client
+
+    async def get_credentials(self, security_scheme_name: str, context=None) -> Optional[str]:
+        scheme_cfg = self._schemes.get(security_scheme_name)
+        if not scheme_cfg:
+            return None
+        cached = self._tokens.get(security_scheme_name)
+        if cached:
+            token, expires_at = cached
+            if time.time() < expires_at - 60:
+                logger.info(f"[Auth] Cache hit for agent {self._agent_name} scheme {security_scheme_name}")
+                return token
+        self._ensure_lock()
+        async with self._lock:
+            cached = self._tokens.get(security_scheme_name)
+            if cached:
+                token, expires_at = cached
+                if time.time() < expires_at - 60:
+                    return token
+            token = await self._login(scheme_cfg)
+            if token:
+                ttl = scheme_cfg.get("token_ttl", 3600)
+                self._tokens[security_scheme_name] = (token, time.time() + ttl)
+                logger.info(f"[Auth] Login succeeded: agent={self._agent_name}, scheme={security_scheme_name}")
+            return token
+
+    async def _login(self, scheme_cfg: dict) -> Optional[str]:
+        login_url = scheme_cfg.get("login_url")
+        if not login_url:
+            return None
+        method = scheme_cfg.get("method", "POST").upper()
+        content_type = scheme_cfg.get("content_type", "application/json")
+        token_field = scheme_cfg.get("token_field", "accessSession")
+        request_fields = scheme_cfg.get("request_fields")
+        if request_fields and isinstance(request_fields, dict):
+            body = {k: decrypt_if_needed(v) if isinstance(v, str) else v
+                    for k, v in request_fields.items()}
+        else:
+            username = scheme_cfg.get("username")
+            password = decrypt_if_needed(scheme_cfg.get("password"))
+            if not username or not password:
+                return None
+            body = {scheme_cfg.get("username_field","username"): username, scheme_cfg.get("password_field","password"): password}
+        client = self._httpx_client or httpx.AsyncClient(timeout=httpx.Timeout(connect=30, read=30, write=30, pool=5.0), verify=False)
+        own_client = self._httpx_client is None
+        try:
+            logger.info(f"[Auth] Login attempt: agent={self._agent_name}, method={method}, url={login_url}, content_type={content_type}, params={_sanitize_body(body)}")
+            req_kwargs = {"method": method, "url": login_url}
+            if content_type == "application/x-www-form-urlencoded":
+                req_kwargs["data"] = body
+            else:
+                req_kwargs["json"] = body
+            resp = await client.request(**req_kwargs)
+            resp.raise_for_status()
+            data = resp.json()
+            token = self._extract_nested_value(data, token_field) if isinstance(data, dict) else None
+            if not token and isinstance(data, dict):
+                token = data.get("accessSession") or data.get("access_session") or data.get("access_token") or data.get("token")
+            return token
+        except Exception as e:
+            logger.error(f"[Auth] Login failed: agent={self._agent_name}, url={login_url}, error={e}")
+            return None
+        finally:
+            if own_client:
+                await client.aclose()
+
+    @staticmethod
+    def _extract_nested_value(data: dict, path: str) -> Optional[str]:
+        if not path:
+            return None
+        current = data
+        for part in path.split("."):
+            if not isinstance(current, dict):
+                return None
+            current = current.get(part)
+            if current is None:
+                return None
+        return current
+
+
+class AgentAuthManager:
+    """Loads agent credentials from config, creates per-agent CredentialService."""
+
+    def __init__(self, config: Optional[Dict[str, dict]] = None, config_path: Optional[str] = None):
+        self._config: Dict[str, dict] = {}
+        self._services: Dict[str, AgentCredentialService] = {}
+        if config:
+            self._config = config
+        elif config_path:
+            self._load_from_file(config_path)
+
+    def _load_from_file(self, path: str):
+        p = Path(path)
+        if not p.exists():
+            return
+        try:
+            with open(p, "r", encoding="utf-8") as f:
+                self._config = json.load(f)
+            logger.info(f"[Auth] Loaded credentials for {len(self._config)} agent(s): {list(self._config.keys())}")
+        except Exception as e:
+            logger.warning(f"[Auth] Failed to load credentials: {e}")
+
+    def get_service(self, agent_name: str) -> Optional[AgentCredentialService]:
+        if agent_name in self._services:
+            return self._services[agent_name]
+        agent_creds = self._config.get(agent_name)
+        if not agent_creds:
+            return None
+        service = AgentCredentialService(agent_name, agent_creds)
+        self._services[agent_name] = service
+        logger.info(f"[Auth] Created credential service for agent: {agent_name}")
+        return service
+
+    def get_config(self, agent_name: str) -> Optional[Dict[str, dict]]:
+        return self._config.get(agent_name)
+
+    def set_httpx_client(self, client: httpx.AsyncClient):
+        for svc in self._services.values():
+            svc.set_httpx_client(client)
+
+
+class CustomAuthInterceptor(ClientCallInterceptor if _A2A_AVAILABLE else object):
+    """Auth interceptor supporting custom header names."""
+
+    def __init__(self, credential_service: AgentCredentialService, scheme_configs: Dict[str, dict]):
+        self._credential_service = credential_service
+        self._scheme_configs = scheme_configs
+
+    async def before(self, args: BeforeArgs) -> None:
+        agent_card = args.agent_card
+        if not agent_card.security_requirements or not agent_card.security_schemes:
+            return
+        for requirement in agent_card.security_requirements:
+            for scheme_name in requirement.schemes:
+                scheme_cfg = self._scheme_configs.get(scheme_name, {})
+                credential = await self._credential_service.get_credentials(scheme_name, args.context)
+                if not credential:
+                    continue
+                if args.context is None:
+                    from a2a.client.client import ClientCallContext
+                    args.context = ClientCallContext()
+                if args.context.service_parameters is None:
+                    args.context.service_parameters = {}
+                auth_header = scheme_cfg.get("auth_header")
+                if auth_header:
+                    prefix = scheme_cfg.get("auth_header_prefix", "")
+                    args.context.service_parameters[auth_header] = f"{prefix}{credential}"
+                    logger.info(f"[CustomAuth] Set header {auth_header} for scheme {scheme_name}")
+                else:
+                    args.context.service_parameters["Authorization"] = f"Bearer {credential}"
+                    logger.info(f"[CustomAuth] Set Bearer header for scheme {scheme_name}")
+                accept_header = scheme_cfg.get("accept_header")
+                if accept_header:
+                    args.context.service_parameters["Accept"] = accept_header
+                    logger.info(f"[CustomAuth] Override Accept header to {accept_header} for agent {getattr(args.agent_card, chr(39)+chr(110)+chr(97)+chr(109)+chr(101)+chr(39), chr(63))}")
+                return
+
+    async def after(self, args: AfterArgs) -> None:
+        pass
+
+
+def _sanitize_body(body: dict) -> dict:
+    """Mask sensitive fields (password, value, accessSession) for safe logging."""
+    sanitized = {}
+    for k, v in body.items():
+        if k.lower() in ("password", "value", "accesssession"):
+            sanitized[k] = "***"
+        else:
+            sanitized[k] = v
+    return sanitized