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 { 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 }