new-api/service/asset_tianyiyun.go

package service

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	"github.com/QuantumNous/new-api/setting/system_setting"
)

type TianyiYunAssetAdapter struct{}

func NewTianyiYunAssetAdapter() AssetAdapter {
	return &TianyiYunAssetAdapter{}
}

func (a *TianyiYunAssetAdapter) Name() string {
	return "tianyiyun_asset"
}

func (a *TianyiYunAssetAdapter) Supports(operation AssetOperation) bool {
	return operation == AssetOperationAssetCreate || operation == AssetOperationAssetGet
}

func (a *TianyiYunAssetAdapter) DoAssetRequest(ctx context.Context, channel *model.Channel, req AssetRequest) (*AssetUpstreamResponse, *AssetError) {
	if !a.Supports(req.Action.Operation) {
		return nil, newAssetError(AssetErrorOperationNotSupported, fmt.Sprintf("asset operation %s is not supported", req.Action.Operation), http.StatusBadRequest)
	}
	if channel == nil {
		return nil, newAssetError(AssetErrorInvalidRequest, "channel is required", http.StatusBadRequest)
	}

	upstreamURL, body, method, assetErr := buildTianyiYunAssetRequest(channel, req)
	if assetErr != nil {
		return nil, assetErr
	}
	fetchSetting := system_setting.GetFetchSetting()
	if err := common.ValidateURLWithFetchSetting(upstreamURL, fetchSetting.EnableSSRFProtection, fetchSetting.AllowPrivateIp, fetchSetting.DomainFilterMode, fetchSetting.IpFilterMode, fetchSetting.DomainList, fetchSetting.IpList, fetchSetting.AllowedPorts, fetchSetting.ApplyIPFilterForDomain); err != nil {
		return nil, newAssetError(AssetErrorServer, fmt.Sprintf("request blocked: %v", err), http.StatusForbidden)
	}

	httpReq, err := http.NewRequestWithContext(ctx, method, upstreamURL, bytes.NewReader(body))
	if err != nil {
		return nil, newAssetError(AssetErrorServer, err.Error(), http.StatusInternalServerError)
	}
	httpReq.Header.Set("Accept", "application/json")
	httpReq.Header.Set("Authorization", "Bearer "+strings.TrimSpace(channel.Key))
	if method == http.MethodPost {
		httpReq.Header.Set("Content-Type", "application/json")
	}
	client, err := GetHttpClientWithProxy(channel.GetSetting().Proxy)
	if err != nil {
		return nil, newAssetError(AssetErrorServer, err.Error(), http.StatusInternalServerError)
	}
	if client == nil {
		client = http.DefaultClient
	}
	resp, err := client.Do(httpReq)
	if err != nil {
		return nil, newAssetError(AssetErrorUpstream, err.Error(), http.StatusBadGateway)
	}
	defer resp.Body.Close()
	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, newAssetError(AssetErrorUpstream, err.Error(), http.StatusBadGateway)
	}
	if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
		return nil, newAssetError(AssetErrorUpstream, string(data), http.StatusBadGateway)
	}
	return &AssetUpstreamResponse{StatusCode: resp.StatusCode, Header: resp.Header, Body: data}, nil
}

func buildTianyiYunAssetRequest(channel *model.Channel, req AssetRequest) (string, []byte, string, *AssetError) {
	baseURL, err := tianyiYunAssetBaseURL(channel)
	if err != nil {
		return "", nil, "", newAssetError(AssetErrorInvalidRequest, err.Error(), http.StatusBadRequest)
	}
	switch req.Action.Operation {
	case AssetOperationAssetCreate:
		payload, assetErr := buildTianyiYunAssetCreatePayload(req.Body)
		if assetErr != nil {
			return "", nil, "", assetErr
		}
		data, err := common.Marshal(payload)
		if err != nil {
			return "", nil, "", newAssetError(AssetErrorServer, err.Error(), http.StatusInternalServerError)
		}
		return baseURL + "/api/assets/upload", data, http.MethodPost, nil
	case AssetOperationAssetGet:
		id := tianyiYunAssetField(req.Body, "Id", "id")
		if id == "" {
			return "", nil, "", newAssetError(AssetErrorInvalidRequest, "asset id is required", http.StatusBadRequest)
		}
		return baseURL + "/api/assets/" + url.PathEscape(id), nil, http.MethodGet, nil
	default:
		return "", nil, "", newAssetError(AssetErrorOperationNotSupported, fmt.Sprintf("asset operation %s is not supported", req.Action.Operation), http.StatusBadRequest)
	}
}

func tianyiYunAssetBaseURL(channel *model.Channel) (string, error) {
	baseURL := strings.TrimSpace(channel.GetBaseURL())
	if baseURL == "" {
		baseURL = "https://ai.ctaigw.cn/v1"
	}
	u, err := url.Parse(baseURL)
	if err != nil || u.Scheme == "" || u.Host == "" {
		return "", fmt.Errorf("invalid TianyiYun asset base URL")
	}
	u.Path = strings.TrimRight(u.Path, "/")
	if !strings.HasSuffix(u.Path, "/v1") {
		u.Path += "/v1"
	}
	u.RawQuery = ""
	u.Fragment = ""
	return strings.TrimRight(u.String(), "/"), nil
}

func buildTianyiYunAssetCreatePayload(body map[string]any) (map[string]string, *AssetError) {
	sourceURL := tianyiYunAssetField(body, "URL", "url")
	if sourceURL == "" {
		return nil, newAssetError(AssetErrorInvalidRequest, "asset URL is required", http.StatusBadRequest)
	}
	parsedURL, err := url.Parse(sourceURL)
	if err != nil || (parsedURL.Scheme != "http" && parsedURL.Scheme != "https") || parsedURL.Host == "" {
		return nil, newAssetError(AssetErrorInvalidRequest, "asset URL must use http or https", http.StatusBadRequest)
	}
	assetType := tianyiYunAssetField(body, "AssetType", "asset_type")
	if assetType != "Image" && assetType != "Video" && assetType != "Audio" {
		return nil, newAssetError(AssetErrorInvalidRequest, "asset type must be Image, Video, or Audio", http.StatusBadRequest)
	}
	fetchSetting := system_setting.GetFetchSetting()
	if err := common.ValidateURLWithFetchSetting(sourceURL, fetchSetting.EnableSSRFProtection, fetchSetting.AllowPrivateIp, fetchSetting.DomainFilterMode, fetchSetting.IpFilterMode, fetchSetting.DomainList, fetchSetting.IpList, fetchSetting.AllowedPorts, fetchSetting.ApplyIPFilterForDomain); err != nil {
		return nil, newAssetError(AssetErrorInvalidRequest, fmt.Sprintf("asset URL blocked: %v", err), http.StatusBadRequest)
	}
	payload := map[string]string{"url": sourceURL, "asset_type": assetType}
	if name := tianyiYunAssetField(body, "Name", "name"); name != "" {
		payload["name"] = name
	}
	return payload, nil
}

func tianyiYunAssetField(body map[string]any, upperKey string, lowerKey string) string {
	if body == nil {
		return ""
	}
	if value, ok := body[upperKey]; ok {
		if text, ok := value.(string); ok {
			return strings.TrimSpace(text)
		}
	}
	if value, ok := body[lowerKey]; ok {
		if text, ok := value.(string); ok {
			return strings.TrimSpace(text)
		}
	}
	return ""
}