DataOps-platform/scripts/update_n8n_cohere_credential.py

"""
更新 n8n 中的 Cohere API Key 凭证

用于修复 "Forbidden" 错误，重新配置凭证
"""

import os
import sys
import json
from typing import Optional, Dict, Any

import requests
from loguru import logger

# 添加项目根目录到路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from app.config.config import BaseConfig

# 配置日志（避免 emoji 编码问题）
logger.remove()
logger.add(
    sys.stdout,
    level="INFO",
    format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {message}",
)


def update_cohere_credential(
    credential_id: str,
    api_url: Optional[str] = None,
    api_key: Optional[str] = None,
    cohere_api_key: str = "4pLcF0CGE7LeDmAudBQHdvAxGaKwNOKfxUGkHb5C",
    credential_name: Optional[str] = None,
) -> Dict[str, Any]:
    """
    更新 n8n 中的 Cohere 凭证

    Args:
        credential_id: 凭证 ID
        api_url: n8n API 地址
        api_key: n8n API Key
        cohere_api_key: Cohere API Key 值
        credential_name: 凭证名称（可选）

    Returns:
        更新结果
    """
    # 获取配置
    if api_url is None or api_key is None:
        config = BaseConfig()
        api_url = api_url or config.N8N_API_URL
        api_key = api_key or config.N8N_API_KEY

    base_url = api_url.rstrip("/")
    headers = {
        "X-N8N-API-KEY": api_key,
        "Content-Type": "application/json",
        "Accept": "application/json",
    }

    # 先获取现有凭证信息
    logger.info(f"获取凭证信息: {credential_id}")
    try:
        get_response = requests.get(
            f"{base_url}/api/v1/credentials/{credential_id}",
            headers=headers,
            timeout=30,
        )

        if get_response.status_code == 200:
            existing_credential = get_response.json()
            logger.info(f"现有凭证名称: {existing_credential.get('name')}")
            credential_name = credential_name or existing_credential.get("name", "Cohere API Key")
        else:
            logger.warning(f"无法获取现有凭证: {get_response.status_code}")
            credential_name = credential_name or "Cohere API Key"

    except Exception as e:
        logger.warning(f"获取凭证信息失败: {str(e)}")
        credential_name = credential_name or "Cohere API Key"

    # 准备更新数据
    # 注意: n8n 凭证更新可能需要特定的数据格式
    update_data = {
        "name": credential_name,
        "type": "cohereApi",
        "data": {
            "apiKey": cohere_api_key,
        },
    }

    logger.info(f"更新凭证: {credential_id}")
    logger.info(f"API Key (前8位): {cohere_api_key[:8]}...")

    try:
        # 尝试更新凭证
        response = requests.put(
            f"{base_url}/api/v1/credentials/{credential_id}",
            headers=headers,
            json=update_data,
            timeout=30,
        )

        logger.debug(f"响应状态码: {response.status_code}")
        logger.debug(f"响应内容: {response.text[:500]}")

        if response.status_code == 200:
            result = response.json()
            logger.success(f"凭证更新成功: {result.get('name')}")
            return {
                "success": True,
                "message": "凭证更新成功",
                "data": result,
            }
        elif response.status_code == 401:
            logger.error("API 认证失败，请检查 n8n API Key")
            return {
                "success": False,
                "message": "API 认证失败，请检查 n8n API Key",
                "error": "Unauthorized",
            }
        elif response.status_code == 403:
            logger.error("API 权限不足，凭证更新可能需要 Owner 权限")
            return {
                "success": False,
                "message": "API 权限不足，请使用 Web UI 手动更新",
                "error": "Forbidden",
            }
        elif response.status_code == 404:
            logger.error(f"凭证不存在: {credential_id}")
            return {
                "success": False,
                "message": f"凭证不存在: {credential_id}",
                "error": "Not Found",
            }
        else:
            logger.warning(f"更新失败: {response.status_code} - {response.text[:200]}")
            return {
                "success": False,
                "message": f"更新失败: {response.status_code}",
                "error": response.text[:200],
                "status_code": response.status_code,
            }

    except requests.exceptions.RequestException as e:
        logger.error(f"请求异常: {str(e)}")
        return {
            "success": False,
            "message": f"请求失败: {str(e)}",
            "error": str(e),
        }


def delete_and_recreate_credential(
    credential_id: Optional[str] = None,
    api_url: Optional[str] = None,
    api_key: Optional[str] = None,
    cohere_api_key: str = "4pLcF0CGE7LeDmAudBQHdvAxGaKwNOKfxUGkHb5C",
) -> Dict[str, Any]:
    """
    删除旧凭证并重新创建

    Args:
        credential_id: 要删除的凭证 ID（如果为 None，则只创建新凭证）
        api_url: n8n API 地址
        api_key: n8n API Key
        cohere_api_key: Cohere API Key 值

    Returns:
        操作结果
    """
    # 获取配置
    if api_url is None or api_key is None:
        config = BaseConfig()
        api_url = api_url or config.N8N_API_URL
        api_key = api_key or config.N8N_API_KEY

    base_url = api_url.rstrip("/")
    headers = {
        "X-N8N-API-KEY": api_key,
        "Content-Type": "application/json",
        "Accept": "application/json",
    }

    # 删除旧凭证（如果提供）
    if credential_id:
        logger.info(f"删除旧凭证: {credential_id}")
        try:
            delete_response = requests.delete(
                f"{base_url}/api/v1/credentials/{credential_id}",
                headers=headers,
                timeout=30,
            )
            if delete_response.status_code in [200, 204]:
                logger.success("旧凭证已删除")
            else:
                logger.warning(f"删除凭证失败: {delete_response.status_code}")
        except Exception as e:
            logger.warning(f"删除凭证异常: {str(e)}")

    # 创建新凭证
    logger.info("创建新凭证...")
    credential_data = {
        "name": "Cohere API Key",
        "type": "cohereApi",
        "data": {
            "apiKey": cohere_api_key,
        },
    }

    try:
        response = requests.post(
            f"{base_url}/api/v1/credentials",
            headers=headers,
            json=credential_data,
            timeout=30,
        )

        if response.status_code in [200, 201]:
            result = response.json()
            logger.success(f"新凭证创建成功: {result.get('id')}")
            return {
                "success": True,
                "message": "凭证重新创建成功",
                "data": result,
                "credential_id": result.get("id"),
            }
        else:
            logger.error(f"创建凭证失败: {response.status_code} - {response.text[:200]}")
            return {
                "success": False,
                "message": f"创建凭证失败: {response.status_code}",
                "error": response.text[:200],
            }

    except requests.exceptions.RequestException as e:
        logger.error(f"请求异常: {str(e)}")
        return {
            "success": False,
            "message": f"请求失败: {str(e)}",
            "error": str(e),
        }


def main():
    """主函数"""
    import argparse

    parser = argparse.ArgumentParser(description="更新 n8n Cohere API Key 凭证")
    parser.add_argument(
        "--credential-id",
        type=str,
        help="要更新的凭证 ID（如果不提供，将创建新凭证）",
    )
    parser.add_argument(
        "--recreate",
        action="store_true",
        help="删除旧凭证并重新创建",
    )
    parser.add_argument(
        "--api-key",
        type=str,
        help="Cohere API Key（默认使用配置中的值）",
        default="4pLcF0CGE7LeDmAudBQHdvAxGaKwNOKfxUGkHb5C",
    )

    args = parser.parse_args()

    logger.info("开始更新 n8n Cohere API Key 凭证...")

    if args.recreate:
        result = delete_and_recreate_credential(
            credential_id=args.credential_id,
            cohere_api_key=args.api_key,
        )
    elif args.credential_id:
        result = update_cohere_credential(
            credential_id=args.credential_id,
            cohere_api_key=args.api_key,
        )
    else:
        logger.error("请提供 --credential-id 或使用 --recreate 选项")
        return 1

    print("\n" + "=" * 60)
    print("执行结果:")
    print("=" * 60)
    print(json.dumps(result, indent=2, ensure_ascii=False))

    if result.get("success"):
        if "credential_id" in result:
            print(f"\n新凭证 ID: {result.get('credential_id')}")
        print("\n请在工作流中使用更新后的凭证。")
    else:
        print("\n" + "=" * 60)
        print("建议:")
        print("=" * 60)
        print("如果 API 更新失败，请使用 Web UI 手动更新:")
        print("1. 访问: https://n8n.citupro.com/home/credentials")
        print("2. 找到 Cohere API Key 凭证")
        print("3. 点击编辑，重新输入 API Key")
        print("4. 确保 API Key 没有多余的空格或换行")

    return 0 if result.get("success") else 1


if __name__ == "__main__":
    sys.exit(main())