增加token检测
This commit is contained in:
parent
2a60d34711
commit
283a2b7a80
@ -4,12 +4,14 @@ import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"server/models"
|
||||
"server/pkg/jwtutil"
|
||||
"server/pkg/tokenprobe"
|
||||
|
||||
"github.com/beego/beego/v2/client/orm"
|
||||
beego "github.com/beego/beego/v2/server/web"
|
||||
@ -154,6 +156,15 @@ func listPoolRows(c *beego.Controller, module string) {
|
||||
status := strings.TrimSpace(c.GetString("status"))
|
||||
|
||||
where, whereArgs := accountPoolListWhere(dataType, status, keyword)
|
||||
if module == "cursor" {
|
||||
u := strings.TrimSpace(c.GetString("usable"))
|
||||
if u == "1" || u == "0" {
|
||||
if v, err := strconv.ParseInt(u, 10, 8); err == nil {
|
||||
where = "(" + where + ") AND is_used = ?"
|
||||
whereArgs = append(whereArgs, int8(v))
|
||||
}
|
||||
}
|
||||
}
|
||||
offset := (page - 1) * pageSize
|
||||
|
||||
var list interface{}
|
||||
@ -696,32 +707,161 @@ func updatePoolRemark(c *beego.Controller, module string) {
|
||||
_ = c.ServeJSON()
|
||||
}
|
||||
|
||||
func (c *PlatformAccountPoolCursorController) List() { listPoolRows(&c.Controller, "cursor") }
|
||||
func (c *PlatformAccountPoolCursorController) Add() { addPoolRow(&c.Controller, "cursor") }
|
||||
func (c *PlatformAccountPoolCursorController) BatchAdd() { batchAddPoolRows(&c.Controller, "cursor") }
|
||||
func (c *PlatformAccountPoolCursorController) Detail() { getPoolDetail(&c.Controller, "cursor") }
|
||||
func (c *PlatformAccountPoolCursorController) Extract() { extractPoolRow(&c.Controller, "cursor") }
|
||||
func probePoolToken(c *beego.Controller, module string) {
|
||||
if _, err := requirePlatformAuth(c); err != nil {
|
||||
poolJSONErr(c, 401, 401, err.Error())
|
||||
return
|
||||
}
|
||||
raw, err := io.ReadAll(c.Ctx.Request.Body)
|
||||
if err != nil {
|
||||
poolJSONErr(c, 400, 400, "参数错误")
|
||||
return
|
||||
}
|
||||
var payload struct {
|
||||
ID uint64 `json:"id"`
|
||||
AccessToken string `json:"accessToken"`
|
||||
Token string `json:"token"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &payload); err != nil {
|
||||
poolJSONErr(c, 400, 400, "参数错误")
|
||||
return
|
||||
}
|
||||
|
||||
var token string
|
||||
switch module {
|
||||
case "cursor":
|
||||
token = strings.TrimSpace(payload.AccessToken)
|
||||
if token == "" {
|
||||
token = strings.TrimSpace(payload.Token)
|
||||
}
|
||||
if token == "" {
|
||||
if payload.ID == 0 {
|
||||
poolJSONErr(c, 400, 400, "请传入 Cursor 的 accessToken（会话 JWT），或传 id 从库中读取")
|
||||
return
|
||||
}
|
||||
var row models.PlatformAccountPoolCursor
|
||||
if err := models.Orm.QueryTable(new(models.PlatformAccountPoolCursor)).Filter("id", payload.ID).One(&row); err != nil {
|
||||
poolJSONErr(c, 404, 404, "记录不存在")
|
||||
return
|
||||
}
|
||||
token = strings.TrimSpace(row.Token)
|
||||
}
|
||||
case "windsurf":
|
||||
if payload.ID == 0 {
|
||||
poolJSONErr(c, 400, 400, "缺少有效 id")
|
||||
return
|
||||
}
|
||||
var row models.PlatformAccountPoolWindsurf
|
||||
if err := models.Orm.QueryTable(new(models.PlatformAccountPoolWindsurf)).Filter("id", payload.ID).One(&row); err != nil {
|
||||
poolJSONErr(c, 404, 404, "记录不存在")
|
||||
return
|
||||
}
|
||||
token = strings.TrimSpace(row.Token)
|
||||
case "krio":
|
||||
if payload.ID == 0 {
|
||||
poolJSONErr(c, 400, 400, "缺少有效 id")
|
||||
return
|
||||
}
|
||||
var row models.PlatformAccountPoolKiro
|
||||
if err := models.Orm.QueryTable(new(models.PlatformAccountPoolKiro)).Filter("id", payload.ID).One(&row); err != nil {
|
||||
poolJSONErr(c, 404, 404, "记录不存在")
|
||||
return
|
||||
}
|
||||
token = strings.TrimSpace(row.Token)
|
||||
default:
|
||||
poolJSONErr(c, 400, 400, "无效模块")
|
||||
return
|
||||
}
|
||||
|
||||
if token == "" {
|
||||
poolJSONErr(c, 400, 400, "该记录无 Token，无法探测")
|
||||
return
|
||||
}
|
||||
|
||||
r := tokenprobe.ProbeOfficial(module, token)
|
||||
data := map[string]interface{}{
|
||||
"ok": r.OK,
|
||||
"detail": r.Detail,
|
||||
"httpStatus": r.HTTPStatus,
|
||||
}
|
||||
if r.ProbeMessage != "" {
|
||||
data["probeMessage"] = r.ProbeMessage
|
||||
}
|
||||
if r.Endpoint != "" {
|
||||
data["endpoint"] = r.Endpoint
|
||||
}
|
||||
if r.BytesRead > 0 {
|
||||
data["bytesRead"] = r.BytesRead
|
||||
}
|
||||
if r.RawPreview != "" {
|
||||
data["rawPreview"] = r.RawPreview
|
||||
}
|
||||
if r.RequestBodyPrefixHex != "" {
|
||||
data["requestBodyPrefixHex"] = r.RequestBodyPrefixHex
|
||||
}
|
||||
if r.StreamProtocol != "" {
|
||||
data["streamProtocol"] = r.StreamProtocol
|
||||
}
|
||||
if r.StreamNote != "" {
|
||||
data["streamNote"] = r.StreamNote
|
||||
}
|
||||
if module == "cursor" && payload.ID > 0 && r.HTTPStatus == http.StatusOK {
|
||||
var isUsed int8
|
||||
if r.OK {
|
||||
isUsed = 1
|
||||
} else {
|
||||
isUsed = 0
|
||||
}
|
||||
if _, uerr := models.Orm.QueryTable(new(models.PlatformAccountPoolCursor)).Filter("id", payload.ID).Update(orm.Params{
|
||||
"is_used": isUsed,
|
||||
"update_time": time.Now(),
|
||||
}); uerr == nil {
|
||||
data["is_used"] = int(isUsed)
|
||||
}
|
||||
}
|
||||
c.Data["json"] = map[string]interface{}{
|
||||
"code": 200,
|
||||
"msg": "success",
|
||||
"data": data,
|
||||
}
|
||||
_ = c.ServeJSON()
|
||||
}
|
||||
|
||||
func (c *PlatformAccountPoolCursorController) List() { listPoolRows(&c.Controller, "cursor") }
|
||||
func (c *PlatformAccountPoolCursorController) Add() { addPoolRow(&c.Controller, "cursor") }
|
||||
func (c *PlatformAccountPoolCursorController) BatchAdd() { batchAddPoolRows(&c.Controller, "cursor") }
|
||||
func (c *PlatformAccountPoolCursorController) Detail() { getPoolDetail(&c.Controller, "cursor") }
|
||||
func (c *PlatformAccountPoolCursorController) Extract() { extractPoolRow(&c.Controller, "cursor") }
|
||||
func (c *PlatformAccountPoolCursorController) Replenish() { replenishPoolRow(&c.Controller, "cursor") }
|
||||
func (c *PlatformAccountPoolCursorController) UpdateRemark() {
|
||||
updatePoolRemark(&c.Controller, "cursor")
|
||||
}
|
||||
func (c *PlatformAccountPoolCursorController) ProbeToken() { probePoolToken(&c.Controller, "cursor") }
|
||||
|
||||
func (c *PlatformAccountPoolWindsurfController) List() { listPoolRows(&c.Controller, "windsurf") }
|
||||
func (c *PlatformAccountPoolWindsurfController) Add() { addPoolRow(&c.Controller, "windsurf") }
|
||||
func (c *PlatformAccountPoolWindsurfController) BatchAdd() { batchAddPoolRows(&c.Controller, "windsurf") }
|
||||
func (c *PlatformAccountPoolWindsurfController) Detail() { getPoolDetail(&c.Controller, "windsurf") }
|
||||
func (c *PlatformAccountPoolWindsurfController) Extract() { extractPoolRow(&c.Controller, "windsurf") }
|
||||
func (c *PlatformAccountPoolWindsurfController) Replenish() { replenishPoolRow(&c.Controller, "windsurf") }
|
||||
func (c *PlatformAccountPoolWindsurfController) List() { listPoolRows(&c.Controller, "windsurf") }
|
||||
func (c *PlatformAccountPoolWindsurfController) Add() { addPoolRow(&c.Controller, "windsurf") }
|
||||
func (c *PlatformAccountPoolWindsurfController) BatchAdd() {
|
||||
batchAddPoolRows(&c.Controller, "windsurf")
|
||||
}
|
||||
func (c *PlatformAccountPoolWindsurfController) Detail() { getPoolDetail(&c.Controller, "windsurf") }
|
||||
func (c *PlatformAccountPoolWindsurfController) Extract() { extractPoolRow(&c.Controller, "windsurf") }
|
||||
func (c *PlatformAccountPoolWindsurfController) Replenish() {
|
||||
replenishPoolRow(&c.Controller, "windsurf")
|
||||
}
|
||||
func (c *PlatformAccountPoolWindsurfController) UpdateRemark() {
|
||||
updatePoolRemark(&c.Controller, "windsurf")
|
||||
}
|
||||
func (c *PlatformAccountPoolWindsurfController) ProbeToken() {
|
||||
probePoolToken(&c.Controller, "windsurf")
|
||||
}
|
||||
|
||||
func (c *PlatformAccountPoolKrioController) List() { listPoolRows(&c.Controller, "krio") }
|
||||
func (c *PlatformAccountPoolKrioController) Add() { addPoolRow(&c.Controller, "krio") }
|
||||
func (c *PlatformAccountPoolKrioController) BatchAdd() { batchAddPoolRows(&c.Controller, "krio") }
|
||||
func (c *PlatformAccountPoolKrioController) Detail() { getPoolDetail(&c.Controller, "krio") }
|
||||
func (c *PlatformAccountPoolKrioController) Extract() { extractPoolRow(&c.Controller, "krio") }
|
||||
func (c *PlatformAccountPoolKrioController) List() { listPoolRows(&c.Controller, "krio") }
|
||||
func (c *PlatformAccountPoolKrioController) Add() { addPoolRow(&c.Controller, "krio") }
|
||||
func (c *PlatformAccountPoolKrioController) BatchAdd() { batchAddPoolRows(&c.Controller, "krio") }
|
||||
func (c *PlatformAccountPoolKrioController) Detail() { getPoolDetail(&c.Controller, "krio") }
|
||||
func (c *PlatformAccountPoolKrioController) Extract() { extractPoolRow(&c.Controller, "krio") }
|
||||
func (c *PlatformAccountPoolKrioController) Replenish() { replenishPoolRow(&c.Controller, "krio") }
|
||||
func (c *PlatformAccountPoolKrioController) UpdateRemark() {
|
||||
updatePoolRemark(&c.Controller, "krio")
|
||||
}
|
||||
func (c *PlatformAccountPoolKrioController) ProbeToken() { probePoolToken(&c.Controller, "krio") }
|
||||
|
||||
7
go.mod
7
go.mod
@ -9,7 +9,11 @@ require (
|
||||
golang.org/x/crypto v0.1.0 // indirect
|
||||
)
|
||||
|
||||
require github.com/go-sql-driver/mysql v1.7.0
|
||||
require (
|
||||
github.com/go-sql-driver/mysql v1.7.0
|
||||
github.com/google/uuid v1.6.0
|
||||
golang.org/x/net v0.7.0
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/beorn7/perks v1.0.1 // indirect
|
||||
@ -24,7 +28,6 @@ require (
|
||||
github.com/prometheus/common v0.42.0 // indirect
|
||||
github.com/prometheus/procfs v0.9.0 // indirect
|
||||
github.com/shiena/ansicolor v0.0.0-20200904210342-c7312218db18 // indirect
|
||||
golang.org/x/net v0.7.0 // indirect
|
||||
golang.org/x/sync v0.1.0 // indirect
|
||||
golang.org/x/sys v0.6.0 // indirect
|
||||
golang.org/x/text v0.7.0 // indirect
|
||||
|
||||
2
go.sum
2
go.sum
@ -234,6 +234,8 @@ github.com/google/pprof v0.0.0-20200708004538-1a94d8640e99/go.mod h1:ZgVRPoUq/hf
|
||||
github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI=
|
||||
github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/google/uuid v1.2.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg=
|
||||
github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk=
|
||||
github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY=
|
||||
|
||||
@ -51,6 +51,7 @@ type PlatformAccountPoolCursor struct {
|
||||
Token string `orm:"column(token);type(text);null" json:"token"`
|
||||
Remark string `orm:"column(remark);size(255);default()" json:"remark"`
|
||||
IsExtracted int8 `orm:"column(is_extracted);default(0)" json:"is_extracted"`
|
||||
IsUsed *int8 `orm:"column(is_used);null" json:"is_used"` // 0=用完/不可用 1=可用 NULL=未探测
|
||||
ExtractedTime *time.Time `orm:"column(extracted_time);type(datetime);null" json:"extracted_time"`
|
||||
ExtractedPlatform *string `orm:"column(extracted_platform);size(32);null" json:"extracted_platform"`
|
||||
CreateTime time.Time `orm:"column(create_time);auto_now_add;type(datetime)" json:"create_time"`
|
||||
|
||||
560
pkg/tokenprobe/cursor_hi.go
Normal file
560
pkg/tokenprobe/cursor_hi.go
Normal file
@ -0,0 +1,560 @@
|
||||
package tokenprobe
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"crypto/sha256"
|
||||
"crypto/tls"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"runtime"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"golang.org/x/net/http2"
|
||||
)
|
||||
|
||||
const (
|
||||
cursorBackendURL = "https://api2.cursor.sh"
|
||||
cursorAgentPath = "/aiserver.v1.ChatService/StreamUnifiedChatWithTools"
|
||||
cursorClientVersion = "2.6.22"
|
||||
cursorHiMaxRead = 512 * 1024
|
||||
// probeHiText 发往官方 Agent 的探测内容（与前端展示 probeMessage 一致）
|
||||
probeHiText = "hi"
|
||||
)
|
||||
|
||||
var cursorProbeHTTPClient = newCursorHTTP2Client()
|
||||
|
||||
func newCursorHTTP2Client() *http.Client {
|
||||
tr := &http.Transport{
|
||||
TLSClientConfig: &tls.Config{MinVersion: tls.VersionTLS12},
|
||||
}
|
||||
// 与 Cursor 官方一致走 HTTP/2
|
||||
if err := http2.ConfigureTransport(tr); err != nil {
|
||||
return &http.Client{Timeout: 40 * time.Second}
|
||||
}
|
||||
return &http.Client{Transport: tr, Timeout: 40 * time.Second}
|
||||
}
|
||||
|
||||
func cursorClientOS() string {
|
||||
switch runtime.GOOS {
|
||||
case "windows":
|
||||
return "win32"
|
||||
case "darwin":
|
||||
return "darwin"
|
||||
default:
|
||||
return "linux"
|
||||
}
|
||||
}
|
||||
|
||||
func cursorClientArch() string {
|
||||
switch runtime.GOARCH {
|
||||
case "amd64":
|
||||
return "x64"
|
||||
case "arm64":
|
||||
return "arm64"
|
||||
default:
|
||||
return runtime.GOARCH
|
||||
}
|
||||
}
|
||||
|
||||
func cursorEnvVersion() string {
|
||||
if v := strings.TrimSpace(os.Getenv("CURSOR_CLIENT_VERSION")); v != "" {
|
||||
return v
|
||||
}
|
||||
return cursorClientVersion
|
||||
}
|
||||
|
||||
// --- protobuf wire (与 cursor_api_demo 对齐) ---
|
||||
|
||||
func pbVarint(v uint64) []byte {
|
||||
var out []byte
|
||||
for v >= 0x80 {
|
||||
out = append(out, byte(v&0x7f|0x80))
|
||||
v >>= 7
|
||||
}
|
||||
out = append(out, byte(v&0x7f))
|
||||
return out
|
||||
}
|
||||
|
||||
func pbField(fieldNum int, wireType int, value interface{}) []byte {
|
||||
tag := uint64(fieldNum<<3 | wireType)
|
||||
out := pbVarint(tag)
|
||||
switch wireType {
|
||||
case 0:
|
||||
var n uint64
|
||||
switch x := value.(type) {
|
||||
case int:
|
||||
n = uint64(x)
|
||||
case int32:
|
||||
n = uint64(x)
|
||||
case uint32:
|
||||
n = uint64(x)
|
||||
case uint64:
|
||||
n = x
|
||||
default:
|
||||
n = uint64(0)
|
||||
}
|
||||
out = append(out, pbVarint(n)...)
|
||||
case 2:
|
||||
var b []byte
|
||||
switch x := value.(type) {
|
||||
case string:
|
||||
b = []byte(x)
|
||||
case []byte:
|
||||
b = x
|
||||
default:
|
||||
b = []byte(fmt.Sprint(x))
|
||||
}
|
||||
out = append(out, pbVarint(uint64(len(b)))...)
|
||||
out = append(out, b...)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func encodeCursorMessage(content string, role int, messageID string, chatModeEnum *int) []byte {
|
||||
msg := pbField(1, 2, content)
|
||||
msg = append(msg, pbField(2, 0, role)...)
|
||||
msg = append(msg, pbField(13, 2, messageID)...)
|
||||
if chatModeEnum != nil {
|
||||
msg = append(msg, pbField(47, 0, *chatModeEnum)...)
|
||||
}
|
||||
return msg
|
||||
}
|
||||
|
||||
func encodeCursorModel(modelName string) []byte {
|
||||
msg := pbField(1, 2, modelName)
|
||||
msg = append(msg, pbField(4, 2, []byte{})...)
|
||||
return msg
|
||||
}
|
||||
|
||||
func encodeCursorSetting() []byte {
|
||||
inner := pbField(1, 2, []byte{})
|
||||
inner = append(inner, pbField(2, 2, []byte{})...)
|
||||
msg := pbField(1, 2, `cursor\aisettings`)
|
||||
msg = append(msg, pbField(3, 2, []byte{})...)
|
||||
msg = append(msg, pbField(6, 2, inner)...)
|
||||
msg = append(msg, pbField(8, 0, 1)...)
|
||||
msg = append(msg, pbField(9, 0, 1)...)
|
||||
return msg
|
||||
}
|
||||
|
||||
func encodeCursorMetadata() []byte {
|
||||
msg := pbField(1, 2, cursorClientOS())
|
||||
msg = append(msg, pbField(2, 2, cursorClientArch())...)
|
||||
msg = append(msg, pbField(3, 2, "unknown")...)
|
||||
msg = append(msg, pbField(4, 2, "go-platform/tokenprobe")...)
|
||||
msg = append(msg, pbField(5, 2, time.Now().Format(time.RFC3339))...)
|
||||
return msg
|
||||
}
|
||||
|
||||
func encodeCursorMessageID(messageID string, role int) []byte {
|
||||
msg := pbField(1, 2, messageID)
|
||||
msg = append(msg, pbField(3, 0, role)...)
|
||||
return msg
|
||||
}
|
||||
|
||||
// defaultAgentTools 与 cursor_agent_client.DEFAULT_TOOLS 一致
|
||||
var defaultAgentTools = []int{5, 6, 3, 15, 7, 8, 42}
|
||||
|
||||
func encodeCursorAgentRequest(userContent, modelName string) []byte {
|
||||
msgID := uuid.NewString()
|
||||
cm := 2 // Agent
|
||||
userMsg := encodeCursorMessage(userContent, 1, msgID, &cm)
|
||||
|
||||
var msg []byte
|
||||
msg = append(msg, pbField(1, 2, userMsg)...)
|
||||
msg = append(msg, pbField(2, 0, 1)...)
|
||||
msg = append(msg, pbField(3, 2, []byte{})...)
|
||||
msg = append(msg, pbField(4, 0, 1)...)
|
||||
msg = append(msg, pbField(5, 2, encodeCursorModel(modelName))...)
|
||||
msg = append(msg, pbField(8, 2, "")...)
|
||||
msg = append(msg, pbField(13, 0, 1)...)
|
||||
msg = append(msg, pbField(15, 2, encodeCursorSetting())...)
|
||||
msg = append(msg, pbField(19, 0, 1)...)
|
||||
msg = append(msg, pbField(23, 2, uuid.NewString())...)
|
||||
msg = append(msg, pbField(26, 2, encodeCursorMetadata())...)
|
||||
msg = append(msg, pbField(27, 0, 1)...)
|
||||
for _, t := range defaultAgentTools {
|
||||
msg = append(msg, pbField(29, 0, t)...)
|
||||
}
|
||||
msg = append(msg, pbField(30, 2, encodeCursorMessageID(msgID, 1))...)
|
||||
msg = append(msg, pbField(35, 0, 0)...)
|
||||
msg = append(msg, pbField(38, 0, 0)...)
|
||||
msg = append(msg, pbField(46, 0, 2)...)
|
||||
msg = append(msg, pbField(47, 2, "")...)
|
||||
msg = append(msg, pbField(48, 0, 0)...)
|
||||
msg = append(msg, pbField(49, 0, 0)...)
|
||||
msg = append(msg, pbField(51, 0, 0)...)
|
||||
msg = append(msg, pbField(53, 0, 1)...)
|
||||
msg = append(msg, pbField(54, 2, "agent")...)
|
||||
return msg
|
||||
}
|
||||
|
||||
func encodeStreamUnifiedChatWithToolsRequest(inner []byte) []byte {
|
||||
return pbField(1, 2, inner)
|
||||
}
|
||||
|
||||
func generateCursorAgentFramedBody(userText, model string) []byte {
|
||||
inner := encodeCursorAgentRequest(userText, model)
|
||||
buf := encodeStreamUnifiedChatWithToolsRequest(inner)
|
||||
magic := byte(0x00)
|
||||
hexLen := fmt.Sprintf("%08x", len(buf))
|
||||
lenBytes, err := hex.DecodeString(hexLen)
|
||||
if err != nil || len(lenBytes) != 4 {
|
||||
lenB := []byte{byte(len(buf) >> 24), byte(len(buf) >> 16), byte(len(buf) >> 8), byte(len(buf))}
|
||||
return append([]byte{magic}, append(lenB, buf...)...)
|
||||
}
|
||||
return append([]byte{magic}, append(lenBytes, buf...)...)
|
||||
}
|
||||
|
||||
func hashed64Hex(input, salt string) string {
|
||||
h := sha256.Sum256([]byte(input + salt))
|
||||
return hex.EncodeToString(h[:])
|
||||
}
|
||||
|
||||
func generateCursorChecksum(authToken string) string {
|
||||
machineID := hashed64Hex(authToken, "machineId")
|
||||
ts := int(time.Now().UnixMilli() / 1_000_000)
|
||||
barr := []byte{
|
||||
byte(ts >> 40), byte(ts >> 32), byte(ts >> 24), byte(ts >> 16), byte(ts >> 8), byte(ts),
|
||||
}
|
||||
t := byte(165)
|
||||
for i := range barr {
|
||||
barr[i] = ((barr[i] ^ t) + byte(i%256)) & 255
|
||||
t = barr[i]
|
||||
}
|
||||
const alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"
|
||||
var enc strings.Builder
|
||||
for i := 0; i < len(barr); i += 3 {
|
||||
a := barr[i]
|
||||
var b, c byte
|
||||
if i+1 < len(barr) {
|
||||
b = barr[i+1]
|
||||
}
|
||||
if i+2 < len(barr) {
|
||||
c = barr[i+2]
|
||||
}
|
||||
enc.WriteByte(alphabet[a>>2])
|
||||
enc.WriteByte(alphabet[((a&3)<<4)|(b>>4)])
|
||||
if i+1 < len(barr) {
|
||||
enc.WriteByte(alphabet[((b&15)<<2)|(c>>6)])
|
||||
}
|
||||
if i+2 < len(barr) {
|
||||
enc.WriteByte(alphabet[c&63])
|
||||
}
|
||||
}
|
||||
return enc.String() + machineID
|
||||
}
|
||||
|
||||
func asciiLowerInPlace(b []byte) {
|
||||
for i := range b {
|
||||
c := b[i]
|
||||
if c >= 'A' && c <= 'Z' {
|
||||
b[i] = c + ('a' - 'A')
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 社区脚本中的「额度用尽」ASCII 前缀（与明文一致，便于在二进制流中 bytes.Contains，无需整句）
|
||||
// 对应明文前缀：Get Cursor Pro for more Agent usage
|
||||
var cursorQuotaExhaustedSigCommunity = []byte{
|
||||
0x47, 0x65, 0x74, 0x20, 0x43, 0x75, 0x72, 0x73,
|
||||
0x6f, 0x72, 0x20, 0x50, 0x72, 0x6f, 0x20, 0x66,
|
||||
0x6f, 0x72, 0x20, 0x6d, 0x6f, 0x72, 0x65, 0x20,
|
||||
0x41, 0x67, 0x65, 0x6e, 0x74, 0x20, 0x75, 0x73,
|
||||
0x61, 0x67, 0x65,
|
||||
}
|
||||
|
||||
// cursorQuotaTipSig 与常见示例一致：raw 全字节里 bytes.Contains(raw, tipSig) → 额度用尽
|
||||
var cursorQuotaTipSig = []byte("Get Cursor Pro for more Agent usage, unlimited Tab, and more.")
|
||||
|
||||
const cursorLimitTipPrefix = "Get Cursor Pro for more Agent usage, unlimited Tab"
|
||||
|
||||
// classifyCursorRawStream 在官方流式二进制/文本中匹配用量与升级提示（ASCII 区不区分大小写 + UTF-8 短语）
|
||||
func classifyCursorRawStream(raw []byte) (blocked bool, reason string) {
|
||||
if len(raw) == 0 {
|
||||
return false, ""
|
||||
}
|
||||
for _, sig := range cursorQuotaExhaustedSigsFromEnv() {
|
||||
if bytes.Contains(raw, sig) {
|
||||
return true, fmt.Sprintf("流中匹配：CURSOR_QUOTA_EXHAUSTED_SIG_HEX 配置的二进制特征（%d 字节）", len(sig))
|
||||
}
|
||||
}
|
||||
if bytes.Contains(raw, cursorQuotaTipSig) {
|
||||
return true, "流中匹配：" + string(cursorQuotaTipSig)
|
||||
}
|
||||
// 社区脚本：仅到「…Agent usage」的 ASCII 前缀（流里可能只有前半段）
|
||||
if bytes.Contains(raw, cursorQuotaExhaustedSigCommunity) {
|
||||
return true, "流中匹配：Get Cursor Pro for more Agent usage…（社区 QuotaExhaustedSignature 前缀）"
|
||||
}
|
||||
if bytes.Contains(raw, []byte(cursorLimitTipPrefix)) {
|
||||
return true, "流中匹配：" + cursorLimitTipPrefix + "…"
|
||||
}
|
||||
low := append([]byte(nil), raw...)
|
||||
asciiLowerInPlace(low)
|
||||
if bytes.Contains(low, []byte("you've hit your usage limit")) ||
|
||||
bytes.Contains(low, []byte("youve hit your usage limit")) ||
|
||||
bytes.Contains(low, []byte("hit your usage limit")) {
|
||||
return true, "流中匹配：hit your usage limit / you've hit your usage limit"
|
||||
}
|
||||
if bytes.Contains(low, []byte("get cursor pro for more agent usage")) {
|
||||
return true, "流中匹配：get cursor pro for more agent usage"
|
||||
}
|
||||
if bytes.Contains(low, []byte("upgrade to pro")) {
|
||||
return true, "流中匹配：upgrade to pro"
|
||||
}
|
||||
if bytes.Contains(low, []byte("get cursor pro")) && bytes.Contains(low, []byte("agent")) {
|
||||
return true, "流中匹配：get cursor pro + agent"
|
||||
}
|
||||
if bytes.Contains(low, []byte("usage limit")) {
|
||||
return true, "流中匹配：usage limit"
|
||||
}
|
||||
if bytes.Contains(low, []byte("unlimited tab")) && bytes.Contains(low, []byte("cursor pro")) {
|
||||
return true, "流中匹配：unlimited tab + cursor pro"
|
||||
}
|
||||
|
||||
flat := strings.ToLower(strings.ToValidUTF8(string(raw), "\uFFFD"))
|
||||
flat = strings.ReplaceAll(flat, "\u2019", "'") // 右单引号
|
||||
flat = strings.ReplaceAll(flat, "`", "'")
|
||||
if strings.Contains(flat, "you've hit your usage limit") {
|
||||
return true, "流中匹配：you've hit your usage limit（UTF-8）"
|
||||
}
|
||||
return false, ""
|
||||
}
|
||||
|
||||
func truncateUTF8Preview(raw []byte, maxBytes int) string {
|
||||
s := strings.ToValidUTF8(string(raw), "\uFFFD")
|
||||
if maxBytes <= 0 || len(s) <= maxBytes {
|
||||
return s
|
||||
}
|
||||
// 按字节截断并保证合法 UTF-8
|
||||
s = s[:maxBytes]
|
||||
for len(s) > 0 && !utf8.ValidString(s) {
|
||||
s = s[:len(s)-1]
|
||||
}
|
||||
return s + "…(已截断)"
|
||||
}
|
||||
|
||||
func prefixHexBody(b []byte, max int) string {
|
||||
if len(b) > max {
|
||||
b = b[:max]
|
||||
}
|
||||
return hex.EncodeToString(b)
|
||||
}
|
||||
|
||||
func looksLikeGzip(raw []byte) bool {
|
||||
return len(raw) >= 3 && raw[0] == 0x1f && raw[1] == 0x8b && raw[2] == 0x08
|
||||
}
|
||||
|
||||
func gunzipBytes(raw []byte) ([]byte, error) {
|
||||
zr, err := gzip.NewReader(bytes.NewReader(raw))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer zr.Close()
|
||||
return io.ReadAll(io.LimitReader(zr, cursorHiMaxRead))
|
||||
}
|
||||
|
||||
func decodeConnectFramedBody(raw []byte) ([]byte, string, bool) {
|
||||
if len(raw) < 5 {
|
||||
return nil, "", false
|
||||
}
|
||||
|
||||
var out bytes.Buffer
|
||||
offset := 0
|
||||
frameCount := 0
|
||||
compressedFrames := 0
|
||||
|
||||
for offset+5 <= len(raw) {
|
||||
flags := raw[offset]
|
||||
n := int(raw[offset+1])<<24 | int(raw[offset+2])<<16 | int(raw[offset+3])<<8 | int(raw[offset+4])
|
||||
offset += 5
|
||||
if n < 0 || offset+n > len(raw) {
|
||||
return nil, "", false
|
||||
}
|
||||
payload := raw[offset : offset+n]
|
||||
offset += n
|
||||
frameCount++
|
||||
|
||||
isCompressed := flags&0x01 == 0x01
|
||||
if isCompressed || looksLikeGzip(payload) {
|
||||
decoded, err := gunzipBytes(payload)
|
||||
if err != nil {
|
||||
out.Write(payload)
|
||||
} else {
|
||||
out.Write(decoded)
|
||||
compressedFrames++
|
||||
}
|
||||
} else {
|
||||
out.Write(payload)
|
||||
}
|
||||
}
|
||||
|
||||
if frameCount == 0 || offset != len(raw) {
|
||||
return nil, "", false
|
||||
}
|
||||
|
||||
note := fmt.Sprintf("响应体已按 Connect 分帧解析（%d 帧", frameCount)
|
||||
if compressedFrames > 0 {
|
||||
note += fmt.Sprintf("，其中 %d 帧已做 gzip 解压", compressedFrames)
|
||||
}
|
||||
note += "）后分析"
|
||||
return out.Bytes(), note, true
|
||||
}
|
||||
|
||||
func decodeCursorResponseBody(raw []byte, contentEncoding string) ([]byte, string) {
|
||||
if decoded, note, ok := decodeConnectFramedBody(raw); ok {
|
||||
return decoded, note
|
||||
}
|
||||
|
||||
enc := strings.ToLower(strings.TrimSpace(contentEncoding))
|
||||
if strings.Contains(enc, "gzip") || looksLikeGzip(raw) {
|
||||
decoded, err := gunzipBytes(raw)
|
||||
if err != nil {
|
||||
if strings.Contains(enc, "gzip") {
|
||||
return raw, "响应头声明 gzip，但解压失败，已回退为原始字节预览"
|
||||
}
|
||||
return raw, "检测到 gzip 魔数，但解压失败，已回退为原始字节预览"
|
||||
}
|
||||
if strings.Contains(enc, "gzip") {
|
||||
return decoded, "响应体已按 gzip 解压后分析"
|
||||
}
|
||||
return decoded, "响应体虽未显式声明 Content-Encoding，但按 gzip 魔数解压后分析"
|
||||
}
|
||||
if enc != "" {
|
||||
return raw, "响应头 Content-Encoding=" + enc + "，当前未额外解码，按原始字节分析"
|
||||
}
|
||||
return raw, "响应体未压缩或未声明压缩，且未识别为 Connect 分帧，按原始字节分析"
|
||||
}
|
||||
|
||||
// cursorStreamProtocol 与官方客户端一致：Connect-RPC + protobuf 体，HTTP/2 流式
|
||||
const cursorStreamProtocol = "Connect-Protocol-Version:1 + application/connect+proto，HTTP/2 二进制流（gRPC 兼容形态，非 JSON REST）"
|
||||
|
||||
// cursorStreamNote 说明 rawPreview / ok 的含义边界（与「仅通 200」结论一致）
|
||||
const cursorStreamNote = `【协议】本 URL 为 Cursor 官方 Agent 流式接口，请求体为 protobuf（requestBodyPrefixHex 可见非表单/JSON）。` +
|
||||
`【响应】正文为分包二进制流，rawPreview 是按 UTF-8 有损解码的片段，绝大多数情况下会像乱码，属正常现象，不能当普通 UTF-8 接口正文解析。` +
|
||||
`【HTTP 200】仅表示 TLS/代理/网络到 api2.cursor.sh 通畅，不代表 protobuf 业务层、鉴权、设备指纹、风控配额或「能持续对话」已全部通过。` +
|
||||
`【ok 字段】当前仅在解码片段上做英文关键词启发式匹配；未命中不代表账户可用，命中也不覆盖「须在 IDE 内完整走流式协议」的场景。` +
|
||||
`【若要等价客户端】需完整实现 Connect 帧解析、会话与校验头、可能的 gzip/分包及双向流，本探测只做粗连通与可观测性辅助。` +
|
||||
`【二进制特征】若提示语被包在 protobuf 字段内、ASCII 子串匹配不到，可在运行环境设置 CURSOR_QUOTA_EXHAUSTED_SIG_HEX=hex1,hex2（逗号分隔十六进制，可选 0x 前缀），在原始响应字节上做 bytes.Contains，无需解析整条 proto；特征需自行对比「额度正常」与「用尽」两次抓包提取。`
|
||||
|
||||
// cursorQuotaExhaustedSigsFromEnv 从环境变量解析额度用尽时的二进制特征（不转 UTF-8）
|
||||
func cursorQuotaExhaustedSigsFromEnv() [][]byte {
|
||||
s := strings.TrimSpace(os.Getenv("CURSOR_QUOTA_EXHAUSTED_SIG_HEX"))
|
||||
if s == "" {
|
||||
return nil
|
||||
}
|
||||
var out [][]byte
|
||||
for _, part := range strings.Split(s, ",") {
|
||||
part = strings.TrimSpace(part)
|
||||
if part == "" {
|
||||
continue
|
||||
}
|
||||
part = strings.TrimPrefix(strings.TrimPrefix(part, "0x"), "0X")
|
||||
b, err := hex.DecodeString(part)
|
||||
if err != nil || len(b) == 0 {
|
||||
continue
|
||||
}
|
||||
out = append(out, b)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func cursorProbeResult(ok bool, detail string, httpStatus int, reqBody, raw, preview []byte) Result {
|
||||
if preview == nil {
|
||||
preview = raw
|
||||
}
|
||||
return Result{
|
||||
OK: ok,
|
||||
Detail: detail,
|
||||
HTTPStatus: httpStatus,
|
||||
ProbeMessage: probeHiText,
|
||||
Endpoint: cursorBackendURL + cursorAgentPath,
|
||||
BytesRead: len(raw),
|
||||
RawPreview: truncateUTF8Preview(preview, 24000),
|
||||
RequestBodyPrefixHex: prefixHexBody(reqBody, 128),
|
||||
StreamProtocol: cursorStreamProtocol,
|
||||
StreamNote: cursorStreamNote,
|
||||
}
|
||||
}
|
||||
|
||||
func probeCursorHiAgent(authToken string) Result {
|
||||
if strings.Contains(authToken, "::") {
|
||||
if i := strings.LastIndex(authToken, "::"); i >= 0 {
|
||||
authToken = strings.TrimSpace(authToken[i+2:])
|
||||
}
|
||||
}
|
||||
if authToken == "" {
|
||||
return Result{OK: false, Detail: "Token 为空"}
|
||||
}
|
||||
|
||||
sessionID := uuid.NewSHA1(uuid.NameSpaceDNS, []byte(authToken)).String()
|
||||
clientKey := hashed64Hex(authToken, "")
|
||||
checksum := generateCursorChecksum(authToken)
|
||||
conversationID := uuid.NewString()
|
||||
reqID := uuid.NewString()
|
||||
|
||||
body := generateCursorAgentFramedBody(probeHiText, "default")
|
||||
fullURL := cursorBackendURL + cursorAgentPath
|
||||
req, err := http.NewRequest(http.MethodPost, fullURL, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
r := cursorProbeResult(false, err.Error(), 0, body, nil, nil)
|
||||
return r
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+authToken)
|
||||
req.Header.Set("Connect-Accept-Encoding", "gzip")
|
||||
req.Header.Set("Connect-Protocol-Version", "1")
|
||||
req.Header.Set("Content-Type", "application/connect+proto")
|
||||
req.Header.Set("User-Agent", "connect-es/1.6.1")
|
||||
req.Header.Set("X-Amzn-Trace-Id", "Root="+reqID)
|
||||
req.Header.Set("X-Client-Key", clientKey)
|
||||
req.Header.Set("X-Cursor-Checksum", checksum)
|
||||
req.Header.Set("X-Cursor-Client-Version", cursorEnvVersion())
|
||||
req.Header.Set("X-Cursor-Client-Type", "ide")
|
||||
req.Header.Set("X-Cursor-Client-Os", cursorClientOS())
|
||||
req.Header.Set("X-Cursor-Client-Arch", cursorClientArch())
|
||||
req.Header.Set("X-Cursor-Client-Os-Version", "unknown")
|
||||
req.Header.Set("X-Cursor-Client-Device-Type", "desktop")
|
||||
req.Header.Set("X-Cursor-Config-Version", uuid.NewString())
|
||||
req.Header.Set("X-Cursor-Timezone", "UTC")
|
||||
req.Header.Set("X-Ghost-Mode", "false")
|
||||
req.Header.Set("X-New-Onboarding-Completed", "true")
|
||||
req.Header.Set("X-Request-Id", reqID)
|
||||
req.Header.Set("X-Session-Id", sessionID)
|
||||
req.Header.Set("X-Conversation-Id", conversationID)
|
||||
req.Host = "api2.cursor.sh"
|
||||
|
||||
resp, err := cursorProbeHTTPClient.Do(req)
|
||||
if err != nil {
|
||||
return cursorProbeResult(false, "请求 Cursor Agent 失败: "+err.Error(), 0, body, nil, nil)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
raw, _ := io.ReadAll(io.LimitReader(resp.Body, cursorHiMaxRead))
|
||||
decoded, decodeNote := decodeCursorResponseBody(raw, resp.Header.Get("Content-Encoding"))
|
||||
blocked, reason := classifyCursorRawStream(decoded)
|
||||
if blocked {
|
||||
return cursorProbeResult(false, reason+"；"+decodeNote, resp.StatusCode, body, raw, decoded)
|
||||
}
|
||||
detail := fmt.Sprintf("HTTP %d（非 200）；%s；说明与协议边界见 streamNote", resp.StatusCode, decodeNote)
|
||||
return cursorProbeResult(false, detail, resp.StatusCode, body, raw, decoded)
|
||||
}
|
||||
|
||||
var buf bytes.Buffer
|
||||
_, _ = io.Copy(&buf, io.LimitReader(resp.Body, cursorHiMaxRead))
|
||||
raw := buf.Bytes()
|
||||
decoded, decodeNote := decodeCursorResponseBody(raw, resp.Header.Get("Content-Encoding"))
|
||||
blocked, reason := classifyCursorRawStream(decoded)
|
||||
if blocked {
|
||||
return cursorProbeResult(false, reason+"；"+decodeNote, resp.StatusCode, body, raw, decoded)
|
||||
}
|
||||
detail := "HTTP 200；未命中内置英文关键词；" + decodeNote + "；二进制流含义与 ok 边界见 streamNote"
|
||||
return cursorProbeResult(true, detail, resp.StatusCode, body, raw, decoded)
|
||||
}
|
||||
216
pkg/tokenprobe/probe.go
Normal file
216
pkg/tokenprobe/probe.go
Normal file
@ -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 ""
|
||||
}
|
||||
@ -170,6 +170,7 @@ func Register() {
|
||||
beego.Router("/platform/accountPool/cursor/extract", &controllers.PlatformAccountPoolCursorController{}, "post:Extract")
|
||||
beego.Router("/platform/accountPool/cursor/updateRemark", &controllers.PlatformAccountPoolCursorController{}, "post:UpdateRemark")
|
||||
beego.Router("/platform/accountPool/cursor/replenish", &controllers.PlatformAccountPoolCursorController{}, "post:Replenish")
|
||||
beego.Router("/platform/accountPool/cursor/probeToken", &controllers.PlatformAccountPoolCursorController{}, "post:ProbeToken")
|
||||
|
||||
beego.Router("/platform/accountPool/windsurf/list", &controllers.PlatformAccountPoolWindsurfController{}, "get:List")
|
||||
beego.Router("/platform/accountPool/windsurf/add", &controllers.PlatformAccountPoolWindsurfController{}, "post:Add")
|
||||
@ -178,6 +179,7 @@ func Register() {
|
||||
beego.Router("/platform/accountPool/windsurf/extract", &controllers.PlatformAccountPoolWindsurfController{}, "post:Extract")
|
||||
beego.Router("/platform/accountPool/windsurf/updateRemark", &controllers.PlatformAccountPoolWindsurfController{}, "post:UpdateRemark")
|
||||
beego.Router("/platform/accountPool/windsurf/replenish", &controllers.PlatformAccountPoolWindsurfController{}, "post:Replenish")
|
||||
beego.Router("/platform/accountPool/windsurf/probeToken", &controllers.PlatformAccountPoolWindsurfController{}, "post:ProbeToken")
|
||||
|
||||
beego.Router("/platform/accountPool/krio/list", &controllers.PlatformAccountPoolKrioController{}, "get:List")
|
||||
beego.Router("/platform/accountPool/krio/add", &controllers.PlatformAccountPoolKrioController{}, "post:Add")
|
||||
@ -186,4 +188,5 @@ func Register() {
|
||||
beego.Router("/platform/accountPool/krio/extract", &controllers.PlatformAccountPoolKrioController{}, "post:Extract")
|
||||
beego.Router("/platform/accountPool/krio/updateRemark", &controllers.PlatformAccountPoolKrioController{}, "post:UpdateRemark")
|
||||
beego.Router("/platform/accountPool/krio/replenish", &controllers.PlatformAccountPoolKrioController{}, "post:Replenish")
|
||||
beego.Router("/platform/accountPool/krio/probeToken", &controllers.PlatformAccountPoolKrioController{}, "post:ProbeToken")
|
||||
}
|
||||
|
||||
BIN
server.exe
BIN
server.exe
Binary file not shown.
Loading…
Reference in New Issue
Block a user