zhanlu_proxy/internal/zhanlu/client.go

package zhanlu

import (
	"bytes"
	"context"
	"crypto/rand"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"git.misaka.ren/M1saka/zhanlu_proxy/internal/auth"
	"git.misaka.ren/M1saka/zhanlu_proxy/internal/sign"
	"git.misaka.ren/M1saka/zhanlu_proxy/internal/util"
)

type Client struct {
	LoginBaseURL  string
	ModelBaseURL  string
	ChatPath      string
	PluginVersion string
	SM2PrivateKey string
	Signer        sign.Signer
	HTTPClient    *http.Client
}

func NewClient(loginBaseURL, modelBaseURL, chatPath, pluginVersion, sm2PrivateKey string, signer sign.Signer, timeout time.Duration) *Client {
	// The Zhanlu gateway rejects HTTP/2 ALPN negotiation (connection EOF on
	// handshake), so force HTTP/1.1 for all upstream requests.
	transport := &http.Transport{ForceAttemptHTTP2: false}
	return &Client{
		LoginBaseURL:  loginBaseURL,
		ModelBaseURL:  modelBaseURL,
		ChatPath:      chatPath,
		PluginVersion: pluginVersion,
		SM2PrivateKey: sm2PrivateKey,
		Signer:        signer,
		HTTPClient:    &http.Client{Timeout: timeout, Transport: transport},
	}
}

// LoginProfile validates AK/SK/token against the Zhanlu gateway
// (POST /api/acepilot/zhanlu/v1/login) and returns the decrypted user profile.
func (c *Client) LoginProfile(ctx context.Context, creds auth.Credentials) (auth.Profile, error) {
	signedURL, err := c.Signer.BuildOpURL("/api/acepilot/zhanlu/v1/login", creds, c.LoginBaseURL, http.MethodPost)
	if err != nil {
		return auth.Profile{}, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, signedURL, strings.NewReader(`{}`))
	if err != nil {
		return auth.Profile{}, err
	}
	c.setPluginHeaders(req)
	resp, err := c.HTTPClient.Do(req)
	if err != nil {
		return auth.Profile{}, err
	}
	defer resp.Body.Close()
	var out struct {
		State string `json:"state"`
		Body  struct {
			Email        string `json:"email"`
			Organization string `json:"organization"`
			Team         string `json:"team"`
			Name         string `json:"name"`
			Telephone    string `json:"telephone"`
		} `json:"body"`
	}
	if err := decodeJSON(resp, &out); err != nil {
		return auth.Profile{}, err
	}
	if out.State != "OK" {
		return auth.Profile{}, fmt.Errorf("zhanlu login failed: state=%s", out.State)
	}
	return auth.Profile{
		Email:        auth.DecryptCredentialOrRaw(strings.TrimSpace(out.Body.Email), creds.Token),
		Organization: auth.DecryptCredentialOrRaw(strings.TrimSpace(out.Body.Organization), creds.Token),
		Team:         auth.DecryptCredentialOrRaw(strings.TrimSpace(out.Body.Team), creds.Token),
		UserName:     auth.DecryptCredentialOrRaw(strings.TrimSpace(out.Body.Name), creds.Token),
		Telephone:    auth.DecryptCredentialOrRaw(strings.TrimSpace(out.Body.Telephone), creds.Token),
	}, nil
}

// ProvisionAPIKey requests a Zhanlu gateway API key
// (POST {modelBaseUrl}/user/api/v2/external/key/get-or-create) using the
// SM2-signed X-Auth-* headers, then returns the apiKey. Empty organization and
// team default to the plugin's "未配置" placeholder.
func (c *Client) ProvisionAPIKey(ctx context.Context, email, organization, team string) (string, error) {
	if strings.TrimSpace(c.SM2PrivateKey) == "" {
		return "", fmt.Errorf("ZHANLU_APIKEY_AUTH_SM2_PRIVATE_KEY is required to provision an API key")
	}
	if strings.TrimSpace(email) == "" {
		return "", fmt.Errorf("profile email is required to provision an API key")
	}
	org := util.FirstNonEmpty(organization, "未配置")
	tm := util.FirstNonEmpty(team, "未配置")
	// Field order matters: the SM2 signature covers the exact JSON body bytes,
	// matching the plugin's JSON.stringify({email, organization, team}).
	body, err := json.Marshal(struct {
		Email        string `json:"email"`
		Organization string `json:"organization"`
		Team         string `json:"team"`
	}{email, org, tm})
	if err != nil {
		return "", err
	}
	endpoint := strings.TrimRight(c.ModelBaseURL, "/") + "/user/api/v2/external/key/get-or-create"

	timestamp := fmt.Sprintf("%d", time.Now().Unix())
	nonce := randomAlnum(32)
	signature, err := auth.SignSM2Authorization(c.SM2PrivateKey, timestamp+":"+nonce+":"+string(body))
	if err != nil {
		return "", err
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-Auth-Signature", signature)
	req.Header.Set("X-Auth-Timestamp", timestamp)
	req.Header.Set("X-Auth-Nonce", nonce)

	resp, err := c.HTTPClient.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		b, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
		return "", fmt.Errorf("zhanlu api key provisioning returned %d: %s", resp.StatusCode, string(b))
	}
	var out map[string]any
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return "", err
	}
	apiKey := findString(out, "apiKey", "key", "api_key")
	if apiKey == "" {
		return "", fmt.Errorf("zhanlu api key provisioning response missing apiKey")
	}
	return apiKey, nil
}

// ChatCompletions posts an OpenAI-compatible body to the Zhanlu gateway chat
// endpoint authenticated with the provisioned API key.
func (c *Client) ChatCompletions(ctx context.Context, apiKey string, body []byte) (*http.Response, error) {
	baseURL := c.ModelBaseURL
	path := c.ChatPath
	if path == "" {
		path = "/chat/completions"
	}
	endpoint := strings.TrimRight(baseURL, "/") + path
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "text/event-stream, application/json")
	req.Header.Set("Authorization", "Bearer "+apiKey)
	return c.HTTPClient.Do(req)
}

// Models lists model ids exposed by the gateway model-info endpoint
// (GET {modelBaseUrl}/gateway/v1/model/info) with the provisioned API key.
func (c *Client) Models(ctx context.Context, apiKey string) ([]string, error) {
	endpoint := strings.TrimRight(c.ModelBaseURL, "/") + "/gateway/v1/model/info"
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Authorization", "Bearer "+apiKey)
	resp, err := c.HTTPClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		b, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
		return nil, fmt.Errorf("zhanlu model info returned %d: %s", resp.StatusCode, string(b))
	}
	var out struct {
		Data []struct {
			ModelName string `json:"model_name"`
			ID        string `json:"id"`
			ModelInfo struct {
				ID string `json:"id"`
			} `json:"model_info"`
		} `json:"data"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return nil, err
	}
	seen := map[string]bool{}
	models := make([]string, 0, len(out.Data))
	for _, m := range out.Data {
		id := util.FirstNonEmpty(m.ModelName, m.ID, m.ModelInfo.ID)
		if id == "" || seen[id] {
			continue
		}
		seen[id] = true
		models = append(models, id)
	}
	return models, nil
}

func (c *Client) setPluginHeaders(req *http.Request) {
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("plugin_type", "zhanlu_ide")
	req.Header.Set("plugin_version", c.PluginVersion)
	req.Header.Set("request", randomRequestID())
}

func decodeJSON(resp *http.Response, out any) error {
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		b, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
		return fmt.Errorf("zhanlu upstream returned %d: %s", resp.StatusCode, string(b))
	}
	return json.NewDecoder(resp.Body).Decode(out)
}

func findString(m map[string]any, keys ...string) string {
	for _, key := range keys {
		if v, ok := m[key].(string); ok && strings.TrimSpace(v) != "" {
			return strings.TrimSpace(v)
		}
	}
	// The plugin searches the root and the data/body/result/payload containers.
	for _, container := range []string{"data", "body", "result", "payload"} {
		if v, ok := m[container].(map[string]any); ok {
			if s := findString(v, keys...); s != "" {
				return s
			}
		}
	}
	return ""
}

const alnum = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"

func randomAlnum(n int) string {
	b := make([]byte, n)
	rand.Read(b)
	for i := range b {
		b[i] = alnum[int(b[i])%len(alnum)]
	}
	return string(b)
}

func randomRequestID() string {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		return fmt.Sprintf("%d", time.Now().UnixNano())
	}
	b[6] = (b[6] & 0x0f) | 0x40
	b[8] = (b[8] & 0x3f) | 0x80
	return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:])
}