181 lines
5.4 KiB
Go
181 lines
5.4 KiB
Go
package service
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"strings"
|
|
|
|
kiropkg "github.com/Wei-Shaw/sub2api/internal/pkg/kiro"
|
|
"github.com/google/uuid"
|
|
)
|
|
|
|
func buildKiroAccountKey(account *Account) string {
|
|
if account == nil {
|
|
return ""
|
|
}
|
|
return kiropkg.BuildAccountKey(
|
|
account.GetCredential("client_id"),
|
|
account.GetCredential("client_id_hash"),
|
|
account.GetCredential("refresh_token"),
|
|
account.GetCredential("profile_arn"),
|
|
account.ID,
|
|
)
|
|
}
|
|
|
|
func buildKiroMachineID(account *Account) string {
|
|
if account == nil {
|
|
return kiropkg.BuildMachineID("", "", "account:nil")
|
|
}
|
|
for _, key := range []string{"machine_id", "machineId"} {
|
|
if machineID, ok := kiropkg.NormalizeMachineID(account.GetCredential(key)); ok {
|
|
return machineID
|
|
}
|
|
}
|
|
fallbackKey := buildKiroMachineIDFallbackKey(account)
|
|
if account.Type == AccountTypeAPIKey {
|
|
return kiropkg.BuildMachineID("", firstKiroCredential(account, "kiro_api_key", "kiroApiKey", "api_key"), fallbackKey)
|
|
}
|
|
return kiropkg.BuildMachineID(account.GetCredential("refresh_token"), "", fallbackKey)
|
|
}
|
|
|
|
func firstKiroCredential(account *Account, keys ...string) string {
|
|
if account == nil {
|
|
return ""
|
|
}
|
|
for _, key := range keys {
|
|
if value := strings.TrimSpace(account.GetCredential(key)); value != "" {
|
|
return value
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func buildKiroMachineIDFallbackKey(account *Account) string {
|
|
if account == nil {
|
|
return "account:nil"
|
|
}
|
|
if account.ID > 0 {
|
|
return fmt.Sprintf("account:%d", account.ID)
|
|
}
|
|
for _, key := range []string{"client_id", "profile_arn"} {
|
|
if value := strings.TrimSpace(account.GetCredential(key)); value != "" {
|
|
return key + ":" + value
|
|
}
|
|
}
|
|
if name := strings.TrimSpace(account.Name); name != "" {
|
|
return "name:" + name
|
|
}
|
|
return "account:unknown"
|
|
}
|
|
|
|
func buildKiroRequestID(resp *http.Response) string {
|
|
if resp == nil {
|
|
return ""
|
|
}
|
|
if requestID := strings.TrimSpace(resp.Header.Get("x-request-id")); requestID != "" {
|
|
return requestID
|
|
}
|
|
if requestID := strings.TrimSpace(resp.Header.Get("x-amzn-requestid")); requestID != "" {
|
|
return requestID
|
|
}
|
|
return strings.TrimSpace(resp.Header.Get("x-amz-request-id"))
|
|
}
|
|
|
|
func isKiroInvalidModelIDBody(respBody []byte) bool { //nolint:unused // exercised by internal tests for error classification.
|
|
var payload struct {
|
|
Reason string `json:"reason"`
|
|
Message string `json:"message"`
|
|
Error struct {
|
|
Reason string `json:"reason"`
|
|
Message string `json:"message"`
|
|
} `json:"error"`
|
|
}
|
|
if json.Unmarshal(respBody, &payload) != nil {
|
|
return looksLikeKiroBadRequestInvalidModelError(strings.ToLower(string(respBody)))
|
|
}
|
|
return strings.EqualFold(strings.TrimSpace(payload.Reason), "INVALID_MODEL_ID") ||
|
|
strings.EqualFold(strings.TrimSpace(payload.Error.Reason), "INVALID_MODEL_ID") ||
|
|
looksLikeKiroBadRequestInvalidModelError(strings.ToLower(payload.Message)) ||
|
|
looksLikeKiroBadRequestInvalidModelError(strings.ToLower(payload.Error.Message))
|
|
}
|
|
|
|
func isKiroSuspendedBody(respBody []byte) bool {
|
|
body := string(respBody)
|
|
return strings.Contains(body, "SUSPENDED") || strings.Contains(body, "TEMPORARILY_SUSPENDED")
|
|
}
|
|
|
|
func isKiroTokenErrorBody(respBody []byte) bool {
|
|
lower := strings.ToLower(string(respBody))
|
|
return strings.Contains(lower, "token") ||
|
|
strings.Contains(lower, "expired") ||
|
|
strings.Contains(lower, "invalid") ||
|
|
strings.Contains(lower, "unauthorized")
|
|
}
|
|
|
|
func kiroProxyURL(account *Account) string {
|
|
if account != nil && account.ProxyID != nil && account.Proxy != nil {
|
|
return account.Proxy.URL()
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func kiroAPIRegion(account *Account) string {
|
|
if account == nil {
|
|
return kiroDefaultRegion
|
|
}
|
|
region := strings.TrimSpace(account.GetCredential("api_region"))
|
|
if region == "" {
|
|
region = kiroDefaultRegion
|
|
}
|
|
return region
|
|
}
|
|
|
|
func applyKiroConditionalHeaders(req *http.Request, account *Account) {
|
|
if req == nil || account == nil {
|
|
return
|
|
}
|
|
if strings.EqualFold(strings.TrimSpace(account.GetCredential("auth_method")), "external_idp") {
|
|
req.Header.Set("TokenType", "EXTERNAL_IDP")
|
|
}
|
|
if strings.EqualFold(strings.TrimSpace(account.GetCredential("provider")), "Internal") {
|
|
req.Header.Set("redirect-for-internal", "true")
|
|
}
|
|
}
|
|
|
|
func resolveKiroPayloadProfileArn(account *Account) string {
|
|
if account == nil {
|
|
return ""
|
|
}
|
|
return strings.TrimSpace(account.GetCredential("profile_arn"))
|
|
}
|
|
|
|
func newKiroJSONRequest(ctx context.Context, endpointURL string, payload []byte, token, accountKey, machineID, amzTarget string, account *Account) (*http.Request, error) {
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpointURL, bytes.NewReader(payload))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Accept", "*/*")
|
|
req.Header.Set("Authorization", "Bearer "+token)
|
|
req.Header.Set("User-Agent", kiropkg.BuildRuntimeUserAgent(accountKey, machineID))
|
|
req.Header.Set("X-Amz-User-Agent", kiropkg.BuildRuntimeAmzUserAgent(accountKey, machineID))
|
|
req.Header.Set("x-amzn-kiro-agent-mode", "vibe")
|
|
req.Header.Set("x-amzn-codewhisperer-optout", "true")
|
|
req.Header.Set("Amz-Sdk-Request", "attempt=1; max=3")
|
|
req.Header.Set("Amz-Sdk-Invocation-Id", uuid.NewString())
|
|
if amzTarget != "" {
|
|
req.Header.Set("X-Amz-Target", amzTarget)
|
|
}
|
|
if account != nil {
|
|
profileArn := strings.TrimSpace(account.GetCredential("profile_arn"))
|
|
if profileArn != "" {
|
|
req.Header.Set("x-amzn-kiro-profile-arn", profileArn)
|
|
}
|
|
}
|
|
applyKiroConditionalHeaders(req, account)
|
|
return req, nil
|
|
}
|