yunzerwebsiteallinone/go/pkg/tokenprobe/probe.go

package tokenprobe

import (
	"bytes"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"
)

var httpClient = &http.Client{
	Timeout: 12 * time.Second,
	Transport: &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	},
}

type Result struct {
	OK                   bool   `json:"ok"`
	Detail               string `json:"detail"`
	HTTPStatus           int    `json:"httpStatus"`
	ProbeMessage         string `json:"probeMessage,omitempty"`
	Endpoint             string `json:"endpoint,omitempty"`
	BytesRead            int    `json:"bytesRead,omitempty"`
	RawPreview           string `json:"rawPreview,omitempty"`
	RequestBodyPrefixHex string `json:"requestBodyPrefixHex,omitempty"`
	StreamProtocol       string `json:"streamProtocol,omitempty"`
	StreamNote           string `json:"streamNote,omitempty"`
}

func ProbeOfficial(module, rawToken string) Result {
	tok := normalizeBearerToken(strings.TrimSpace(rawToken))
	if tok == "" {
		return Result{OK: false, Detail: "Token 为空"}
	}
	switch module {
	case "cursor":
		return probeCursor(tok)
	case "windsurf":
		return probeWindsurf(tok)
	case "krio":
		return probeKiro(tok)
	default:
		return Result{OK: false, Detail: "未知模块"}
	}
}

func normalizeBearerToken(s string) string {
	s = strings.TrimSpace(s)
	if i := strings.LastIndex(s, "::"); i >= 0 {
		return strings.TrimSpace(s[i+2:])
	}
	return s
}

func probeCursor(token string) Result {
	return probeCursorHiAgent(token)
}

func probeWindsurf(apiKey string) Result {
	payload := map[string]interface{}{
		"metadata": map[string]string{
			"apiKey":           apiKey,
			"ideName":          "windsurf",
			"ideVersion":       "0.0.0",
			"extensionName":    "windsurf",
			"extensionVersion": "0.0.0",
			"locale":           "zh",
		},
	}
	raw, err := json.Marshal(payload)
	if err != nil {
		return Result{OK: false, Detail: err.Error()}
	}
	req, err := http.NewRequest(
		http.MethodPost,
		"https://server.codeium.com/exa.seat_management_pb.SeatManagementService/GetUserStatus",
		bytes.NewReader(raw),
	)
	if err != nil {
		return Result{OK: false, Detail: err.Error()}
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Connect-Protocol-Version", "1")

	resp, err := httpClient.Do(req)
	if err != nil {
		return Result{OK: false, Detail: "请求失败: " + err.Error()}
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(io.LimitReader(resp.Body, 8192))

	switch resp.StatusCode {
	case http.StatusOK:
		var wrap map[string]interface{}
		if json.Unmarshal(body, &wrap) == nil {
			if _, ok := wrap["userStatus"]; ok {
				return Result{OK: true, Detail: "Codeium 云端接口响应正常", HTTPStatus: resp.StatusCode}
			}
		}
		if bytes.Contains(body, []byte(`"planStatus"`)) || bytes.Contains(body, []byte(`"userStatus"`)) {
			return Result{OK: true, Detail: "Codeium 云端接口响应正常", HTTPStatus: resp.StatusCode}
		}
		return Result{OK: true, Detail: fmt.Sprintf("HTTP %d，已收到响应", resp.StatusCode), HTTPStatus: resp.StatusCode}
	case http.StatusUnauthorized, http.StatusForbidden:
		return Result{OK: false, Detail: fmt.Sprintf("API Key 无效或已失效（HTTP %d）", resp.StatusCode), HTTPStatus: resp.StatusCode}
	default:
		snip := strings.TrimSpace(string(body))
		if len(snip) > 220 {
			snip = snip[:220] + "…"
		}
		return Result{OK: false, Detail: fmt.Sprintf("HTTP %d %s", resp.StatusCode, snip), HTTPStatus: resp.StatusCode}
	}
}

func probeKiro(accessToken string) Result {
	arn := findProfileArnInJWT(accessToken)
	if arn == "" {
		return Result{
			OK:     false,
			Detail: "无法从 Token 中解析 profileArn，Kiro 暂无法自动探测",
		}
	}

	q := url.Values{}
	q.Set("origin", "AI_EDITOR")
	q.Set("profileArn", arn)
	q.Set("resourceType", "AGENTIC_REQUEST")
	u := "https://q.us-east-1.amazonaws.com/getUsageLimits?" + q.Encode()

	req, err := http.NewRequest(http.MethodGet, u, nil)
	if err != nil {
		return Result{OK: false, Detail: err.Error()}
	}
	req.Header.Set("Authorization", "Bearer "+normalizeBearerToken(accessToken))
	req.Header.Set("Accept", "application/json")

	resp, err := httpClient.Do(req)
	if err != nil {
		return Result{OK: false, Detail: "请求失败: " + err.Error()}
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))

	switch resp.StatusCode {
	case http.StatusOK:
		return Result{OK: true, Detail: "Kiro（AWS Q）用量接口响应正常", HTTPStatus: resp.StatusCode}
	case http.StatusUnauthorized, http.StatusForbidden:
		return Result{OK: false, Detail: fmt.Sprintf("Token 无效或已过期（HTTP %d）", resp.StatusCode), HTTPStatus: resp.StatusCode}
	default:
		snip := strings.TrimSpace(string(body))
		if len(snip) > 220 {
			snip = snip[:220] + "…"
		}
		return Result{OK: false, Detail: fmt.Sprintf("HTTP %d %s", resp.StatusCode, snip), HTTPStatus: resp.StatusCode}
	}
}

func decodeJWTPayloadMap(raw string) (map[string]interface{}, error) {
	tok := normalizeBearerToken(strings.TrimSpace(raw))
	parts := strings.Split(tok, ".")
	if len(parts) < 2 {
		return nil, fmt.Errorf("not a JWT")
	}
	b, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		return nil, err
	}
	var m map[string]interface{}
	if err := json.Unmarshal(b, &m); err != nil {
		return nil, err
	}
	return m, nil
}

func findProfileArnInJWT(raw string) string {
	m, err := decodeJWTPayloadMap(raw)
	if err != nil {
		return ""
	}
	return findProfileArnValue(m)
}

func findProfileArnValue(v interface{}) string {
	switch x := v.(type) {
	case map[string]interface{}:
		for k, val := range x {
			lk := strings.ToLower(k)
			if lk == "profilearn" || lk == "profile_arn" {
				if s, ok := val.(string); ok && strings.Contains(s, "arn:") {
					return s
				}
			}
		}
		for _, val := range x {
			if s := findProfileArnValue(val); s != "" {
				return s
			}
		}
	case []interface{}:
		for _, el := range x {
			if s := findProfileArnValue(el); s != "" {
				return s
			}
		}
	case string:
		if strings.Contains(x, "arn:aws:codewhisperer") && strings.Contains(x, ":profile/") {
			return x
		}
	}
	return ""
}