223 lines
6.6 KiB
Go
223 lines
6.6 KiB
Go
// 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 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":
|
||
// 直接使用 cursor_hi.go 中已有的完整探测函数
|
||
return probeCursorHiAgent(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
|
||
}
|
||
|
||
// probeCursor Cursor Token 探测（直接使用 cursor_hi.go 的实现）
|
||
func probeCursor(token string) Result {
|
||
return probeCursorHiAgent(token)
|
||
}
|
||
|
||
// probeWindsurf WindSurf 探测
|
||
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}
|
||
}
|
||
}
|
||
|
||
// probeKiro Kiro 探测
|
||
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}
|
||
}
|
||
}
|
||
|
||
// decodeJWTPayloadMap 解析 JWT payload
|
||
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
|
||
}
|
||
|
||
// findProfileArnInJWT 从 JWT 中查找 profileArn
|
||
func findProfileArnInJWT(raw string) string {
|
||
m, err := decodeJWTPayloadMap(raw)
|
||
if err != nil {
|
||
return ""
|
||
}
|
||
return findProfileArnValue(m)
|
||
}
|
||
|
||
// findProfileArnValue 递归查找 profileArn
|
||
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 ""
|
||
}
|