first commit

2026-06-03 10:09:03 +08:00
commit 81f5039458
6266 changed files with 188236 additions and 0 deletions
@@ -0,0 +1,216 @@
+// Package tokenprobe 使用号池内 Token 调用各厂商接口做可用性探测。
+// Cursor 走 api2.cursor.sh 的 Connect + protobuf 二进制流（非 JSON 文本接口）。
+package tokenprobe
+
+import (
+	"bytes"
+	"encoding/base64"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+	"strings"
+	"time"
+)
+
+var httpClient = &http.Client{Timeout: 25 * time.Second}
+
+// Result 探测结果（Cursor 会填充 ProbeMessage / Endpoint / BytesRead / RawPreview 等）
+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 / StreamNote 仅 Cursor Agent 探测填充，说明二进制流与结论边界
+	StreamProtocol string `json:"streamProtocol,omitempty"`
+	StreamNote     string `json:"streamNote,omitempty"`
+}
+
+// ProbeOfficial 按号池模块探测 Token（cursor / windsurf / krio）
+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 暂无法自动探测（需完整登录 JWT）",
+		}
+	}
+
+	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 ""
+}