Cyrene/backend/ai-core/internal/tools/crypto_tool.go

package tools

import (
	"context"
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/base64"
	"fmt"
	"hash"
	"net/url"
)

// CryptoTool provides cryptographic and encoding utilities for the LLM.
// Supports hashing, base64, and URL encoding.
type CryptoTool struct{}

// NewCryptoTool creates a crypto/encoding tool.
func NewCryptoTool() *CryptoTool {
	return &CryptoTool{}
}

// Definition returns the tool definition for LLM function calling.
func (t *CryptoTool) Definition() ToolDefinition {
	return ToolDefinition{
		Name:        "crypto",
		Description: "加密哈希与编码工具。计算MD5/SHA哈希值，执行Base64编码/解码，URL编码/解码。",
		Parameters: map[string]interface{}{
			"type": "object",
			"properties": map[string]interface{}{
				"action": map[string]interface{}{
					"type":        "string",
					"enum":        []string{"hash", "base64_encode", "base64_decode", "url_encode", "url_decode"},
					"description": "操作类型。hash: 计算哈希值；base64_encode: Base64编码；base64_decode: Base64解码；url_encode: URL编码；url_decode: URL解码",
				},
				"input": map[string]interface{}{
					"type":        "string",
					"description": "输入数据，需要处理的字符串",
				},
				"algorithm": map[string]interface{}{
					"type":        "string",
					"enum":        []string{"md5", "sha1", "sha256", "sha512"},
					"description": "哈希算法（用于 hash 操作），默认 sha256",
				},
			},
			"required": []string{"action", "input"},
		},
	}
}

// Execute performs crypto/encoding operations.
func (t *CryptoTool) Execute(ctx context.Context, arguments map[string]interface{}) (*ToolResult, error) {
	action, ok := arguments["action"].(string)
	if !ok || action == "" {
		return &ToolResult{
			ToolName: "crypto",
			Success:  false,
			Error:    "缺少 action 参数",
		}, nil
	}

	input, ok := arguments["input"].(string)
	if !ok {
		return &ToolResult{
			ToolName: "crypto",
			Success:  false,
			Error:    "缺少 input 参数",
		}, nil
	}

	switch action {
	case "hash":
		return t.handleHash(arguments)
	case "base64_encode":
		return t.handleBase64Encode(input)
	case "base64_decode":
		return t.handleBase64Decode(input)
	case "url_encode":
		return t.handleURLEncode(input)
	case "url_decode":
		return t.handleURLDecode(input)
	default:
		return &ToolResult{
			ToolName: "crypto",
			Success:  false,
			Error:    fmt.Sprintf("未知操作: %s，支持: hash, base64_encode, base64_decode, url_encode, url_decode", action),
		}, nil
	}
}

// handleHash computes a hash of the input using the specified algorithm.
func (t *CryptoTool) handleHash(arguments map[string]interface{}) (*ToolResult, error) {
	input, _ := arguments["input"].(string)
	algorithm, _ := arguments["algorithm"].(string)
	if algorithm == "" {
		algorithm = "sha256"
	}

	var h hash.Hash
	switch algorithm {
	case "md5":
		h = md5.New()
	case "sha1":
		h = sha1.New()
	case "sha256":
		h = sha256.New()
	case "sha512":
		h = sha512.New()
	default:
		return &ToolResult{
			ToolName: "crypto",
			Success:  false,
			Error:    fmt.Sprintf("不支持的哈希算法: %s，支持: md5, sha1, sha256, sha512", algorithm),
		}, nil
	}

	h.Write([]byte(input))
	hashBytes := h.Sum(nil)
	hashHex := fmt.Sprintf("%x", hashBytes)

	return &ToolResult{
		ToolName: "crypto",
		Success:  true,
		Data: fmt.Sprintf("哈希算法: %s\n输入长度: %d 字节\n哈希值 (hex): %s\n哈希长度: %d 位",
			algorithm, len(input), hashHex, len(hashBytes)*8),
	}, nil
}

// handleBase64Encode encodes input to Base64.
func (t *CryptoTool) handleBase64Encode(input string) (*ToolResult, error) {
	encoded := base64.StdEncoding.EncodeToString([]byte(input))

	return &ToolResult{
		ToolName: "crypto",
		Success:  true,
		Data: fmt.Sprintf("Base64 编码结果:\n原始 (%d 字节): %s\n编码 (%d 字符): %s",
			len(input), truncate(input, 100), len(encoded), encoded),
	}, nil
}

// handleBase64Decode decodes a Base64 string.
func (t *CryptoTool) handleBase64Decode(input string) (*ToolResult, error) {
	// Try standard encoding first, then URL-safe
	decoded, err := base64.StdEncoding.DecodeString(input)
	if err != nil {
		decoded, err = base64.RawStdEncoding.DecodeString(input)
		if err != nil {
			decoded, err = base64.URLEncoding.DecodeString(input)
			if err != nil {
				decoded, err = base64.RawURLEncoding.DecodeString(input)
				if err != nil {
					return &ToolResult{
						ToolName: "crypto",
						Success:  false,
						Error:    fmt.Sprintf("Base64 解码失败: 输入不是有效的 Base64 字符串"),
					}, nil
				}
			}
		}
	}

	return &ToolResult{
		ToolName: "crypto",
		Success:  true,
		Data: fmt.Sprintf("Base64 解码结果:\n原始 (%d 字符): %s\n解码 (%d 字节): %s",
			len(input), truncate(input, 100), len(decoded), truncate(string(decoded), 200)),
	}, nil
}

// handleURLEncode URL-encodes the input string.
func (t *CryptoTool) handleURLEncode(input string) (*ToolResult, error) {
	encoded := url.QueryEscape(input)

	return &ToolResult{
		ToolName: "crypto",
		Success:  true,
		Data: fmt.Sprintf("URL 编码结果:\n原始 (%d 字节): %s\n编码 (%d 字节): %s",
			len(input), truncate(input, 100), len(encoded), encoded),
	}, nil
}

// handleURLDecode URL-decodes the input string.
func (t *CryptoTool) handleURLDecode(input string) (*ToolResult, error) {
	decoded, err := url.QueryUnescape(input)
	if err != nil {
		return &ToolResult{
			ToolName: "crypto",
			Success:  false,
			Error:    fmt.Sprintf("URL 解码失败: %v", err),
		}, nil
	}

	return &ToolResult{
		ToolName: "crypto",
		Success:  true,
		Data: fmt.Sprintf("URL 解码结果:\n原始 (%d 字节): %s\n解码 (%d 字节): %s",
			len(input), truncate(input, 100), len(decoded), truncate(decoded, 200)),
	}, nil
}

// truncate truncates a string to maxLen characters, adding "..." if truncated.
func truncate(s string, maxLen int) string {
	runes := []rune(s)
	if len(runes) <= maxLen {
		return s
	}
	return string(runes[:maxLen]) + "..."
}