mirror of
https://github.com/router-for-me/CLIProxyAPIPlus.git
synced 2026-03-08 06:43:41 +00:00
feat(kiro): add AWS Builder ID authentication support
- Add --kiro-aws-login flag for AWS Builder ID device code flow - Add DoKiroAWSLogin function for AWS SSO OIDC authentication - Complete Kiro integration with AWS, Google OAuth, and social auth - Add kiro executor, translator, and SDK components - Update browser support for Kiro authentication flows
This commit is contained in:
2
.gitignore
vendored
2
.gitignore
vendored
@@ -1,5 +1,6 @@
|
||||
# Binaries
|
||||
cli-proxy-api
|
||||
cliproxy
|
||||
*.exe
|
||||
|
||||
# Configuration
|
||||
@@ -31,6 +32,7 @@ GEMINI.md
|
||||
.vscode/*
|
||||
.claude/*
|
||||
.serena/*
|
||||
.mcp/cache/
|
||||
|
||||
# macOS
|
||||
.DS_Store
|
||||
|
||||
@@ -62,10 +62,16 @@ func main() {
|
||||
var iflowCookie bool
|
||||
var noBrowser bool
|
||||
var antigravityLogin bool
|
||||
var kiroLogin bool
|
||||
var kiroGoogleLogin bool
|
||||
var kiroAWSLogin bool
|
||||
var kiroImport bool
|
||||
var projectID string
|
||||
var vertexImport string
|
||||
var configPath string
|
||||
var password string
|
||||
var noIncognito bool
|
||||
var useIncognito bool
|
||||
|
||||
// Define command-line flags for different operation modes.
|
||||
flag.BoolVar(&login, "login", false, "Login Google Account")
|
||||
@@ -75,7 +81,13 @@ func main() {
|
||||
flag.BoolVar(&iflowLogin, "iflow-login", false, "Login to iFlow using OAuth")
|
||||
flag.BoolVar(&iflowCookie, "iflow-cookie", false, "Login to iFlow using Cookie")
|
||||
flag.BoolVar(&noBrowser, "no-browser", false, "Don't open browser automatically for OAuth")
|
||||
flag.BoolVar(&useIncognito, "incognito", false, "Open browser in incognito/private mode for OAuth (useful for multiple accounts)")
|
||||
flag.BoolVar(&noIncognito, "no-incognito", false, "Force disable incognito mode (uses existing browser session)")
|
||||
flag.BoolVar(&antigravityLogin, "antigravity-login", false, "Login to Antigravity using OAuth")
|
||||
flag.BoolVar(&kiroLogin, "kiro-login", false, "Login to Kiro using Google OAuth")
|
||||
flag.BoolVar(&kiroGoogleLogin, "kiro-google-login", false, "Login to Kiro using Google OAuth (same as --kiro-login)")
|
||||
flag.BoolVar(&kiroAWSLogin, "kiro-aws-login", false, "Login to Kiro using AWS Builder ID (device code flow)")
|
||||
flag.BoolVar(&kiroImport, "kiro-import", false, "Import Kiro token from Kiro IDE (~/.aws/sso/cache/kiro-auth-token.json)")
|
||||
flag.StringVar(&projectID, "project_id", "", "Project ID (Gemini only, not required)")
|
||||
flag.StringVar(&configPath, "config", DefaultConfigPath, "Configure File Path")
|
||||
flag.StringVar(&vertexImport, "vertex-import", "", "Import Vertex service account key JSON file")
|
||||
@@ -448,6 +460,44 @@ func main() {
|
||||
cmd.DoIFlowLogin(cfg, options)
|
||||
} else if iflowCookie {
|
||||
cmd.DoIFlowCookieAuth(cfg, options)
|
||||
} else if kiroLogin {
|
||||
// For Kiro auth, default to incognito mode for multi-account support
|
||||
// Users can explicitly override with --no-incognito
|
||||
// Note: This config mutation is safe - auth commands exit after completion
|
||||
// and don't share config with StartService (which is in the else branch)
|
||||
if useIncognito {
|
||||
cfg.IncognitoBrowser = true
|
||||
} else if noIncognito {
|
||||
cfg.IncognitoBrowser = false
|
||||
} else {
|
||||
cfg.IncognitoBrowser = true // Kiro default
|
||||
}
|
||||
cmd.DoKiroLogin(cfg, options)
|
||||
} else if kiroGoogleLogin {
|
||||
// For Kiro auth, default to incognito mode for multi-account support
|
||||
// Users can explicitly override with --no-incognito
|
||||
// Note: This config mutation is safe - auth commands exit after completion
|
||||
if useIncognito {
|
||||
cfg.IncognitoBrowser = true
|
||||
} else if noIncognito {
|
||||
cfg.IncognitoBrowser = false
|
||||
} else {
|
||||
cfg.IncognitoBrowser = true // Kiro default
|
||||
}
|
||||
cmd.DoKiroGoogleLogin(cfg, options)
|
||||
} else if kiroAWSLogin {
|
||||
// For Kiro auth, default to incognito mode for multi-account support
|
||||
// Users can explicitly override with --no-incognito
|
||||
if useIncognito {
|
||||
cfg.IncognitoBrowser = true
|
||||
} else if noIncognito {
|
||||
cfg.IncognitoBrowser = false
|
||||
} else {
|
||||
cfg.IncognitoBrowser = true // Kiro default
|
||||
}
|
||||
cmd.DoKiroAWSLogin(cfg, options)
|
||||
} else if kiroImport {
|
||||
cmd.DoKiroImport(cfg, options)
|
||||
} else {
|
||||
// In cloud deploy mode without config file, just wait for shutdown signals
|
||||
if isCloudDeploy && !configFileExists {
|
||||
|
||||
@@ -32,6 +32,11 @@ api-keys:
|
||||
# Enable debug logging
|
||||
debug: false
|
||||
|
||||
# Open OAuth URLs in incognito/private browser mode.
|
||||
# Useful when you want to login with a different account without logging out from your current session.
|
||||
# Default: false (but Kiro auth defaults to true for multi-account support)
|
||||
incognito-browser: true
|
||||
|
||||
# When true, write application logs to rotating files instead of stdout
|
||||
logging-to-file: false
|
||||
|
||||
@@ -99,6 +104,16 @@ ws-auth: false
|
||||
# - "*-think" # wildcard matching suffix (e.g. claude-opus-4-5-thinking)
|
||||
# - "*haiku*" # wildcard matching substring (e.g. claude-3-5-haiku-20241022)
|
||||
|
||||
# Kiro (AWS CodeWhisperer) configuration
|
||||
# Note: Kiro API currently only operates in us-east-1 region
|
||||
#kiro:
|
||||
# - token-file: "~/.aws/sso/cache/kiro-auth-token.json" # path to Kiro token file
|
||||
# agent-task-type: "" # optional: "vibe" or empty (API default)
|
||||
# - access-token: "aoaAAAAA..." # or provide tokens directly
|
||||
# refresh-token: "aorAAAAA..."
|
||||
# profile-arn: "arn:aws:codewhisperer:us-east-1:..."
|
||||
# proxy-url: "socks5://proxy.example.com:1080" # optional: proxy override
|
||||
|
||||
# OpenAI compatibility providers
|
||||
# openai-compatibility:
|
||||
# - name: "openrouter" # The name of the provider; it will be used in the user agent and other places.
|
||||
|
||||
3
go.mod
3
go.mod
@@ -13,14 +13,15 @@ require (
|
||||
github.com/joho/godotenv v1.5.1
|
||||
github.com/klauspost/compress v1.17.4
|
||||
github.com/minio/minio-go/v7 v7.0.66
|
||||
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c
|
||||
github.com/sirupsen/logrus v1.9.3
|
||||
github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966
|
||||
github.com/tidwall/gjson v1.18.0
|
||||
github.com/tidwall/sjson v1.2.5
|
||||
github.com/tiktoken-go/tokenizer v0.7.0
|
||||
golang.org/x/crypto v0.43.0
|
||||
golang.org/x/net v0.46.0
|
||||
golang.org/x/oauth2 v0.30.0
|
||||
golang.org/x/term v0.36.0
|
||||
gopkg.in/natefinch/lumberjack.v2 v2.2.1
|
||||
gopkg.in/yaml.v3 v3.0.1
|
||||
)
|
||||
|
||||
5
go.sum
5
go.sum
@@ -116,6 +116,8 @@ github.com/pelletier/go-toml/v2 v2.2.2 h1:aYUidT7k73Pcl9nb2gScu7NSrKCSHIDE89b3+6
|
||||
github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs=
|
||||
github.com/pjbgf/sha1cd v0.5.0 h1:a+UkboSi1znleCDUNT3M5YxjOnN1fz2FhN48FlwCxs0=
|
||||
github.com/pjbgf/sha1cd v0.5.0/go.mod h1:lhpGlyHLpQZoxMv8HcgXvZEhcGs0PG/vsZnEJ7H0iCM=
|
||||
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ=
|
||||
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU=
|
||||
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
|
||||
@@ -126,8 +128,6 @@ github.com/sergi/go-diff v1.4.0 h1:n/SP9D5ad1fORl+llWyN+D6qoUETXNZARKjyY2/KVCw=
|
||||
github.com/sergi/go-diff v1.4.0/go.mod h1:A0bzQcvG0E7Rwjx0REVgAGH58e96+X0MeOfepqsbeW4=
|
||||
github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
|
||||
github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
|
||||
github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966 h1:JIAuq3EEf9cgbU6AtGPK4CTG3Zf6CKMNqf0MHTggAUA=
|
||||
github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966/go.mod h1:sUM3LWHvSMaG192sy56D9F7CNvL7jUJVXoqM1QKLnog=
|
||||
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
|
||||
github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
|
||||
github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
|
||||
@@ -169,6 +169,7 @@ golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKl
|
||||
golang.org/x/sync v0.17.0 h1:l60nONMj9l5drqw6jlhIELNv9I0A4OFgRsG9k2oT9Ug=
|
||||
golang.org/x/sync v0.17.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
|
||||
golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.37.0 h1:fdNQudmxPjkdUTPnLn5mdQv7Zwvbvpaxqs831goi9kQ=
|
||||
golang.org/x/sys v0.37.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
|
||||
|
||||
@@ -36,9 +36,32 @@ import (
|
||||
)
|
||||
|
||||
var (
|
||||
oauthStatus = make(map[string]string)
|
||||
oauthStatus = make(map[string]string)
|
||||
oauthStatusMutex sync.RWMutex
|
||||
)
|
||||
|
||||
// getOAuthStatus safely retrieves an OAuth status
|
||||
func getOAuthStatus(key string) (string, bool) {
|
||||
oauthStatusMutex.RLock()
|
||||
defer oauthStatusMutex.RUnlock()
|
||||
status, ok := oauthStatus[key]
|
||||
return status, ok
|
||||
}
|
||||
|
||||
// setOAuthStatus safely sets an OAuth status
|
||||
func setOAuthStatus(key string, status string) {
|
||||
oauthStatusMutex.Lock()
|
||||
defer oauthStatusMutex.Unlock()
|
||||
oauthStatus[key] = status
|
||||
}
|
||||
|
||||
// deleteOAuthStatus safely deletes an OAuth status
|
||||
func deleteOAuthStatus(key string) {
|
||||
oauthStatusMutex.Lock()
|
||||
defer oauthStatusMutex.Unlock()
|
||||
delete(oauthStatus, key)
|
||||
}
|
||||
|
||||
var lastRefreshKeys = []string{"last_refresh", "lastRefresh", "last_refreshed_at", "lastRefreshedAt"}
|
||||
|
||||
const (
|
||||
@@ -760,7 +783,7 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
|
||||
deadline := time.Now().Add(timeout)
|
||||
for {
|
||||
if time.Now().After(deadline) {
|
||||
oauthStatus[state] = "Timeout waiting for OAuth callback"
|
||||
setOAuthStatus(state, "Timeout waiting for OAuth callback")
|
||||
return nil, fmt.Errorf("timeout waiting for OAuth callback")
|
||||
}
|
||||
data, errRead := os.ReadFile(path)
|
||||
@@ -785,13 +808,13 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
|
||||
if errStr := resultMap["error"]; errStr != "" {
|
||||
oauthErr := claude.NewOAuthError(errStr, "", http.StatusBadRequest)
|
||||
log.Error(claude.GetUserFriendlyMessage(oauthErr))
|
||||
oauthStatus[state] = "Bad request"
|
||||
setOAuthStatus(state, "Bad request")
|
||||
return
|
||||
}
|
||||
if resultMap["state"] != state {
|
||||
authErr := claude.NewAuthenticationError(claude.ErrInvalidState, fmt.Errorf("expected %s, got %s", state, resultMap["state"]))
|
||||
log.Error(claude.GetUserFriendlyMessage(authErr))
|
||||
oauthStatus[state] = "State code error"
|
||||
setOAuthStatus(state, "State code error")
|
||||
return
|
||||
}
|
||||
|
||||
@@ -824,7 +847,7 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
|
||||
if errDo != nil {
|
||||
authErr := claude.NewAuthenticationError(claude.ErrCodeExchangeFailed, errDo)
|
||||
log.Errorf("Failed to exchange authorization code for tokens: %v", authErr)
|
||||
oauthStatus[state] = "Failed to exchange authorization code for tokens"
|
||||
setOAuthStatus(state, "Failed to exchange authorization code for tokens")
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
@@ -835,7 +858,7 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
log.Errorf("token exchange failed with status %d: %s", resp.StatusCode, string(respBody))
|
||||
oauthStatus[state] = fmt.Sprintf("token exchange failed with status %d", resp.StatusCode)
|
||||
setOAuthStatus(state, fmt.Sprintf("token exchange failed with status %d", resp.StatusCode))
|
||||
return
|
||||
}
|
||||
var tResp struct {
|
||||
@@ -848,7 +871,7 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
|
||||
}
|
||||
if errU := json.Unmarshal(respBody, &tResp); errU != nil {
|
||||
log.Errorf("failed to parse token response: %v", errU)
|
||||
oauthStatus[state] = "Failed to parse token response"
|
||||
setOAuthStatus(state, "Failed to parse token response")
|
||||
return
|
||||
}
|
||||
bundle := &claude.ClaudeAuthBundle{
|
||||
@@ -873,7 +896,7 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
|
||||
savedPath, errSave := h.saveTokenRecord(ctx, record)
|
||||
if errSave != nil {
|
||||
log.Fatalf("Failed to save authentication tokens: %v", errSave)
|
||||
oauthStatus[state] = "Failed to save authentication tokens"
|
||||
setOAuthStatus(state, "Failed to save authentication tokens")
|
||||
return
|
||||
}
|
||||
|
||||
@@ -882,10 +905,10 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
|
||||
fmt.Println("API key obtained and saved")
|
||||
}
|
||||
fmt.Println("You can now use Claude services through this CLI")
|
||||
delete(oauthStatus, state)
|
||||
deleteOAuthStatus(state)
|
||||
}()
|
||||
|
||||
oauthStatus[state] = ""
|
||||
setOAuthStatus(state, "")
|
||||
c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state})
|
||||
}
|
||||
|
||||
@@ -944,7 +967,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
|
||||
for {
|
||||
if time.Now().After(deadline) {
|
||||
log.Error("oauth flow timed out")
|
||||
oauthStatus[state] = "OAuth flow timed out"
|
||||
setOAuthStatus(state, "OAuth flow timed out")
|
||||
return
|
||||
}
|
||||
if data, errR := os.ReadFile(waitFile); errR == nil {
|
||||
@@ -953,13 +976,13 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
|
||||
_ = os.Remove(waitFile)
|
||||
if errStr := m["error"]; errStr != "" {
|
||||
log.Errorf("Authentication failed: %s", errStr)
|
||||
oauthStatus[state] = "Authentication failed"
|
||||
setOAuthStatus(state, "Authentication failed")
|
||||
return
|
||||
}
|
||||
authCode = m["code"]
|
||||
if authCode == "" {
|
||||
log.Errorf("Authentication failed: code not found")
|
||||
oauthStatus[state] = "Authentication failed: code not found"
|
||||
setOAuthStatus(state, "Authentication failed: code not found")
|
||||
return
|
||||
}
|
||||
break
|
||||
@@ -971,7 +994,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
|
||||
token, err := conf.Exchange(ctx, authCode)
|
||||
if err != nil {
|
||||
log.Errorf("Failed to exchange token: %v", err)
|
||||
oauthStatus[state] = "Failed to exchange token"
|
||||
setOAuthStatus(state, "Failed to exchange token")
|
||||
return
|
||||
}
|
||||
|
||||
@@ -982,7 +1005,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
|
||||
req, errNewRequest := http.NewRequestWithContext(ctx, "GET", "https://www.googleapis.com/oauth2/v1/userinfo?alt=json", nil)
|
||||
if errNewRequest != nil {
|
||||
log.Errorf("Could not get user info: %v", errNewRequest)
|
||||
oauthStatus[state] = "Could not get user info"
|
||||
setOAuthStatus(state, "Could not get user info")
|
||||
return
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
@@ -991,7 +1014,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
|
||||
resp, errDo := authHTTPClient.Do(req)
|
||||
if errDo != nil {
|
||||
log.Errorf("Failed to execute request: %v", errDo)
|
||||
oauthStatus[state] = "Failed to execute request"
|
||||
setOAuthStatus(state, "Failed to execute request")
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
@@ -1003,7 +1026,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
|
||||
bodyBytes, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
log.Errorf("Get user info request failed with status %d: %s", resp.StatusCode, string(bodyBytes))
|
||||
oauthStatus[state] = fmt.Sprintf("Get user info request failed with status %d", resp.StatusCode)
|
||||
setOAuthStatus(state, fmt.Sprintf("Get user info request failed with status %d", resp.StatusCode))
|
||||
return
|
||||
}
|
||||
|
||||
@@ -1012,7 +1035,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
|
||||
fmt.Printf("Authenticated user email: %s\n", email)
|
||||
} else {
|
||||
fmt.Println("Failed to get user email from token")
|
||||
oauthStatus[state] = "Failed to get user email from token"
|
||||
setOAuthStatus(state, "Failed to get user email from token")
|
||||
}
|
||||
|
||||
// Marshal/unmarshal oauth2.Token to generic map and enrich fields
|
||||
@@ -1020,7 +1043,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
|
||||
jsonData, _ := json.Marshal(token)
|
||||
if errUnmarshal := json.Unmarshal(jsonData, &ifToken); errUnmarshal != nil {
|
||||
log.Errorf("Failed to unmarshal token: %v", errUnmarshal)
|
||||
oauthStatus[state] = "Failed to unmarshal token"
|
||||
setOAuthStatus(state, "Failed to unmarshal token")
|
||||
return
|
||||
}
|
||||
|
||||
@@ -1046,7 +1069,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
|
||||
gemClient, errGetClient := gemAuth.GetAuthenticatedClient(ctx, &ts, h.cfg, true)
|
||||
if errGetClient != nil {
|
||||
log.Fatalf("failed to get authenticated client: %v", errGetClient)
|
||||
oauthStatus[state] = "Failed to get authenticated client"
|
||||
setOAuthStatus(state, "Failed to get authenticated client")
|
||||
return
|
||||
}
|
||||
fmt.Println("Authentication successful.")
|
||||
@@ -1056,12 +1079,12 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
|
||||
projects, errAll := onboardAllGeminiProjects(ctx, gemClient, &ts)
|
||||
if errAll != nil {
|
||||
log.Errorf("Failed to complete Gemini CLI onboarding: %v", errAll)
|
||||
oauthStatus[state] = "Failed to complete Gemini CLI onboarding"
|
||||
setOAuthStatus(state, "Failed to complete Gemini CLI onboarding")
|
||||
return
|
||||
}
|
||||
if errVerify := ensureGeminiProjectsEnabled(ctx, gemClient, projects); errVerify != nil {
|
||||
log.Errorf("Failed to verify Cloud AI API status: %v", errVerify)
|
||||
oauthStatus[state] = "Failed to verify Cloud AI API status"
|
||||
setOAuthStatus(state, "Failed to verify Cloud AI API status")
|
||||
return
|
||||
}
|
||||
ts.ProjectID = strings.Join(projects, ",")
|
||||
@@ -1069,26 +1092,26 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
|
||||
} else {
|
||||
if errEnsure := ensureGeminiProjectAndOnboard(ctx, gemClient, &ts, requestedProjectID); errEnsure != nil {
|
||||
log.Errorf("Failed to complete Gemini CLI onboarding: %v", errEnsure)
|
||||
oauthStatus[state] = "Failed to complete Gemini CLI onboarding"
|
||||
setOAuthStatus(state, "Failed to complete Gemini CLI onboarding")
|
||||
return
|
||||
}
|
||||
|
||||
if strings.TrimSpace(ts.ProjectID) == "" {
|
||||
log.Error("Onboarding did not return a project ID")
|
||||
oauthStatus[state] = "Failed to resolve project ID"
|
||||
setOAuthStatus(state, "Failed to resolve project ID")
|
||||
return
|
||||
}
|
||||
|
||||
isChecked, errCheck := checkCloudAPIIsEnabled(ctx, gemClient, ts.ProjectID)
|
||||
if errCheck != nil {
|
||||
log.Errorf("Failed to verify Cloud AI API status: %v", errCheck)
|
||||
oauthStatus[state] = "Failed to verify Cloud AI API status"
|
||||
setOAuthStatus(state, "Failed to verify Cloud AI API status")
|
||||
return
|
||||
}
|
||||
ts.Checked = isChecked
|
||||
if !isChecked {
|
||||
log.Error("Cloud AI API is not enabled for the selected project")
|
||||
oauthStatus[state] = "Cloud AI API not enabled"
|
||||
setOAuthStatus(state, "Cloud AI API not enabled")
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -1111,15 +1134,15 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
|
||||
savedPath, errSave := h.saveTokenRecord(ctx, record)
|
||||
if errSave != nil {
|
||||
log.Fatalf("Failed to save token to file: %v", errSave)
|
||||
oauthStatus[state] = "Failed to save token to file"
|
||||
setOAuthStatus(state, "Failed to save token to file")
|
||||
return
|
||||
}
|
||||
|
||||
delete(oauthStatus, state)
|
||||
deleteOAuthStatus(state)
|
||||
fmt.Printf("You can now use Gemini CLI services through this CLI; token saved to %s\n", savedPath)
|
||||
}()
|
||||
|
||||
oauthStatus[state] = ""
|
||||
setOAuthStatus(state, "")
|
||||
c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state})
|
||||
}
|
||||
|
||||
@@ -1180,7 +1203,7 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
|
||||
if time.Now().After(deadline) {
|
||||
authErr := codex.NewAuthenticationError(codex.ErrCallbackTimeout, fmt.Errorf("timeout waiting for OAuth callback"))
|
||||
log.Error(codex.GetUserFriendlyMessage(authErr))
|
||||
oauthStatus[state] = "Timeout waiting for OAuth callback"
|
||||
setOAuthStatus(state, "Timeout waiting for OAuth callback")
|
||||
return
|
||||
}
|
||||
if data, errR := os.ReadFile(waitFile); errR == nil {
|
||||
@@ -1190,12 +1213,12 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
|
||||
if errStr := m["error"]; errStr != "" {
|
||||
oauthErr := codex.NewOAuthError(errStr, "", http.StatusBadRequest)
|
||||
log.Error(codex.GetUserFriendlyMessage(oauthErr))
|
||||
oauthStatus[state] = "Bad Request"
|
||||
setOAuthStatus(state, "Bad Request")
|
||||
return
|
||||
}
|
||||
if m["state"] != state {
|
||||
authErr := codex.NewAuthenticationError(codex.ErrInvalidState, fmt.Errorf("expected %s, got %s", state, m["state"]))
|
||||
oauthStatus[state] = "State code error"
|
||||
setOAuthStatus(state, "State code error")
|
||||
log.Error(codex.GetUserFriendlyMessage(authErr))
|
||||
return
|
||||
}
|
||||
@@ -1226,14 +1249,14 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
|
||||
resp, errDo := httpClient.Do(req)
|
||||
if errDo != nil {
|
||||
authErr := codex.NewAuthenticationError(codex.ErrCodeExchangeFailed, errDo)
|
||||
oauthStatus[state] = "Failed to exchange authorization code for tokens"
|
||||
setOAuthStatus(state, "Failed to exchange authorization code for tokens")
|
||||
log.Errorf("Failed to exchange authorization code for tokens: %v", authErr)
|
||||
return
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
oauthStatus[state] = fmt.Sprintf("Token exchange failed with status %d", resp.StatusCode)
|
||||
setOAuthStatus(state, fmt.Sprintf("Token exchange failed with status %d", resp.StatusCode))
|
||||
log.Errorf("token exchange failed with status %d: %s", resp.StatusCode, string(respBody))
|
||||
return
|
||||
}
|
||||
@@ -1244,7 +1267,7 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
|
||||
ExpiresIn int `json:"expires_in"`
|
||||
}
|
||||
if errU := json.Unmarshal(respBody, &tokenResp); errU != nil {
|
||||
oauthStatus[state] = "Failed to parse token response"
|
||||
setOAuthStatus(state, "Failed to parse token response")
|
||||
log.Errorf("failed to parse token response: %v", errU)
|
||||
return
|
||||
}
|
||||
@@ -1282,7 +1305,7 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
|
||||
}
|
||||
savedPath, errSave := h.saveTokenRecord(ctx, record)
|
||||
if errSave != nil {
|
||||
oauthStatus[state] = "Failed to save authentication tokens"
|
||||
setOAuthStatus(state, "Failed to save authentication tokens")
|
||||
log.Fatalf("Failed to save authentication tokens: %v", errSave)
|
||||
return
|
||||
}
|
||||
@@ -1291,10 +1314,10 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
|
||||
fmt.Println("API key obtained and saved")
|
||||
}
|
||||
fmt.Println("You can now use Codex services through this CLI")
|
||||
delete(oauthStatus, state)
|
||||
deleteOAuthStatus(state)
|
||||
}()
|
||||
|
||||
oauthStatus[state] = ""
|
||||
setOAuthStatus(state, "")
|
||||
c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state})
|
||||
}
|
||||
|
||||
@@ -1360,7 +1383,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
|
||||
for {
|
||||
if time.Now().After(deadline) {
|
||||
log.Error("oauth flow timed out")
|
||||
oauthStatus[state] = "OAuth flow timed out"
|
||||
setOAuthStatus(state, "OAuth flow timed out")
|
||||
return
|
||||
}
|
||||
if data, errReadFile := os.ReadFile(waitFile); errReadFile == nil {
|
||||
@@ -1369,18 +1392,18 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
|
||||
_ = os.Remove(waitFile)
|
||||
if errStr := strings.TrimSpace(payload["error"]); errStr != "" {
|
||||
log.Errorf("Authentication failed: %s", errStr)
|
||||
oauthStatus[state] = "Authentication failed"
|
||||
setOAuthStatus(state, "Authentication failed")
|
||||
return
|
||||
}
|
||||
if payloadState := strings.TrimSpace(payload["state"]); payloadState != "" && payloadState != state {
|
||||
log.Errorf("Authentication failed: state mismatch")
|
||||
oauthStatus[state] = "Authentication failed: state mismatch"
|
||||
setOAuthStatus(state, "Authentication failed: state mismatch")
|
||||
return
|
||||
}
|
||||
authCode = strings.TrimSpace(payload["code"])
|
||||
if authCode == "" {
|
||||
log.Error("Authentication failed: code not found")
|
||||
oauthStatus[state] = "Authentication failed: code not found"
|
||||
setOAuthStatus(state, "Authentication failed: code not found")
|
||||
return
|
||||
}
|
||||
break
|
||||
@@ -1399,7 +1422,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
|
||||
req, errNewRequest := http.NewRequestWithContext(ctx, http.MethodPost, "https://oauth2.googleapis.com/token", strings.NewReader(form.Encode()))
|
||||
if errNewRequest != nil {
|
||||
log.Errorf("Failed to build token request: %v", errNewRequest)
|
||||
oauthStatus[state] = "Failed to build token request"
|
||||
setOAuthStatus(state, "Failed to build token request")
|
||||
return
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
@@ -1407,7 +1430,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
|
||||
resp, errDo := httpClient.Do(req)
|
||||
if errDo != nil {
|
||||
log.Errorf("Failed to execute token request: %v", errDo)
|
||||
oauthStatus[state] = "Failed to exchange token"
|
||||
setOAuthStatus(state, "Failed to exchange token")
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
@@ -1419,7 +1442,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
|
||||
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
|
||||
bodyBytes, _ := io.ReadAll(resp.Body)
|
||||
log.Errorf("Antigravity token exchange failed with status %d: %s", resp.StatusCode, string(bodyBytes))
|
||||
oauthStatus[state] = fmt.Sprintf("Token exchange failed: %d", resp.StatusCode)
|
||||
setOAuthStatus(state, fmt.Sprintf("Token exchange failed: %d", resp.StatusCode))
|
||||
return
|
||||
}
|
||||
|
||||
@@ -1431,7 +1454,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
|
||||
}
|
||||
if errDecode := json.NewDecoder(resp.Body).Decode(&tokenResp); errDecode != nil {
|
||||
log.Errorf("Failed to parse token response: %v", errDecode)
|
||||
oauthStatus[state] = "Failed to parse token response"
|
||||
setOAuthStatus(state, "Failed to parse token response")
|
||||
return
|
||||
}
|
||||
|
||||
@@ -1440,7 +1463,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
|
||||
infoReq, errInfoReq := http.NewRequestWithContext(ctx, http.MethodGet, "https://www.googleapis.com/oauth2/v1/userinfo?alt=json", nil)
|
||||
if errInfoReq != nil {
|
||||
log.Errorf("Failed to build user info request: %v", errInfoReq)
|
||||
oauthStatus[state] = "Failed to build user info request"
|
||||
setOAuthStatus(state, "Failed to build user info request")
|
||||
return
|
||||
}
|
||||
infoReq.Header.Set("Authorization", "Bearer "+tokenResp.AccessToken)
|
||||
@@ -1448,7 +1471,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
|
||||
infoResp, errInfo := httpClient.Do(infoReq)
|
||||
if errInfo != nil {
|
||||
log.Errorf("Failed to execute user info request: %v", errInfo)
|
||||
oauthStatus[state] = "Failed to execute user info request"
|
||||
setOAuthStatus(state, "Failed to execute user info request")
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
@@ -1467,7 +1490,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
|
||||
} else {
|
||||
bodyBytes, _ := io.ReadAll(infoResp.Body)
|
||||
log.Errorf("User info request failed with status %d: %s", infoResp.StatusCode, string(bodyBytes))
|
||||
oauthStatus[state] = fmt.Sprintf("User info request failed: %d", infoResp.StatusCode)
|
||||
setOAuthStatus(state, fmt.Sprintf("User info request failed: %d", infoResp.StatusCode))
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -1515,11 +1538,11 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
|
||||
savedPath, errSave := h.saveTokenRecord(ctx, record)
|
||||
if errSave != nil {
|
||||
log.Fatalf("Failed to save token to file: %v", errSave)
|
||||
oauthStatus[state] = "Failed to save token to file"
|
||||
setOAuthStatus(state, "Failed to save token to file")
|
||||
return
|
||||
}
|
||||
|
||||
delete(oauthStatus, state)
|
||||
deleteOAuthStatus(state)
|
||||
fmt.Printf("Authentication successful! Token saved to %s\n", savedPath)
|
||||
if projectID != "" {
|
||||
fmt.Printf("Using GCP project: %s\n", projectID)
|
||||
@@ -1527,7 +1550,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
|
||||
fmt.Println("You can now use Antigravity services through this CLI")
|
||||
}()
|
||||
|
||||
oauthStatus[state] = ""
|
||||
setOAuthStatus(state, "")
|
||||
c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state})
|
||||
}
|
||||
|
||||
@@ -1552,7 +1575,7 @@ func (h *Handler) RequestQwenToken(c *gin.Context) {
|
||||
fmt.Println("Waiting for authentication...")
|
||||
tokenData, errPollForToken := qwenAuth.PollForToken(deviceFlow.DeviceCode, deviceFlow.CodeVerifier)
|
||||
if errPollForToken != nil {
|
||||
oauthStatus[state] = "Authentication failed"
|
||||
setOAuthStatus(state, "Authentication failed")
|
||||
fmt.Printf("Authentication failed: %v\n", errPollForToken)
|
||||
return
|
||||
}
|
||||
@@ -1571,16 +1594,16 @@ func (h *Handler) RequestQwenToken(c *gin.Context) {
|
||||
savedPath, errSave := h.saveTokenRecord(ctx, record)
|
||||
if errSave != nil {
|
||||
log.Fatalf("Failed to save authentication tokens: %v", errSave)
|
||||
oauthStatus[state] = "Failed to save authentication tokens"
|
||||
setOAuthStatus(state, "Failed to save authentication tokens")
|
||||
return
|
||||
}
|
||||
|
||||
fmt.Printf("Authentication successful! Token saved to %s\n", savedPath)
|
||||
fmt.Println("You can now use Qwen services through this CLI")
|
||||
delete(oauthStatus, state)
|
||||
deleteOAuthStatus(state)
|
||||
}()
|
||||
|
||||
oauthStatus[state] = ""
|
||||
setOAuthStatus(state, "")
|
||||
c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state})
|
||||
}
|
||||
|
||||
@@ -1619,7 +1642,7 @@ func (h *Handler) RequestIFlowToken(c *gin.Context) {
|
||||
var resultMap map[string]string
|
||||
for {
|
||||
if time.Now().After(deadline) {
|
||||
oauthStatus[state] = "Authentication failed"
|
||||
setOAuthStatus(state, "Authentication failed")
|
||||
fmt.Println("Authentication failed: timeout waiting for callback")
|
||||
return
|
||||
}
|
||||
@@ -1632,26 +1655,26 @@ func (h *Handler) RequestIFlowToken(c *gin.Context) {
|
||||
}
|
||||
|
||||
if errStr := strings.TrimSpace(resultMap["error"]); errStr != "" {
|
||||
oauthStatus[state] = "Authentication failed"
|
||||
setOAuthStatus(state, "Authentication failed")
|
||||
fmt.Printf("Authentication failed: %s\n", errStr)
|
||||
return
|
||||
}
|
||||
if resultState := strings.TrimSpace(resultMap["state"]); resultState != state {
|
||||
oauthStatus[state] = "Authentication failed"
|
||||
setOAuthStatus(state, "Authentication failed")
|
||||
fmt.Println("Authentication failed: state mismatch")
|
||||
return
|
||||
}
|
||||
|
||||
code := strings.TrimSpace(resultMap["code"])
|
||||
if code == "" {
|
||||
oauthStatus[state] = "Authentication failed"
|
||||
setOAuthStatus(state, "Authentication failed")
|
||||
fmt.Println("Authentication failed: code missing")
|
||||
return
|
||||
}
|
||||
|
||||
tokenData, errExchange := authSvc.ExchangeCodeForTokens(ctx, code, redirectURI)
|
||||
if errExchange != nil {
|
||||
oauthStatus[state] = "Authentication failed"
|
||||
setOAuthStatus(state, "Authentication failed")
|
||||
fmt.Printf("Authentication failed: %v\n", errExchange)
|
||||
return
|
||||
}
|
||||
@@ -1673,7 +1696,7 @@ func (h *Handler) RequestIFlowToken(c *gin.Context) {
|
||||
|
||||
savedPath, errSave := h.saveTokenRecord(ctx, record)
|
||||
if errSave != nil {
|
||||
oauthStatus[state] = "Failed to save authentication tokens"
|
||||
setOAuthStatus(state, "Failed to save authentication tokens")
|
||||
log.Fatalf("Failed to save authentication tokens: %v", errSave)
|
||||
return
|
||||
}
|
||||
@@ -1683,10 +1706,10 @@ func (h *Handler) RequestIFlowToken(c *gin.Context) {
|
||||
fmt.Println("API key obtained and saved")
|
||||
}
|
||||
fmt.Println("You can now use iFlow services through this CLI")
|
||||
delete(oauthStatus, state)
|
||||
deleteOAuthStatus(state)
|
||||
}()
|
||||
|
||||
oauthStatus[state] = ""
|
||||
setOAuthStatus(state, "")
|
||||
c.JSON(http.StatusOK, gin.H{"status": "ok", "url": authURL, "state": state})
|
||||
}
|
||||
|
||||
@@ -2110,7 +2133,7 @@ func checkCloudAPIIsEnabled(ctx context.Context, httpClient *http.Client, projec
|
||||
|
||||
func (h *Handler) GetAuthStatus(c *gin.Context) {
|
||||
state := c.Query("state")
|
||||
if err, ok := oauthStatus[state]; ok {
|
||||
if err, ok := getOAuthStatus(state); ok {
|
||||
if err != "" {
|
||||
c.JSON(200, gin.H{"status": "error", "error": err})
|
||||
} else {
|
||||
@@ -2120,5 +2143,5 @@ func (h *Handler) GetAuthStatus(c *gin.Context) {
|
||||
} else {
|
||||
c.JSON(200, gin.H{"status": "ok"})
|
||||
}
|
||||
delete(oauthStatus, state)
|
||||
deleteOAuthStatus(state)
|
||||
}
|
||||
|
||||
@@ -902,7 +902,7 @@ func (s *Server) UpdateClients(cfg *config.Config) {
|
||||
for _, p := range cfg.OpenAICompatibility {
|
||||
providerNames = append(providerNames, p.Name)
|
||||
}
|
||||
s.handlers.OpenAICompatProviders = providerNames
|
||||
s.handlers.SetOpenAICompatProviders(providerNames)
|
||||
|
||||
s.handlers.UpdateClients(&cfg.SDKConfig)
|
||||
|
||||
|
||||
@@ -242,6 +242,11 @@ func (s *OAuthServer) handleSuccess(w http.ResponseWriter, r *http.Request) {
|
||||
platformURL = "https://console.anthropic.com/"
|
||||
}
|
||||
|
||||
// Validate platformURL to prevent XSS - only allow http/https URLs
|
||||
if !isValidURL(platformURL) {
|
||||
platformURL = "https://console.anthropic.com/"
|
||||
}
|
||||
|
||||
// Generate success page HTML with dynamic content
|
||||
successHTML := s.generateSuccessHTML(setupRequired, platformURL)
|
||||
|
||||
@@ -251,6 +256,12 @@ func (s *OAuthServer) handleSuccess(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
}
|
||||
|
||||
// isValidURL checks if the URL is a valid http/https URL to prevent XSS
|
||||
func isValidURL(urlStr string) bool {
|
||||
urlStr = strings.TrimSpace(urlStr)
|
||||
return strings.HasPrefix(urlStr, "https://") || strings.HasPrefix(urlStr, "http://")
|
||||
}
|
||||
|
||||
// generateSuccessHTML creates the HTML content for the success page.
|
||||
// It customizes the page based on whether additional setup is required
|
||||
// and includes a link to the platform.
|
||||
|
||||
@@ -239,6 +239,11 @@ func (s *OAuthServer) handleSuccess(w http.ResponseWriter, r *http.Request) {
|
||||
platformURL = "https://platform.openai.com"
|
||||
}
|
||||
|
||||
// Validate platformURL to prevent XSS - only allow http/https URLs
|
||||
if !isValidURL(platformURL) {
|
||||
platformURL = "https://platform.openai.com"
|
||||
}
|
||||
|
||||
// Generate success page HTML with dynamic content
|
||||
successHTML := s.generateSuccessHTML(setupRequired, platformURL)
|
||||
|
||||
@@ -248,6 +253,12 @@ func (s *OAuthServer) handleSuccess(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
}
|
||||
|
||||
// isValidURL checks if the URL is a valid http/https URL to prevent XSS
|
||||
func isValidURL(urlStr string) bool {
|
||||
urlStr = strings.TrimSpace(urlStr)
|
||||
return strings.HasPrefix(urlStr, "https://") || strings.HasPrefix(urlStr, "http://")
|
||||
}
|
||||
|
||||
// generateSuccessHTML creates the HTML content for the success page.
|
||||
// It customizes the page based on whether additional setup is required
|
||||
// and includes a link to the platform.
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -28,10 +29,21 @@ const (
|
||||
iFlowAPIKeyEndpoint = "https://platform.iflow.cn/api/openapi/apikey"
|
||||
|
||||
// Client credentials provided by iFlow for the Code Assist integration.
|
||||
iFlowOAuthClientID = "10009311001"
|
||||
iFlowOAuthClientSecret = "4Z3YjXycVsQvyGF1etiNlIBB4RsqSDtW"
|
||||
iFlowOAuthClientID = "10009311001"
|
||||
// Default client secret (can be overridden via IFLOW_CLIENT_SECRET env var)
|
||||
defaultIFlowClientSecret = "4Z3YjXycVsQvyGF1etiNlIBB4RsqSDtW"
|
||||
)
|
||||
|
||||
// getIFlowClientSecret returns the iFlow OAuth client secret.
|
||||
// It first checks the IFLOW_CLIENT_SECRET environment variable,
|
||||
// falling back to the default value if not set.
|
||||
func getIFlowClientSecret() string {
|
||||
if secret := os.Getenv("IFLOW_CLIENT_SECRET"); secret != "" {
|
||||
return secret
|
||||
}
|
||||
return defaultIFlowClientSecret
|
||||
}
|
||||
|
||||
// DefaultAPIBaseURL is the canonical chat completions endpoint.
|
||||
const DefaultAPIBaseURL = "https://apis.iflow.cn/v1"
|
||||
|
||||
@@ -72,7 +84,7 @@ func (ia *IFlowAuth) ExchangeCodeForTokens(ctx context.Context, code, redirectUR
|
||||
form.Set("code", code)
|
||||
form.Set("redirect_uri", redirectURI)
|
||||
form.Set("client_id", iFlowOAuthClientID)
|
||||
form.Set("client_secret", iFlowOAuthClientSecret)
|
||||
form.Set("client_secret", getIFlowClientSecret())
|
||||
|
||||
req, err := ia.newTokenRequest(ctx, form)
|
||||
if err != nil {
|
||||
@@ -88,7 +100,7 @@ func (ia *IFlowAuth) RefreshTokens(ctx context.Context, refreshToken string) (*I
|
||||
form.Set("grant_type", "refresh_token")
|
||||
form.Set("refresh_token", refreshToken)
|
||||
form.Set("client_id", iFlowOAuthClientID)
|
||||
form.Set("client_secret", iFlowOAuthClientSecret)
|
||||
form.Set("client_secret", getIFlowClientSecret())
|
||||
|
||||
req, err := ia.newTokenRequest(ctx, form)
|
||||
if err != nil {
|
||||
@@ -104,7 +116,7 @@ func (ia *IFlowAuth) newTokenRequest(ctx context.Context, form url.Values) (*htt
|
||||
return nil, fmt.Errorf("iflow token: create request failed: %w", err)
|
||||
}
|
||||
|
||||
basic := base64.StdEncoding.EncodeToString([]byte(iFlowOAuthClientID + ":" + iFlowOAuthClientSecret))
|
||||
basic := base64.StdEncoding.EncodeToString([]byte(iFlowOAuthClientID + ":" + getIFlowClientSecret()))
|
||||
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("Authorization", "Basic "+basic)
|
||||
|
||||
301
internal/auth/kiro/aws.go
Normal file
301
internal/auth/kiro/aws.go
Normal file
@@ -0,0 +1,301 @@
|
||||
// Package kiro provides authentication functionality for AWS CodeWhisperer (Kiro) API.
|
||||
// It includes interfaces and implementations for token storage and authentication methods.
|
||||
package kiro
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// PKCECodes holds PKCE verification codes for OAuth2 PKCE flow
|
||||
type PKCECodes struct {
|
||||
// CodeVerifier is the cryptographically random string used to correlate
|
||||
// the authorization request to the token request
|
||||
CodeVerifier string `json:"code_verifier"`
|
||||
// CodeChallenge is the SHA256 hash of the code verifier, base64url-encoded
|
||||
CodeChallenge string `json:"code_challenge"`
|
||||
}
|
||||
|
||||
// KiroTokenData holds OAuth token information from AWS CodeWhisperer (Kiro)
|
||||
type KiroTokenData struct {
|
||||
// AccessToken is the OAuth2 access token for API access
|
||||
AccessToken string `json:"accessToken"`
|
||||
// RefreshToken is used to obtain new access tokens
|
||||
RefreshToken string `json:"refreshToken"`
|
||||
// ProfileArn is the AWS CodeWhisperer profile ARN
|
||||
ProfileArn string `json:"profileArn"`
|
||||
// ExpiresAt is the timestamp when the token expires
|
||||
ExpiresAt string `json:"expiresAt"`
|
||||
// AuthMethod indicates the authentication method used (e.g., "builder-id", "social")
|
||||
AuthMethod string `json:"authMethod"`
|
||||
// Provider indicates the OAuth provider (e.g., "AWS", "Google")
|
||||
Provider string `json:"provider"`
|
||||
// ClientID is the OIDC client ID (needed for token refresh)
|
||||
ClientID string `json:"clientId,omitempty"`
|
||||
// ClientSecret is the OIDC client secret (needed for token refresh)
|
||||
ClientSecret string `json:"clientSecret,omitempty"`
|
||||
// Email is the user's email address (used for file naming)
|
||||
Email string `json:"email,omitempty"`
|
||||
}
|
||||
|
||||
// KiroAuthBundle aggregates authentication data after OAuth flow completion
|
||||
type KiroAuthBundle struct {
|
||||
// TokenData contains the OAuth tokens from the authentication flow
|
||||
TokenData KiroTokenData `json:"token_data"`
|
||||
// LastRefresh is the timestamp of the last token refresh
|
||||
LastRefresh string `json:"last_refresh"`
|
||||
}
|
||||
|
||||
// KiroUsageInfo represents usage information from CodeWhisperer API
|
||||
type KiroUsageInfo struct {
|
||||
// SubscriptionTitle is the subscription plan name (e.g., "KIRO FREE")
|
||||
SubscriptionTitle string `json:"subscription_title"`
|
||||
// CurrentUsage is the current credit usage
|
||||
CurrentUsage float64 `json:"current_usage"`
|
||||
// UsageLimit is the maximum credit limit
|
||||
UsageLimit float64 `json:"usage_limit"`
|
||||
// NextReset is the timestamp of the next usage reset
|
||||
NextReset string `json:"next_reset"`
|
||||
}
|
||||
|
||||
// KiroModel represents a model available through the CodeWhisperer API
|
||||
type KiroModel struct {
|
||||
// ModelID is the unique identifier for the model
|
||||
ModelID string `json:"modelId"`
|
||||
// ModelName is the human-readable name
|
||||
ModelName string `json:"modelName"`
|
||||
// Description is the model description
|
||||
Description string `json:"description"`
|
||||
// RateMultiplier is the credit multiplier for this model
|
||||
RateMultiplier float64 `json:"rateMultiplier"`
|
||||
// RateUnit is the unit for rate calculation (e.g., "credit")
|
||||
RateUnit string `json:"rateUnit"`
|
||||
// MaxInputTokens is the maximum input token limit
|
||||
MaxInputTokens int `json:"maxInputTokens,omitempty"`
|
||||
}
|
||||
|
||||
// KiroIDETokenFile is the default path to Kiro IDE's token file
|
||||
const KiroIDETokenFile = ".aws/sso/cache/kiro-auth-token.json"
|
||||
|
||||
// LoadKiroIDEToken loads token data from Kiro IDE's token file.
|
||||
func LoadKiroIDEToken() (*KiroTokenData, error) {
|
||||
homeDir, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get home directory: %w", err)
|
||||
}
|
||||
|
||||
tokenPath := filepath.Join(homeDir, KiroIDETokenFile)
|
||||
data, err := os.ReadFile(tokenPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read Kiro IDE token file (%s): %w", tokenPath, err)
|
||||
}
|
||||
|
||||
var token KiroTokenData
|
||||
if err := json.Unmarshal(data, &token); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse Kiro IDE token: %w", err)
|
||||
}
|
||||
|
||||
if token.AccessToken == "" {
|
||||
return nil, fmt.Errorf("access token is empty in Kiro IDE token file")
|
||||
}
|
||||
|
||||
return &token, nil
|
||||
}
|
||||
|
||||
// LoadKiroTokenFromPath loads token data from a custom path.
|
||||
// This supports multiple accounts by allowing different token files.
|
||||
func LoadKiroTokenFromPath(tokenPath string) (*KiroTokenData, error) {
|
||||
// Expand ~ to home directory
|
||||
if len(tokenPath) > 0 && tokenPath[0] == '~' {
|
||||
homeDir, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get home directory: %w", err)
|
||||
}
|
||||
tokenPath = filepath.Join(homeDir, tokenPath[1:])
|
||||
}
|
||||
|
||||
data, err := os.ReadFile(tokenPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read token file (%s): %w", tokenPath, err)
|
||||
}
|
||||
|
||||
var token KiroTokenData
|
||||
if err := json.Unmarshal(data, &token); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse token file: %w", err)
|
||||
}
|
||||
|
||||
if token.AccessToken == "" {
|
||||
return nil, fmt.Errorf("access token is empty in token file")
|
||||
}
|
||||
|
||||
return &token, nil
|
||||
}
|
||||
|
||||
// ListKiroTokenFiles lists all Kiro token files in the cache directory.
|
||||
// This supports multiple accounts by finding all token files.
|
||||
func ListKiroTokenFiles() ([]string, error) {
|
||||
homeDir, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get home directory: %w", err)
|
||||
}
|
||||
|
||||
cacheDir := filepath.Join(homeDir, ".aws", "sso", "cache")
|
||||
|
||||
// Check if directory exists
|
||||
if _, err := os.Stat(cacheDir); os.IsNotExist(err) {
|
||||
return nil, nil // No token files
|
||||
}
|
||||
|
||||
entries, err := os.ReadDir(cacheDir)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read cache directory: %w", err)
|
||||
}
|
||||
|
||||
var tokenFiles []string
|
||||
for _, entry := range entries {
|
||||
if entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
name := entry.Name()
|
||||
// Look for kiro token files only (avoid matching unrelated AWS SSO cache files)
|
||||
if strings.HasSuffix(name, ".json") && strings.HasPrefix(name, "kiro") {
|
||||
tokenFiles = append(tokenFiles, filepath.Join(cacheDir, name))
|
||||
}
|
||||
}
|
||||
|
||||
return tokenFiles, nil
|
||||
}
|
||||
|
||||
// LoadAllKiroTokens loads all Kiro tokens from the cache directory.
|
||||
// This supports multiple accounts.
|
||||
func LoadAllKiroTokens() ([]*KiroTokenData, error) {
|
||||
files, err := ListKiroTokenFiles()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var tokens []*KiroTokenData
|
||||
for _, file := range files {
|
||||
token, err := LoadKiroTokenFromPath(file)
|
||||
if err != nil {
|
||||
// Skip invalid token files
|
||||
continue
|
||||
}
|
||||
tokens = append(tokens, token)
|
||||
}
|
||||
|
||||
return tokens, nil
|
||||
}
|
||||
|
||||
// JWTClaims represents the claims we care about from a JWT token.
|
||||
// JWT tokens from Kiro/AWS contain user information in the payload.
|
||||
type JWTClaims struct {
|
||||
Email string `json:"email,omitempty"`
|
||||
Sub string `json:"sub,omitempty"`
|
||||
PreferredUser string `json:"preferred_username,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
Iss string `json:"iss,omitempty"`
|
||||
}
|
||||
|
||||
// ExtractEmailFromJWT extracts the user's email from a JWT access token.
|
||||
// JWT tokens typically have format: header.payload.signature
|
||||
// The payload is base64url-encoded JSON containing user claims.
|
||||
func ExtractEmailFromJWT(accessToken string) string {
|
||||
if accessToken == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
// JWT format: header.payload.signature
|
||||
parts := strings.Split(accessToken, ".")
|
||||
if len(parts) != 3 {
|
||||
return ""
|
||||
}
|
||||
|
||||
// Decode the payload (second part)
|
||||
payload := parts[1]
|
||||
|
||||
// Add padding if needed (base64url requires padding)
|
||||
switch len(payload) % 4 {
|
||||
case 2:
|
||||
payload += "=="
|
||||
case 3:
|
||||
payload += "="
|
||||
}
|
||||
|
||||
decoded, err := base64.URLEncoding.DecodeString(payload)
|
||||
if err != nil {
|
||||
// Try RawURLEncoding (no padding)
|
||||
decoded, err = base64.RawURLEncoding.DecodeString(parts[1])
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
var claims JWTClaims
|
||||
if err := json.Unmarshal(decoded, &claims); err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
// Return email if available
|
||||
if claims.Email != "" {
|
||||
return claims.Email
|
||||
}
|
||||
|
||||
// Fallback to preferred_username (some providers use this)
|
||||
if claims.PreferredUser != "" && strings.Contains(claims.PreferredUser, "@") {
|
||||
return claims.PreferredUser
|
||||
}
|
||||
|
||||
// Fallback to sub if it looks like an email
|
||||
if claims.Sub != "" && strings.Contains(claims.Sub, "@") {
|
||||
return claims.Sub
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
// SanitizeEmailForFilename sanitizes an email address for use in a filename.
|
||||
// Replaces special characters with underscores and prevents path traversal attacks.
|
||||
// Also handles URL-encoded characters to prevent encoded path traversal attempts.
|
||||
func SanitizeEmailForFilename(email string) string {
|
||||
if email == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
result := email
|
||||
|
||||
// First, handle URL-encoded path traversal attempts (%2F, %2E, %5C, etc.)
|
||||
// This prevents encoded characters from bypassing the sanitization.
|
||||
// Note: We replace % last to catch any remaining encodings including double-encoding (%252F)
|
||||
result = strings.ReplaceAll(result, "%2F", "_") // /
|
||||
result = strings.ReplaceAll(result, "%2f", "_")
|
||||
result = strings.ReplaceAll(result, "%5C", "_") // \
|
||||
result = strings.ReplaceAll(result, "%5c", "_")
|
||||
result = strings.ReplaceAll(result, "%2E", "_") // .
|
||||
result = strings.ReplaceAll(result, "%2e", "_")
|
||||
result = strings.ReplaceAll(result, "%00", "_") // null byte
|
||||
result = strings.ReplaceAll(result, "%", "_") // Catch remaining % to prevent double-encoding attacks
|
||||
|
||||
// Replace characters that are problematic in filenames
|
||||
// Keep @ and . in middle but replace other special characters
|
||||
for _, char := range []string{"/", "\\", ":", "*", "?", "\"", "<", ">", "|", " ", "\x00"} {
|
||||
result = strings.ReplaceAll(result, char, "_")
|
||||
}
|
||||
|
||||
// Prevent path traversal: replace leading dots in each path component
|
||||
// This handles cases like "../../../etc/passwd" → "_.._.._.._etc_passwd"
|
||||
parts := strings.Split(result, "_")
|
||||
for i, part := range parts {
|
||||
for strings.HasPrefix(part, ".") {
|
||||
part = "_" + part[1:]
|
||||
}
|
||||
parts[i] = part
|
||||
}
|
||||
result = strings.Join(parts, "_")
|
||||
|
||||
return result
|
||||
}
|
||||
314
internal/auth/kiro/aws_auth.go
Normal file
314
internal/auth/kiro/aws_auth.go
Normal file
@@ -0,0 +1,314 @@
|
||||
// Package kiro provides OAuth2 authentication functionality for AWS CodeWhisperer (Kiro) API.
|
||||
// This package implements token loading, refresh, and API communication with CodeWhisperer.
|
||||
package kiro
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
const (
|
||||
// awsKiroEndpoint is used for CodeWhisperer management APIs (GetUsageLimits, ListProfiles, etc.)
|
||||
// Note: This is different from the Amazon Q streaming endpoint (q.us-east-1.amazonaws.com)
|
||||
// used in kiro_executor.go for GenerateAssistantResponse. Both endpoints are correct
|
||||
// for their respective API operations.
|
||||
awsKiroEndpoint = "https://codewhisperer.us-east-1.amazonaws.com"
|
||||
defaultTokenFile = "~/.aws/sso/cache/kiro-auth-token.json"
|
||||
targetGetUsage = "AmazonCodeWhispererService.GetUsageLimits"
|
||||
targetListModels = "AmazonCodeWhispererService.ListAvailableModels"
|
||||
targetGenerateChat = "AmazonCodeWhispererStreamingService.GenerateAssistantResponse"
|
||||
)
|
||||
|
||||
// KiroAuth handles AWS CodeWhisperer authentication and API communication.
|
||||
// It provides methods for loading tokens, refreshing expired tokens,
|
||||
// and communicating with the CodeWhisperer API.
|
||||
type KiroAuth struct {
|
||||
httpClient *http.Client
|
||||
endpoint string
|
||||
}
|
||||
|
||||
// NewKiroAuth creates a new Kiro authentication service.
|
||||
// It initializes the HTTP client with proxy settings from the configuration.
|
||||
//
|
||||
// Parameters:
|
||||
// - cfg: The application configuration containing proxy settings
|
||||
//
|
||||
// Returns:
|
||||
// - *KiroAuth: A new Kiro authentication service instance
|
||||
func NewKiroAuth(cfg *config.Config) *KiroAuth {
|
||||
return &KiroAuth{
|
||||
httpClient: util.SetProxy(&cfg.SDKConfig, &http.Client{Timeout: 120 * time.Second}),
|
||||
endpoint: awsKiroEndpoint,
|
||||
}
|
||||
}
|
||||
|
||||
// LoadTokenFromFile loads token data from a file path.
|
||||
// This method reads and parses the token file, expanding ~ to the home directory.
|
||||
//
|
||||
// Parameters:
|
||||
// - tokenFile: Path to the token file (supports ~ expansion)
|
||||
//
|
||||
// Returns:
|
||||
// - *KiroTokenData: The parsed token data
|
||||
// - error: An error if file reading or parsing fails
|
||||
func (k *KiroAuth) LoadTokenFromFile(tokenFile string) (*KiroTokenData, error) {
|
||||
// Expand ~ to home directory
|
||||
if strings.HasPrefix(tokenFile, "~") {
|
||||
home, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get home directory: %w", err)
|
||||
}
|
||||
tokenFile = filepath.Join(home, tokenFile[1:])
|
||||
}
|
||||
|
||||
data, err := os.ReadFile(tokenFile)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read token file: %w", err)
|
||||
}
|
||||
|
||||
var tokenData KiroTokenData
|
||||
if err := json.Unmarshal(data, &tokenData); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse token file: %w", err)
|
||||
}
|
||||
|
||||
return &tokenData, nil
|
||||
}
|
||||
|
||||
// IsTokenExpired checks if the token has expired.
|
||||
// This method parses the expiration timestamp and compares it with the current time.
|
||||
//
|
||||
// Parameters:
|
||||
// - tokenData: The token data to check
|
||||
//
|
||||
// Returns:
|
||||
// - bool: True if the token has expired, false otherwise
|
||||
func (k *KiroAuth) IsTokenExpired(tokenData *KiroTokenData) bool {
|
||||
if tokenData.ExpiresAt == "" {
|
||||
return true
|
||||
}
|
||||
|
||||
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
|
||||
if err != nil {
|
||||
// Try alternate format
|
||||
expiresAt, err = time.Parse("2006-01-02T15:04:05.000Z", tokenData.ExpiresAt)
|
||||
if err != nil {
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
return time.Now().After(expiresAt)
|
||||
}
|
||||
|
||||
// makeRequest sends a request to the CodeWhisperer API.
|
||||
// This is an internal method for making authenticated API calls.
|
||||
//
|
||||
// Parameters:
|
||||
// - ctx: The context for the request
|
||||
// - target: The API target (e.g., "AmazonCodeWhispererService.GetUsageLimits")
|
||||
// - accessToken: The OAuth access token
|
||||
// - payload: The request payload
|
||||
//
|
||||
// Returns:
|
||||
// - []byte: The response body
|
||||
// - error: An error if the request fails
|
||||
func (k *KiroAuth) makeRequest(ctx context.Context, target string, accessToken string, payload interface{}) ([]byte, error) {
|
||||
jsonBody, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, k.endpoint, strings.NewReader(string(jsonBody)))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/x-amz-json-1.0")
|
||||
req.Header.Set("x-amz-target", target)
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := k.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("request failed: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
if errClose := resp.Body.Close(); errClose != nil {
|
||||
log.Errorf("failed to close response body: %v", errClose)
|
||||
}
|
||||
}()
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read response: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("API error (status %d): %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
return body, nil
|
||||
}
|
||||
|
||||
// GetUsageLimits retrieves usage information from the CodeWhisperer API.
|
||||
// This method fetches the current usage statistics and subscription information.
|
||||
//
|
||||
// Parameters:
|
||||
// - ctx: The context for the request
|
||||
// - tokenData: The token data containing access token and profile ARN
|
||||
//
|
||||
// Returns:
|
||||
// - *KiroUsageInfo: The usage information
|
||||
// - error: An error if the request fails
|
||||
func (k *KiroAuth) GetUsageLimits(ctx context.Context, tokenData *KiroTokenData) (*KiroUsageInfo, error) {
|
||||
payload := map[string]interface{}{
|
||||
"origin": "AI_EDITOR",
|
||||
"profileArn": tokenData.ProfileArn,
|
||||
"resourceType": "AGENTIC_REQUEST",
|
||||
}
|
||||
|
||||
body, err := k.makeRequest(ctx, targetGetUsage, tokenData.AccessToken, payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var result struct {
|
||||
SubscriptionInfo struct {
|
||||
SubscriptionTitle string `json:"subscriptionTitle"`
|
||||
} `json:"subscriptionInfo"`
|
||||
UsageBreakdownList []struct {
|
||||
CurrentUsageWithPrecision float64 `json:"currentUsageWithPrecision"`
|
||||
UsageLimitWithPrecision float64 `json:"usageLimitWithPrecision"`
|
||||
} `json:"usageBreakdownList"`
|
||||
NextDateReset float64 `json:"nextDateReset"`
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(body, &result); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse usage response: %w", err)
|
||||
}
|
||||
|
||||
usage := &KiroUsageInfo{
|
||||
SubscriptionTitle: result.SubscriptionInfo.SubscriptionTitle,
|
||||
NextReset: fmt.Sprintf("%v", result.NextDateReset),
|
||||
}
|
||||
|
||||
if len(result.UsageBreakdownList) > 0 {
|
||||
usage.CurrentUsage = result.UsageBreakdownList[0].CurrentUsageWithPrecision
|
||||
usage.UsageLimit = result.UsageBreakdownList[0].UsageLimitWithPrecision
|
||||
}
|
||||
|
||||
return usage, nil
|
||||
}
|
||||
|
||||
// ListAvailableModels retrieves available models from the CodeWhisperer API.
|
||||
// This method fetches the list of AI models available for the authenticated user.
|
||||
//
|
||||
// Parameters:
|
||||
// - ctx: The context for the request
|
||||
// - tokenData: The token data containing access token and profile ARN
|
||||
//
|
||||
// Returns:
|
||||
// - []*KiroModel: The list of available models
|
||||
// - error: An error if the request fails
|
||||
func (k *KiroAuth) ListAvailableModels(ctx context.Context, tokenData *KiroTokenData) ([]*KiroModel, error) {
|
||||
payload := map[string]interface{}{
|
||||
"origin": "AI_EDITOR",
|
||||
"profileArn": tokenData.ProfileArn,
|
||||
}
|
||||
|
||||
body, err := k.makeRequest(ctx, targetListModels, tokenData.AccessToken, payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var result struct {
|
||||
Models []struct {
|
||||
ModelID string `json:"modelId"`
|
||||
ModelName string `json:"modelName"`
|
||||
Description string `json:"description"`
|
||||
RateMultiplier float64 `json:"rateMultiplier"`
|
||||
RateUnit string `json:"rateUnit"`
|
||||
TokenLimits struct {
|
||||
MaxInputTokens int `json:"maxInputTokens"`
|
||||
} `json:"tokenLimits"`
|
||||
} `json:"models"`
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(body, &result); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse models response: %w", err)
|
||||
}
|
||||
|
||||
models := make([]*KiroModel, 0, len(result.Models))
|
||||
for _, m := range result.Models {
|
||||
models = append(models, &KiroModel{
|
||||
ModelID: m.ModelID,
|
||||
ModelName: m.ModelName,
|
||||
Description: m.Description,
|
||||
RateMultiplier: m.RateMultiplier,
|
||||
RateUnit: m.RateUnit,
|
||||
MaxInputTokens: m.TokenLimits.MaxInputTokens,
|
||||
})
|
||||
}
|
||||
|
||||
return models, nil
|
||||
}
|
||||
|
||||
// CreateTokenStorage creates a new KiroTokenStorage from token data.
|
||||
// This method converts the token data into a storage structure suitable for persistence.
|
||||
//
|
||||
// Parameters:
|
||||
// - tokenData: The token data to convert
|
||||
//
|
||||
// Returns:
|
||||
// - *KiroTokenStorage: A new token storage instance
|
||||
func (k *KiroAuth) CreateTokenStorage(tokenData *KiroTokenData) *KiroTokenStorage {
|
||||
return &KiroTokenStorage{
|
||||
AccessToken: tokenData.AccessToken,
|
||||
RefreshToken: tokenData.RefreshToken,
|
||||
ProfileArn: tokenData.ProfileArn,
|
||||
ExpiresAt: tokenData.ExpiresAt,
|
||||
AuthMethod: tokenData.AuthMethod,
|
||||
Provider: tokenData.Provider,
|
||||
LastRefresh: time.Now().Format(time.RFC3339),
|
||||
}
|
||||
}
|
||||
|
||||
// ValidateToken checks if the token is valid by making a test API call.
|
||||
// This method verifies the token by attempting to fetch usage limits.
|
||||
//
|
||||
// Parameters:
|
||||
// - ctx: The context for the request
|
||||
// - tokenData: The token data to validate
|
||||
//
|
||||
// Returns:
|
||||
// - error: An error if the token is invalid
|
||||
func (k *KiroAuth) ValidateToken(ctx context.Context, tokenData *KiroTokenData) error {
|
||||
_, err := k.GetUsageLimits(ctx, tokenData)
|
||||
return err
|
||||
}
|
||||
|
||||
// UpdateTokenStorage updates an existing token storage with new token data.
|
||||
// This method refreshes the token storage with newly obtained access and refresh tokens.
|
||||
//
|
||||
// Parameters:
|
||||
// - storage: The existing token storage to update
|
||||
// - tokenData: The new token data to apply
|
||||
func (k *KiroAuth) UpdateTokenStorage(storage *KiroTokenStorage, tokenData *KiroTokenData) {
|
||||
storage.AccessToken = tokenData.AccessToken
|
||||
storage.RefreshToken = tokenData.RefreshToken
|
||||
storage.ProfileArn = tokenData.ProfileArn
|
||||
storage.ExpiresAt = tokenData.ExpiresAt
|
||||
storage.AuthMethod = tokenData.AuthMethod
|
||||
storage.Provider = tokenData.Provider
|
||||
storage.LastRefresh = time.Now().Format(time.RFC3339)
|
||||
}
|
||||
161
internal/auth/kiro/aws_test.go
Normal file
161
internal/auth/kiro/aws_test.go
Normal file
@@ -0,0 +1,161 @@
|
||||
package kiro
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestExtractEmailFromJWT(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
token string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "Empty token",
|
||||
token: "",
|
||||
expected: "",
|
||||
},
|
||||
{
|
||||
name: "Invalid token format",
|
||||
token: "not.a.valid.jwt",
|
||||
expected: "",
|
||||
},
|
||||
{
|
||||
name: "Invalid token - not base64",
|
||||
token: "xxx.yyy.zzz",
|
||||
expected: "",
|
||||
},
|
||||
{
|
||||
name: "Valid JWT with email",
|
||||
token: createTestJWT(map[string]any{"email": "test@example.com", "sub": "user123"}),
|
||||
expected: "test@example.com",
|
||||
},
|
||||
{
|
||||
name: "JWT without email but with preferred_username",
|
||||
token: createTestJWT(map[string]any{"preferred_username": "user@domain.com", "sub": "user123"}),
|
||||
expected: "user@domain.com",
|
||||
},
|
||||
{
|
||||
name: "JWT with email-like sub",
|
||||
token: createTestJWT(map[string]any{"sub": "another@test.com"}),
|
||||
expected: "another@test.com",
|
||||
},
|
||||
{
|
||||
name: "JWT without any email fields",
|
||||
token: createTestJWT(map[string]any{"sub": "user123", "name": "Test User"}),
|
||||
expected: "",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := ExtractEmailFromJWT(tt.token)
|
||||
if result != tt.expected {
|
||||
t.Errorf("ExtractEmailFromJWT() = %q, want %q", result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSanitizeEmailForFilename(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
email string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "Empty email",
|
||||
email: "",
|
||||
expected: "",
|
||||
},
|
||||
{
|
||||
name: "Simple email",
|
||||
email: "user@example.com",
|
||||
expected: "user@example.com",
|
||||
},
|
||||
{
|
||||
name: "Email with space",
|
||||
email: "user name@example.com",
|
||||
expected: "user_name@example.com",
|
||||
},
|
||||
{
|
||||
name: "Email with special chars",
|
||||
email: "user:name@example.com",
|
||||
expected: "user_name@example.com",
|
||||
},
|
||||
{
|
||||
name: "Email with multiple special chars",
|
||||
email: "user/name:test@example.com",
|
||||
expected: "user_name_test@example.com",
|
||||
},
|
||||
{
|
||||
name: "Path traversal attempt",
|
||||
email: "../../../etc/passwd",
|
||||
expected: "_.__.__._etc_passwd",
|
||||
},
|
||||
{
|
||||
name: "Path traversal with backslash",
|
||||
email: `..\..\..\..\windows\system32`,
|
||||
expected: "_.__.__.__._windows_system32",
|
||||
},
|
||||
{
|
||||
name: "Null byte injection attempt",
|
||||
email: "user\x00@evil.com",
|
||||
expected: "user_@evil.com",
|
||||
},
|
||||
// URL-encoded path traversal tests
|
||||
{
|
||||
name: "URL-encoded slash",
|
||||
email: "user%2Fpath@example.com",
|
||||
expected: "user_path@example.com",
|
||||
},
|
||||
{
|
||||
name: "URL-encoded backslash",
|
||||
email: "user%5Cpath@example.com",
|
||||
expected: "user_path@example.com",
|
||||
},
|
||||
{
|
||||
name: "URL-encoded dot",
|
||||
email: "%2E%2E%2Fetc%2Fpasswd",
|
||||
expected: "___etc_passwd",
|
||||
},
|
||||
{
|
||||
name: "URL-encoded null",
|
||||
email: "user%00@evil.com",
|
||||
expected: "user_@evil.com",
|
||||
},
|
||||
{
|
||||
name: "Double URL-encoding attack",
|
||||
email: "%252F%252E%252E",
|
||||
expected: "_252F_252E_252E", // % replaced with _, remaining chars preserved (safe)
|
||||
},
|
||||
{
|
||||
name: "Mixed case URL-encoding",
|
||||
email: "%2f%2F%5c%5C",
|
||||
expected: "____",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := SanitizeEmailForFilename(tt.email)
|
||||
if result != tt.expected {
|
||||
t.Errorf("SanitizeEmailForFilename() = %q, want %q", result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// createTestJWT creates a test JWT token with the given claims
|
||||
func createTestJWT(claims map[string]any) string {
|
||||
header := base64.RawURLEncoding.EncodeToString([]byte(`{"alg":"RS256","typ":"JWT"}`))
|
||||
|
||||
payloadBytes, _ := json.Marshal(claims)
|
||||
payload := base64.RawURLEncoding.EncodeToString(payloadBytes)
|
||||
|
||||
signature := base64.RawURLEncoding.EncodeToString([]byte("fake-signature"))
|
||||
|
||||
return header + "." + payload + "." + signature
|
||||
}
|
||||
296
internal/auth/kiro/oauth.go
Normal file
296
internal/auth/kiro/oauth.go
Normal file
@@ -0,0 +1,296 @@
|
||||
// Package kiro provides OAuth2 authentication for Kiro using native Google login.
|
||||
package kiro
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"html"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
const (
|
||||
// Kiro auth endpoint
|
||||
kiroAuthEndpoint = "https://prod.us-east-1.auth.desktop.kiro.dev"
|
||||
|
||||
// Default callback port
|
||||
defaultCallbackPort = 9876
|
||||
|
||||
// Auth timeout
|
||||
authTimeout = 10 * time.Minute
|
||||
)
|
||||
|
||||
// KiroTokenResponse represents the response from Kiro token endpoint.
|
||||
type KiroTokenResponse struct {
|
||||
AccessToken string `json:"accessToken"`
|
||||
RefreshToken string `json:"refreshToken"`
|
||||
ProfileArn string `json:"profileArn"`
|
||||
ExpiresIn int `json:"expiresIn"`
|
||||
}
|
||||
|
||||
// KiroOAuth handles the OAuth flow for Kiro authentication.
|
||||
type KiroOAuth struct {
|
||||
httpClient *http.Client
|
||||
cfg *config.Config
|
||||
}
|
||||
|
||||
// NewKiroOAuth creates a new Kiro OAuth handler.
|
||||
func NewKiroOAuth(cfg *config.Config) *KiroOAuth {
|
||||
client := &http.Client{Timeout: 30 * time.Second}
|
||||
if cfg != nil {
|
||||
client = util.SetProxy(&cfg.SDKConfig, client)
|
||||
}
|
||||
return &KiroOAuth{
|
||||
httpClient: client,
|
||||
cfg: cfg,
|
||||
}
|
||||
}
|
||||
|
||||
// generateCodeVerifier generates a random code verifier for PKCE.
|
||||
func generateCodeVerifier() (string, error) {
|
||||
b := make([]byte, 32)
|
||||
if _, err := rand.Read(b); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return base64.RawURLEncoding.EncodeToString(b), nil
|
||||
}
|
||||
|
||||
// generateCodeChallenge generates the code challenge from verifier.
|
||||
func generateCodeChallenge(verifier string) string {
|
||||
h := sha256.Sum256([]byte(verifier))
|
||||
return base64.RawURLEncoding.EncodeToString(h[:])
|
||||
}
|
||||
|
||||
// generateState generates a random state parameter.
|
||||
func generateState() (string, error) {
|
||||
b := make([]byte, 16)
|
||||
if _, err := rand.Read(b); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return base64.RawURLEncoding.EncodeToString(b), nil
|
||||
}
|
||||
|
||||
// AuthResult contains the authorization code and state from callback.
|
||||
type AuthResult struct {
|
||||
Code string
|
||||
State string
|
||||
Error string
|
||||
}
|
||||
|
||||
// startCallbackServer starts a local HTTP server to receive the OAuth callback.
|
||||
func (o *KiroOAuth) startCallbackServer(ctx context.Context, expectedState string) (string, <-chan AuthResult, error) {
|
||||
// Try to find an available port - use localhost like Kiro does
|
||||
listener, err := net.Listen("tcp", fmt.Sprintf("localhost:%d", defaultCallbackPort))
|
||||
if err != nil {
|
||||
// Try with dynamic port (RFC 8252 allows dynamic ports for native apps)
|
||||
log.Warnf("kiro oauth: default port %d is busy, falling back to dynamic port", defaultCallbackPort)
|
||||
listener, err = net.Listen("tcp", "localhost:0")
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("failed to start callback server: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
port := listener.Addr().(*net.TCPAddr).Port
|
||||
// Use http scheme for local callback server
|
||||
redirectURI := fmt.Sprintf("http://localhost:%d/oauth/callback", port)
|
||||
resultChan := make(chan AuthResult, 1)
|
||||
|
||||
server := &http.Server{
|
||||
ReadHeaderTimeout: 10 * time.Second,
|
||||
}
|
||||
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/oauth/callback", func(w http.ResponseWriter, r *http.Request) {
|
||||
code := r.URL.Query().Get("code")
|
||||
state := r.URL.Query().Get("state")
|
||||
errParam := r.URL.Query().Get("error")
|
||||
|
||||
if errParam != "" {
|
||||
w.Header().Set("Content-Type", "text/html")
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
fmt.Fprintf(w, `<html><body><h1>Login Failed</h1><p>%s</p><p>You can close this window.</p></body></html>`, html.EscapeString(errParam))
|
||||
resultChan <- AuthResult{Error: errParam}
|
||||
return
|
||||
}
|
||||
|
||||
if state != expectedState {
|
||||
w.Header().Set("Content-Type", "text/html")
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
fmt.Fprint(w, `<html><body><h1>Login Failed</h1><p>Invalid state parameter</p><p>You can close this window.</p></body></html>`)
|
||||
resultChan <- AuthResult{Error: "state mismatch"}
|
||||
return
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "text/html")
|
||||
fmt.Fprint(w, `<html><body><h1>Login Successful!</h1><p>You can close this window and return to the terminal.</p></body></html>`)
|
||||
resultChan <- AuthResult{Code: code, State: state}
|
||||
})
|
||||
|
||||
server.Handler = mux
|
||||
|
||||
go func() {
|
||||
if err := server.Serve(listener); err != nil && err != http.ErrServerClosed {
|
||||
log.Debugf("callback server error: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
go func() {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
case <-time.After(authTimeout):
|
||||
case <-resultChan:
|
||||
}
|
||||
_ = server.Shutdown(context.Background())
|
||||
}()
|
||||
|
||||
return redirectURI, resultChan, nil
|
||||
}
|
||||
|
||||
// LoginWithBuilderID performs OAuth login with AWS Builder ID using device code flow.
|
||||
func (o *KiroOAuth) LoginWithBuilderID(ctx context.Context) (*KiroTokenData, error) {
|
||||
ssoClient := NewSSOOIDCClient(o.cfg)
|
||||
return ssoClient.LoginWithBuilderID(ctx)
|
||||
}
|
||||
|
||||
// exchangeCodeForToken exchanges the authorization code for tokens.
|
||||
func (o *KiroOAuth) exchangeCodeForToken(ctx context.Context, code, codeVerifier, redirectURI string) (*KiroTokenData, error) {
|
||||
payload := map[string]string{
|
||||
"code": code,
|
||||
"code_verifier": codeVerifier,
|
||||
"redirect_uri": redirectURI,
|
||||
}
|
||||
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
||||
}
|
||||
|
||||
tokenURL := kiroAuthEndpoint + "/oauth/token"
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenURL, strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", "cli-proxy-api/1.0.0")
|
||||
|
||||
resp, err := o.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("token request failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read response: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
log.Debugf("token exchange failed (status %d): %s", resp.StatusCode, string(respBody))
|
||||
return nil, fmt.Errorf("token exchange failed (status %d)", resp.StatusCode)
|
||||
}
|
||||
|
||||
var tokenResp KiroTokenResponse
|
||||
if err := json.Unmarshal(respBody, &tokenResp); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse token response: %w", err)
|
||||
}
|
||||
|
||||
// Validate ExpiresIn - use default 1 hour if invalid
|
||||
expiresIn := tokenResp.ExpiresIn
|
||||
if expiresIn <= 0 {
|
||||
expiresIn = 3600
|
||||
}
|
||||
expiresAt := time.Now().Add(time.Duration(expiresIn) * time.Second)
|
||||
|
||||
return &KiroTokenData{
|
||||
AccessToken: tokenResp.AccessToken,
|
||||
RefreshToken: tokenResp.RefreshToken,
|
||||
ProfileArn: tokenResp.ProfileArn,
|
||||
ExpiresAt: expiresAt.Format(time.RFC3339),
|
||||
AuthMethod: "social",
|
||||
Provider: "", // Caller should preserve original provider
|
||||
}, nil
|
||||
}
|
||||
|
||||
// RefreshToken refreshes an expired access token.
|
||||
func (o *KiroOAuth) RefreshToken(ctx context.Context, refreshToken string) (*KiroTokenData, error) {
|
||||
payload := map[string]string{
|
||||
"refreshToken": refreshToken,
|
||||
}
|
||||
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
||||
}
|
||||
|
||||
refreshURL := kiroAuthEndpoint + "/refreshToken"
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, refreshURL, strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", "cli-proxy-api/1.0.0")
|
||||
|
||||
resp, err := o.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("refresh request failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read response: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
log.Debugf("token refresh failed (status %d): %s", resp.StatusCode, string(respBody))
|
||||
return nil, fmt.Errorf("token refresh failed (status %d)", resp.StatusCode)
|
||||
}
|
||||
|
||||
var tokenResp KiroTokenResponse
|
||||
if err := json.Unmarshal(respBody, &tokenResp); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse token response: %w", err)
|
||||
}
|
||||
|
||||
// Validate ExpiresIn - use default 1 hour if invalid
|
||||
expiresIn := tokenResp.ExpiresIn
|
||||
if expiresIn <= 0 {
|
||||
expiresIn = 3600
|
||||
}
|
||||
expiresAt := time.Now().Add(time.Duration(expiresIn) * time.Second)
|
||||
|
||||
return &KiroTokenData{
|
||||
AccessToken: tokenResp.AccessToken,
|
||||
RefreshToken: tokenResp.RefreshToken,
|
||||
ProfileArn: tokenResp.ProfileArn,
|
||||
ExpiresAt: expiresAt.Format(time.RFC3339),
|
||||
AuthMethod: "social",
|
||||
Provider: "", // Caller should preserve original provider
|
||||
}, nil
|
||||
}
|
||||
|
||||
// LoginWithGoogle performs OAuth login with Google using Kiro's social auth.
|
||||
// This uses a custom protocol handler (kiro://) to receive the callback.
|
||||
func (o *KiroOAuth) LoginWithGoogle(ctx context.Context) (*KiroTokenData, error) {
|
||||
socialClient := NewSocialAuthClient(o.cfg)
|
||||
return socialClient.LoginWithGoogle(ctx)
|
||||
}
|
||||
|
||||
// LoginWithGitHub performs OAuth login with GitHub using Kiro's social auth.
|
||||
// This uses a custom protocol handler (kiro://) to receive the callback.
|
||||
func (o *KiroOAuth) LoginWithGitHub(ctx context.Context) (*KiroTokenData, error) {
|
||||
socialClient := NewSocialAuthClient(o.cfg)
|
||||
return socialClient.LoginWithGitHub(ctx)
|
||||
}
|
||||
725
internal/auth/kiro/protocol_handler.go
Normal file
725
internal/auth/kiro/protocol_handler.go
Normal file
@@ -0,0 +1,725 @@
|
||||
// Package kiro provides custom protocol handler registration for Kiro OAuth.
|
||||
// This enables the CLI to intercept kiro:// URIs for social authentication (Google/GitHub).
|
||||
package kiro
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"html"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
const (
|
||||
// KiroProtocol is the custom URI scheme used by Kiro
|
||||
KiroProtocol = "kiro"
|
||||
|
||||
// KiroAuthority is the URI authority for authentication callbacks
|
||||
KiroAuthority = "kiro.kiroAgent"
|
||||
|
||||
// KiroAuthPath is the path for successful authentication
|
||||
KiroAuthPath = "/authenticate-success"
|
||||
|
||||
// KiroRedirectURI is the full redirect URI for social auth
|
||||
KiroRedirectURI = "kiro://kiro.kiroAgent/authenticate-success"
|
||||
|
||||
// DefaultHandlerPort is the default port for the local callback server
|
||||
DefaultHandlerPort = 19876
|
||||
|
||||
// HandlerTimeout is how long to wait for the OAuth callback
|
||||
HandlerTimeout = 10 * time.Minute
|
||||
)
|
||||
|
||||
// ProtocolHandler manages the custom kiro:// protocol handler for OAuth callbacks.
|
||||
type ProtocolHandler struct {
|
||||
port int
|
||||
server *http.Server
|
||||
listener net.Listener
|
||||
resultChan chan *AuthCallback
|
||||
stopChan chan struct{}
|
||||
mu sync.Mutex
|
||||
running bool
|
||||
}
|
||||
|
||||
// AuthCallback contains the OAuth callback parameters.
|
||||
type AuthCallback struct {
|
||||
Code string
|
||||
State string
|
||||
Error string
|
||||
}
|
||||
|
||||
// NewProtocolHandler creates a new protocol handler.
|
||||
func NewProtocolHandler() *ProtocolHandler {
|
||||
return &ProtocolHandler{
|
||||
port: DefaultHandlerPort,
|
||||
resultChan: make(chan *AuthCallback, 1),
|
||||
stopChan: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Start starts the local callback server that receives redirects from the protocol handler.
|
||||
func (h *ProtocolHandler) Start(ctx context.Context) (int, error) {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
|
||||
if h.running {
|
||||
return h.port, nil
|
||||
}
|
||||
|
||||
// Drain any stale results from previous runs
|
||||
select {
|
||||
case <-h.resultChan:
|
||||
default:
|
||||
}
|
||||
|
||||
// Reset stopChan for reuse - close old channel first to unblock any waiting goroutines
|
||||
if h.stopChan != nil {
|
||||
select {
|
||||
case <-h.stopChan:
|
||||
// Already closed
|
||||
default:
|
||||
close(h.stopChan)
|
||||
}
|
||||
}
|
||||
h.stopChan = make(chan struct{})
|
||||
|
||||
// Try ports in known range (must match handler script port range)
|
||||
var listener net.Listener
|
||||
var err error
|
||||
portRange := []int{DefaultHandlerPort, DefaultHandlerPort + 1, DefaultHandlerPort + 2, DefaultHandlerPort + 3, DefaultHandlerPort + 4}
|
||||
|
||||
for _, port := range portRange {
|
||||
listener, err = net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port))
|
||||
if err == nil {
|
||||
break
|
||||
}
|
||||
log.Debugf("kiro protocol handler: port %d busy, trying next", port)
|
||||
}
|
||||
|
||||
if listener == nil {
|
||||
return 0, fmt.Errorf("failed to start callback server: all ports %d-%d are busy", DefaultHandlerPort, DefaultHandlerPort+4)
|
||||
}
|
||||
|
||||
h.listener = listener
|
||||
h.port = listener.Addr().(*net.TCPAddr).Port
|
||||
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/oauth/callback", h.handleCallback)
|
||||
|
||||
h.server = &http.Server{
|
||||
Handler: mux,
|
||||
ReadHeaderTimeout: 10 * time.Second,
|
||||
}
|
||||
|
||||
go func() {
|
||||
if err := h.server.Serve(listener); err != nil && err != http.ErrServerClosed {
|
||||
log.Debugf("kiro protocol handler server error: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
h.running = true
|
||||
log.Debugf("kiro protocol handler started on port %d", h.port)
|
||||
|
||||
// Auto-shutdown after context done, timeout, or explicit stop
|
||||
// Capture references to prevent race with new Start() calls
|
||||
currentStopChan := h.stopChan
|
||||
currentServer := h.server
|
||||
currentListener := h.listener
|
||||
go func() {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
case <-time.After(HandlerTimeout):
|
||||
case <-currentStopChan:
|
||||
return // Already stopped, exit goroutine
|
||||
}
|
||||
// Only stop if this is still the current server/listener instance
|
||||
h.mu.Lock()
|
||||
if h.server == currentServer && h.listener == currentListener {
|
||||
h.mu.Unlock()
|
||||
h.Stop()
|
||||
} else {
|
||||
h.mu.Unlock()
|
||||
}
|
||||
}()
|
||||
|
||||
return h.port, nil
|
||||
}
|
||||
|
||||
// Stop stops the callback server.
|
||||
func (h *ProtocolHandler) Stop() {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
|
||||
if !h.running {
|
||||
return
|
||||
}
|
||||
|
||||
// Signal the auto-shutdown goroutine to exit.
|
||||
// This select pattern is safe because stopChan is only modified while holding h.mu,
|
||||
// and we hold the lock here. The select prevents panic from double-close.
|
||||
select {
|
||||
case <-h.stopChan:
|
||||
// Already closed
|
||||
default:
|
||||
close(h.stopChan)
|
||||
}
|
||||
|
||||
if h.server != nil {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
_ = h.server.Shutdown(ctx)
|
||||
}
|
||||
|
||||
h.running = false
|
||||
log.Debug("kiro protocol handler stopped")
|
||||
}
|
||||
|
||||
// WaitForCallback waits for the OAuth callback and returns the result.
|
||||
func (h *ProtocolHandler) WaitForCallback(ctx context.Context) (*AuthCallback, error) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
case <-time.After(HandlerTimeout):
|
||||
return nil, fmt.Errorf("timeout waiting for OAuth callback")
|
||||
case result := <-h.resultChan:
|
||||
return result, nil
|
||||
}
|
||||
}
|
||||
|
||||
// GetPort returns the port the handler is listening on.
|
||||
func (h *ProtocolHandler) GetPort() int {
|
||||
return h.port
|
||||
}
|
||||
|
||||
// handleCallback processes the OAuth callback from the protocol handler script.
|
||||
func (h *ProtocolHandler) handleCallback(w http.ResponseWriter, r *http.Request) {
|
||||
code := r.URL.Query().Get("code")
|
||||
state := r.URL.Query().Get("state")
|
||||
errParam := r.URL.Query().Get("error")
|
||||
|
||||
result := &AuthCallback{
|
||||
Code: code,
|
||||
State: state,
|
||||
Error: errParam,
|
||||
}
|
||||
|
||||
// Send result
|
||||
select {
|
||||
case h.resultChan <- result:
|
||||
default:
|
||||
// Channel full, ignore duplicate callbacks
|
||||
}
|
||||
|
||||
// Send success response
|
||||
w.Header().Set("Content-Type", "text/html; charset=utf-8")
|
||||
if errParam != "" {
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
fmt.Fprintf(w, `<!DOCTYPE html>
|
||||
<html>
|
||||
<head><title>Login Failed</title></head>
|
||||
<body>
|
||||
<h1>Login Failed</h1>
|
||||
<p>Error: %s</p>
|
||||
<p>You can close this window.</p>
|
||||
</body>
|
||||
</html>`, html.EscapeString(errParam))
|
||||
} else {
|
||||
fmt.Fprint(w, `<!DOCTYPE html>
|
||||
<html>
|
||||
<head><title>Login Successful</title></head>
|
||||
<body>
|
||||
<h1>Login Successful!</h1>
|
||||
<p>You can close this window and return to the terminal.</p>
|
||||
<script>window.close();</script>
|
||||
</body>
|
||||
</html>`)
|
||||
}
|
||||
}
|
||||
|
||||
// IsProtocolHandlerInstalled checks if the kiro:// protocol handler is installed.
|
||||
func IsProtocolHandlerInstalled() bool {
|
||||
switch runtime.GOOS {
|
||||
case "linux":
|
||||
return isLinuxHandlerInstalled()
|
||||
case "windows":
|
||||
return isWindowsHandlerInstalled()
|
||||
case "darwin":
|
||||
return isDarwinHandlerInstalled()
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// InstallProtocolHandler installs the kiro:// protocol handler for the current platform.
|
||||
func InstallProtocolHandler(handlerPort int) error {
|
||||
switch runtime.GOOS {
|
||||
case "linux":
|
||||
return installLinuxHandler(handlerPort)
|
||||
case "windows":
|
||||
return installWindowsHandler(handlerPort)
|
||||
case "darwin":
|
||||
return installDarwinHandler(handlerPort)
|
||||
default:
|
||||
return fmt.Errorf("unsupported platform: %s", runtime.GOOS)
|
||||
}
|
||||
}
|
||||
|
||||
// UninstallProtocolHandler removes the kiro:// protocol handler.
|
||||
func UninstallProtocolHandler() error {
|
||||
switch runtime.GOOS {
|
||||
case "linux":
|
||||
return uninstallLinuxHandler()
|
||||
case "windows":
|
||||
return uninstallWindowsHandler()
|
||||
case "darwin":
|
||||
return uninstallDarwinHandler()
|
||||
default:
|
||||
return fmt.Errorf("unsupported platform: %s", runtime.GOOS)
|
||||
}
|
||||
}
|
||||
|
||||
// --- Linux Implementation ---
|
||||
|
||||
func getLinuxDesktopPath() string {
|
||||
homeDir, _ := os.UserHomeDir()
|
||||
return filepath.Join(homeDir, ".local", "share", "applications", "kiro-oauth-handler.desktop")
|
||||
}
|
||||
|
||||
func getLinuxHandlerScriptPath() string {
|
||||
homeDir, _ := os.UserHomeDir()
|
||||
return filepath.Join(homeDir, ".local", "bin", "kiro-oauth-handler")
|
||||
}
|
||||
|
||||
func isLinuxHandlerInstalled() bool {
|
||||
desktopPath := getLinuxDesktopPath()
|
||||
_, err := os.Stat(desktopPath)
|
||||
return err == nil
|
||||
}
|
||||
|
||||
func installLinuxHandler(handlerPort int) error {
|
||||
// Create directories
|
||||
homeDir, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
binDir := filepath.Join(homeDir, ".local", "bin")
|
||||
appDir := filepath.Join(homeDir, ".local", "share", "applications")
|
||||
|
||||
if err := os.MkdirAll(binDir, 0755); err != nil {
|
||||
return fmt.Errorf("failed to create bin directory: %w", err)
|
||||
}
|
||||
if err := os.MkdirAll(appDir, 0755); err != nil {
|
||||
return fmt.Errorf("failed to create applications directory: %w", err)
|
||||
}
|
||||
|
||||
// Create handler script - tries multiple ports to handle dynamic port allocation
|
||||
scriptPath := getLinuxHandlerScriptPath()
|
||||
scriptContent := fmt.Sprintf(`#!/bin/bash
|
||||
# Kiro OAuth Protocol Handler
|
||||
# Handles kiro:// URIs - tries CLI first, then forwards to Kiro IDE
|
||||
|
||||
URL="$1"
|
||||
|
||||
# Check curl availability
|
||||
if ! command -v curl &> /dev/null; then
|
||||
echo "Error: curl is required for Kiro OAuth handler" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Extract code and state from URL
|
||||
[[ "$URL" =~ code=([^&]+) ]] && CODE="${BASH_REMATCH[1]}"
|
||||
[[ "$URL" =~ state=([^&]+) ]] && STATE="${BASH_REMATCH[1]}"
|
||||
[[ "$URL" =~ error=([^&]+) ]] && ERROR="${BASH_REMATCH[1]}"
|
||||
|
||||
# Try CLI proxy on multiple possible ports (default + dynamic range)
|
||||
CLI_OK=0
|
||||
for PORT in %d %d %d %d %d; do
|
||||
if [ -n "$ERROR" ]; then
|
||||
curl -sf --connect-timeout 1 "http://127.0.0.1:$PORT/oauth/callback?error=$ERROR" && CLI_OK=1 && break
|
||||
elif [ -n "$CODE" ] && [ -n "$STATE" ]; then
|
||||
curl -sf --connect-timeout 1 "http://127.0.0.1:$PORT/oauth/callback?code=$CODE&state=$STATE" && CLI_OK=1 && break
|
||||
fi
|
||||
done
|
||||
|
||||
# If CLI not available, forward to Kiro IDE
|
||||
if [ $CLI_OK -eq 0 ] && [ -x "/usr/share/kiro/kiro" ]; then
|
||||
/usr/share/kiro/kiro --open-url "$URL" &
|
||||
fi
|
||||
`, handlerPort, handlerPort+1, handlerPort+2, handlerPort+3, handlerPort+4)
|
||||
|
||||
if err := os.WriteFile(scriptPath, []byte(scriptContent), 0755); err != nil {
|
||||
return fmt.Errorf("failed to write handler script: %w", err)
|
||||
}
|
||||
|
||||
// Create .desktop file
|
||||
desktopPath := getLinuxDesktopPath()
|
||||
desktopContent := fmt.Sprintf(`[Desktop Entry]
|
||||
Name=Kiro OAuth Handler
|
||||
Comment=Handle kiro:// protocol for CLI Proxy API authentication
|
||||
Exec=%s %%u
|
||||
Type=Application
|
||||
Terminal=false
|
||||
NoDisplay=true
|
||||
MimeType=x-scheme-handler/kiro;
|
||||
Categories=Utility;
|
||||
`, scriptPath)
|
||||
|
||||
if err := os.WriteFile(desktopPath, []byte(desktopContent), 0644); err != nil {
|
||||
return fmt.Errorf("failed to write desktop file: %w", err)
|
||||
}
|
||||
|
||||
// Register handler with xdg-mime
|
||||
cmd := exec.Command("xdg-mime", "default", "kiro-oauth-handler.desktop", "x-scheme-handler/kiro")
|
||||
if err := cmd.Run(); err != nil {
|
||||
log.Warnf("xdg-mime registration failed (may need manual setup): %v", err)
|
||||
}
|
||||
|
||||
// Update desktop database
|
||||
cmd = exec.Command("update-desktop-database", appDir)
|
||||
_ = cmd.Run() // Ignore errors, not critical
|
||||
|
||||
log.Info("Kiro protocol handler installed for Linux")
|
||||
return nil
|
||||
}
|
||||
|
||||
func uninstallLinuxHandler() error {
|
||||
desktopPath := getLinuxDesktopPath()
|
||||
scriptPath := getLinuxHandlerScriptPath()
|
||||
|
||||
if err := os.Remove(desktopPath); err != nil && !os.IsNotExist(err) {
|
||||
return fmt.Errorf("failed to remove desktop file: %w", err)
|
||||
}
|
||||
if err := os.Remove(scriptPath); err != nil && !os.IsNotExist(err) {
|
||||
return fmt.Errorf("failed to remove handler script: %w", err)
|
||||
}
|
||||
|
||||
log.Info("Kiro protocol handler uninstalled")
|
||||
return nil
|
||||
}
|
||||
|
||||
// --- Windows Implementation ---
|
||||
|
||||
func isWindowsHandlerInstalled() bool {
|
||||
// Check registry key existence
|
||||
cmd := exec.Command("reg", "query", `HKCU\Software\Classes\kiro`, "/ve")
|
||||
return cmd.Run() == nil
|
||||
}
|
||||
|
||||
func installWindowsHandler(handlerPort int) error {
|
||||
homeDir, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Create handler script (PowerShell)
|
||||
scriptDir := filepath.Join(homeDir, ".cliproxyapi")
|
||||
if err := os.MkdirAll(scriptDir, 0755); err != nil {
|
||||
return fmt.Errorf("failed to create script directory: %w", err)
|
||||
}
|
||||
|
||||
scriptPath := filepath.Join(scriptDir, "kiro-oauth-handler.ps1")
|
||||
scriptContent := fmt.Sprintf(`# Kiro OAuth Protocol Handler for Windows
|
||||
param([string]$url)
|
||||
|
||||
# Load required assembly for HttpUtility
|
||||
Add-Type -AssemblyName System.Web
|
||||
|
||||
# Parse URL parameters
|
||||
$uri = [System.Uri]$url
|
||||
$query = [System.Web.HttpUtility]::ParseQueryString($uri.Query)
|
||||
$code = $query["code"]
|
||||
$state = $query["state"]
|
||||
$errorParam = $query["error"]
|
||||
|
||||
# Try multiple ports (default + dynamic range)
|
||||
$ports = @(%d, %d, %d, %d, %d)
|
||||
$success = $false
|
||||
|
||||
foreach ($port in $ports) {
|
||||
if ($success) { break }
|
||||
$callbackUrl = "http://127.0.0.1:$port/oauth/callback"
|
||||
try {
|
||||
if ($errorParam) {
|
||||
$fullUrl = $callbackUrl + "?error=" + $errorParam
|
||||
Invoke-WebRequest -Uri $fullUrl -UseBasicParsing -TimeoutSec 1 -ErrorAction Stop | Out-Null
|
||||
$success = $true
|
||||
} elseif ($code -and $state) {
|
||||
$fullUrl = $callbackUrl + "?code=" + $code + "&state=" + $state
|
||||
Invoke-WebRequest -Uri $fullUrl -UseBasicParsing -TimeoutSec 1 -ErrorAction Stop | Out-Null
|
||||
$success = $true
|
||||
}
|
||||
} catch {
|
||||
# Try next port
|
||||
}
|
||||
}
|
||||
`, handlerPort, handlerPort+1, handlerPort+2, handlerPort+3, handlerPort+4)
|
||||
|
||||
if err := os.WriteFile(scriptPath, []byte(scriptContent), 0644); err != nil {
|
||||
return fmt.Errorf("failed to write handler script: %w", err)
|
||||
}
|
||||
|
||||
// Create batch wrapper
|
||||
batchPath := filepath.Join(scriptDir, "kiro-oauth-handler.bat")
|
||||
batchContent := fmt.Sprintf("@echo off\npowershell -ExecutionPolicy Bypass -File \"%s\" \"%%1\"\n", scriptPath)
|
||||
|
||||
if err := os.WriteFile(batchPath, []byte(batchContent), 0644); err != nil {
|
||||
return fmt.Errorf("failed to write batch wrapper: %w", err)
|
||||
}
|
||||
|
||||
// Register in Windows registry
|
||||
commands := [][]string{
|
||||
{"reg", "add", `HKCU\Software\Classes\kiro`, "/ve", "/d", "URL:Kiro Protocol", "/f"},
|
||||
{"reg", "add", `HKCU\Software\Classes\kiro`, "/v", "URL Protocol", "/d", "", "/f"},
|
||||
{"reg", "add", `HKCU\Software\Classes\kiro\shell`, "/f"},
|
||||
{"reg", "add", `HKCU\Software\Classes\kiro\shell\open`, "/f"},
|
||||
{"reg", "add", `HKCU\Software\Classes\kiro\shell\open\command`, "/ve", "/d", fmt.Sprintf("\"%s\" \"%%1\"", batchPath), "/f"},
|
||||
}
|
||||
|
||||
for _, args := range commands {
|
||||
cmd := exec.Command(args[0], args[1:]...)
|
||||
if err := cmd.Run(); err != nil {
|
||||
return fmt.Errorf("failed to run registry command: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
log.Info("Kiro protocol handler installed for Windows")
|
||||
return nil
|
||||
}
|
||||
|
||||
func uninstallWindowsHandler() error {
|
||||
// Remove registry keys
|
||||
cmd := exec.Command("reg", "delete", `HKCU\Software\Classes\kiro`, "/f")
|
||||
if err := cmd.Run(); err != nil {
|
||||
log.Warnf("failed to remove registry key: %v", err)
|
||||
}
|
||||
|
||||
// Remove scripts
|
||||
homeDir, _ := os.UserHomeDir()
|
||||
scriptDir := filepath.Join(homeDir, ".cliproxyapi")
|
||||
_ = os.Remove(filepath.Join(scriptDir, "kiro-oauth-handler.ps1"))
|
||||
_ = os.Remove(filepath.Join(scriptDir, "kiro-oauth-handler.bat"))
|
||||
|
||||
log.Info("Kiro protocol handler uninstalled")
|
||||
return nil
|
||||
}
|
||||
|
||||
// --- macOS Implementation ---
|
||||
|
||||
func getDarwinAppPath() string {
|
||||
homeDir, _ := os.UserHomeDir()
|
||||
return filepath.Join(homeDir, "Applications", "KiroOAuthHandler.app")
|
||||
}
|
||||
|
||||
func isDarwinHandlerInstalled() bool {
|
||||
appPath := getDarwinAppPath()
|
||||
_, err := os.Stat(appPath)
|
||||
return err == nil
|
||||
}
|
||||
|
||||
func installDarwinHandler(handlerPort int) error {
|
||||
// Create app bundle structure
|
||||
appPath := getDarwinAppPath()
|
||||
contentsPath := filepath.Join(appPath, "Contents")
|
||||
macOSPath := filepath.Join(contentsPath, "MacOS")
|
||||
|
||||
if err := os.MkdirAll(macOSPath, 0755); err != nil {
|
||||
return fmt.Errorf("failed to create app bundle: %w", err)
|
||||
}
|
||||
|
||||
// Create Info.plist
|
||||
plistPath := filepath.Join(contentsPath, "Info.plist")
|
||||
plistContent := `<?xml version="1.0" encoding="UTF-8"?>
|
||||
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
|
||||
<plist version="1.0">
|
||||
<dict>
|
||||
<key>CFBundleIdentifier</key>
|
||||
<string>com.cliproxyapi.kiro-oauth-handler</string>
|
||||
<key>CFBundleName</key>
|
||||
<string>KiroOAuthHandler</string>
|
||||
<key>CFBundleExecutable</key>
|
||||
<string>kiro-oauth-handler</string>
|
||||
<key>CFBundleVersion</key>
|
||||
<string>1.0</string>
|
||||
<key>CFBundleURLTypes</key>
|
||||
<array>
|
||||
<dict>
|
||||
<key>CFBundleURLName</key>
|
||||
<string>Kiro Protocol</string>
|
||||
<key>CFBundleURLSchemes</key>
|
||||
<array>
|
||||
<string>kiro</string>
|
||||
</array>
|
||||
</dict>
|
||||
</array>
|
||||
<key>LSBackgroundOnly</key>
|
||||
<true/>
|
||||
</dict>
|
||||
</plist>`
|
||||
|
||||
if err := os.WriteFile(plistPath, []byte(plistContent), 0644); err != nil {
|
||||
return fmt.Errorf("failed to write Info.plist: %w", err)
|
||||
}
|
||||
|
||||
// Create executable script - tries multiple ports to handle dynamic port allocation
|
||||
execPath := filepath.Join(macOSPath, "kiro-oauth-handler")
|
||||
execContent := fmt.Sprintf(`#!/bin/bash
|
||||
# Kiro OAuth Protocol Handler for macOS
|
||||
|
||||
URL="$1"
|
||||
|
||||
# Check curl availability (should always exist on macOS)
|
||||
if [ ! -x /usr/bin/curl ]; then
|
||||
echo "Error: curl is required for Kiro OAuth handler" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Extract code and state from URL
|
||||
[[ "$URL" =~ code=([^&]+) ]] && CODE="${BASH_REMATCH[1]}"
|
||||
[[ "$URL" =~ state=([^&]+) ]] && STATE="${BASH_REMATCH[1]}"
|
||||
[[ "$URL" =~ error=([^&]+) ]] && ERROR="${BASH_REMATCH[1]}"
|
||||
|
||||
# Try multiple ports (default + dynamic range)
|
||||
for PORT in %d %d %d %d %d; do
|
||||
if [ -n "$ERROR" ]; then
|
||||
/usr/bin/curl -sf --connect-timeout 1 "http://127.0.0.1:$PORT/oauth/callback?error=$ERROR" && exit 0
|
||||
elif [ -n "$CODE" ] && [ -n "$STATE" ]; then
|
||||
/usr/bin/curl -sf --connect-timeout 1 "http://127.0.0.1:$PORT/oauth/callback?code=$CODE&state=$STATE" && exit 0
|
||||
fi
|
||||
done
|
||||
`, handlerPort, handlerPort+1, handlerPort+2, handlerPort+3, handlerPort+4)
|
||||
|
||||
if err := os.WriteFile(execPath, []byte(execContent), 0755); err != nil {
|
||||
return fmt.Errorf("failed to write executable: %w", err)
|
||||
}
|
||||
|
||||
// Register the app with Launch Services
|
||||
cmd := exec.Command("/System/Library/Frameworks/CoreServices.framework/Frameworks/LaunchServices.framework/Support/lsregister",
|
||||
"-f", appPath)
|
||||
if err := cmd.Run(); err != nil {
|
||||
log.Warnf("lsregister failed (handler may still work): %v", err)
|
||||
}
|
||||
|
||||
log.Info("Kiro protocol handler installed for macOS")
|
||||
return nil
|
||||
}
|
||||
|
||||
func uninstallDarwinHandler() error {
|
||||
appPath := getDarwinAppPath()
|
||||
|
||||
// Unregister from Launch Services
|
||||
cmd := exec.Command("/System/Library/Frameworks/CoreServices.framework/Frameworks/LaunchServices.framework/Support/lsregister",
|
||||
"-u", appPath)
|
||||
_ = cmd.Run()
|
||||
|
||||
// Remove app bundle
|
||||
if err := os.RemoveAll(appPath); err != nil && !os.IsNotExist(err) {
|
||||
return fmt.Errorf("failed to remove app bundle: %w", err)
|
||||
}
|
||||
|
||||
log.Info("Kiro protocol handler uninstalled")
|
||||
return nil
|
||||
}
|
||||
|
||||
// ParseKiroURI parses a kiro:// URI and extracts the callback parameters.
|
||||
func ParseKiroURI(rawURI string) (*AuthCallback, error) {
|
||||
u, err := url.Parse(rawURI)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid URI: %w", err)
|
||||
}
|
||||
|
||||
if u.Scheme != KiroProtocol {
|
||||
return nil, fmt.Errorf("invalid scheme: expected %s, got %s", KiroProtocol, u.Scheme)
|
||||
}
|
||||
|
||||
if u.Host != KiroAuthority {
|
||||
return nil, fmt.Errorf("invalid authority: expected %s, got %s", KiroAuthority, u.Host)
|
||||
}
|
||||
|
||||
query := u.Query()
|
||||
return &AuthCallback{
|
||||
Code: query.Get("code"),
|
||||
State: query.Get("state"),
|
||||
Error: query.Get("error"),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// GetHandlerInstructions returns platform-specific instructions for manual handler setup.
|
||||
func GetHandlerInstructions() string {
|
||||
switch runtime.GOOS {
|
||||
case "linux":
|
||||
return `To manually set up the Kiro protocol handler on Linux:
|
||||
|
||||
1. Create ~/.local/share/applications/kiro-oauth-handler.desktop:
|
||||
[Desktop Entry]
|
||||
Name=Kiro OAuth Handler
|
||||
Exec=~/.local/bin/kiro-oauth-handler %u
|
||||
Type=Application
|
||||
Terminal=false
|
||||
MimeType=x-scheme-handler/kiro;
|
||||
|
||||
2. Create ~/.local/bin/kiro-oauth-handler (make it executable):
|
||||
#!/bin/bash
|
||||
URL="$1"
|
||||
# ... (see generated script for full content)
|
||||
|
||||
3. Run: xdg-mime default kiro-oauth-handler.desktop x-scheme-handler/kiro`
|
||||
|
||||
case "windows":
|
||||
return `To manually set up the Kiro protocol handler on Windows:
|
||||
|
||||
1. Open Registry Editor (regedit.exe)
|
||||
2. Create key: HKEY_CURRENT_USER\Software\Classes\kiro
|
||||
3. Set default value to: URL:Kiro Protocol
|
||||
4. Create string value "URL Protocol" with empty data
|
||||
5. Create subkey: shell\open\command
|
||||
6. Set default value to: "C:\path\to\handler.bat" "%1"`
|
||||
|
||||
case "darwin":
|
||||
return `To manually set up the Kiro protocol handler on macOS:
|
||||
|
||||
1. Create ~/Applications/KiroOAuthHandler.app bundle
|
||||
2. Add Info.plist with CFBundleURLTypes containing "kiro" scheme
|
||||
3. Create executable in Contents/MacOS/
|
||||
4. Run: /System/Library/.../lsregister -f ~/Applications/KiroOAuthHandler.app`
|
||||
|
||||
default:
|
||||
return "Protocol handler setup is not supported on this platform."
|
||||
}
|
||||
}
|
||||
|
||||
// SetupProtocolHandlerIfNeeded checks and installs the protocol handler if needed.
|
||||
func SetupProtocolHandlerIfNeeded(handlerPort int) error {
|
||||
if IsProtocolHandlerInstalled() {
|
||||
log.Debug("Kiro protocol handler already installed")
|
||||
return nil
|
||||
}
|
||||
|
||||
fmt.Println("\n╔══════════════════════════════════════════════════════════╗")
|
||||
fmt.Println("║ Kiro Protocol Handler Setup Required ║")
|
||||
fmt.Println("╚══════════════════════════════════════════════════════════╝")
|
||||
fmt.Println("\nTo enable Google/GitHub login, we need to install a protocol handler.")
|
||||
fmt.Println("This allows your browser to redirect back to the CLI after authentication.")
|
||||
fmt.Println("\nInstalling protocol handler...")
|
||||
|
||||
if err := InstallProtocolHandler(handlerPort); err != nil {
|
||||
fmt.Printf("\n⚠ Automatic installation failed: %v\n", err)
|
||||
fmt.Println("\nManual setup instructions:")
|
||||
fmt.Println(strings.Repeat("-", 60))
|
||||
fmt.Println(GetHandlerInstructions())
|
||||
return err
|
||||
}
|
||||
|
||||
fmt.Println("\n✓ Protocol handler installed successfully!")
|
||||
return nil
|
||||
}
|
||||
403
internal/auth/kiro/social_auth.go
Normal file
403
internal/auth/kiro/social_auth.go
Normal file
@@ -0,0 +1,403 @@
|
||||
// Package kiro provides social authentication (Google/GitHub) for Kiro via AuthServiceClient.
|
||||
package kiro
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"os/exec"
|
||||
"runtime"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/browser"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"golang.org/x/term"
|
||||
)
|
||||
|
||||
const (
|
||||
// Kiro AuthService endpoint
|
||||
kiroAuthServiceEndpoint = "https://prod.us-east-1.auth.desktop.kiro.dev"
|
||||
|
||||
// OAuth timeout
|
||||
socialAuthTimeout = 10 * time.Minute
|
||||
)
|
||||
|
||||
// SocialProvider represents the social login provider.
|
||||
type SocialProvider string
|
||||
|
||||
const (
|
||||
// ProviderGoogle is Google OAuth provider
|
||||
ProviderGoogle SocialProvider = "Google"
|
||||
// ProviderGitHub is GitHub OAuth provider
|
||||
ProviderGitHub SocialProvider = "Github"
|
||||
// Note: AWS Builder ID is NOT supported by Kiro's auth service.
|
||||
// It only supports: Google, Github, Cognito
|
||||
// AWS Builder ID must use device code flow via SSO OIDC.
|
||||
)
|
||||
|
||||
// CreateTokenRequest is sent to Kiro's /oauth/token endpoint.
|
||||
type CreateTokenRequest struct {
|
||||
Code string `json:"code"`
|
||||
CodeVerifier string `json:"code_verifier"`
|
||||
RedirectURI string `json:"redirect_uri"`
|
||||
InvitationCode string `json:"invitation_code,omitempty"`
|
||||
}
|
||||
|
||||
// SocialTokenResponse from Kiro's /oauth/token endpoint for social auth.
|
||||
type SocialTokenResponse struct {
|
||||
AccessToken string `json:"accessToken"`
|
||||
RefreshToken string `json:"refreshToken"`
|
||||
ProfileArn string `json:"profileArn"`
|
||||
ExpiresIn int `json:"expiresIn"`
|
||||
}
|
||||
|
||||
// RefreshTokenRequest is sent to Kiro's /refreshToken endpoint.
|
||||
type RefreshTokenRequest struct {
|
||||
RefreshToken string `json:"refreshToken"`
|
||||
}
|
||||
|
||||
// SocialAuthClient handles social authentication with Kiro.
|
||||
type SocialAuthClient struct {
|
||||
httpClient *http.Client
|
||||
cfg *config.Config
|
||||
protocolHandler *ProtocolHandler
|
||||
}
|
||||
|
||||
// NewSocialAuthClient creates a new social auth client.
|
||||
func NewSocialAuthClient(cfg *config.Config) *SocialAuthClient {
|
||||
client := &http.Client{Timeout: 30 * time.Second}
|
||||
if cfg != nil {
|
||||
client = util.SetProxy(&cfg.SDKConfig, client)
|
||||
}
|
||||
return &SocialAuthClient{
|
||||
httpClient: client,
|
||||
cfg: cfg,
|
||||
protocolHandler: NewProtocolHandler(),
|
||||
}
|
||||
}
|
||||
|
||||
// generatePKCE generates PKCE code verifier and challenge.
|
||||
func generatePKCE() (verifier, challenge string, err error) {
|
||||
// Generate 32 bytes of random data for verifier
|
||||
b := make([]byte, 32)
|
||||
if _, err := rand.Read(b); err != nil {
|
||||
return "", "", fmt.Errorf("failed to generate random bytes: %w", err)
|
||||
}
|
||||
verifier = base64.RawURLEncoding.EncodeToString(b)
|
||||
|
||||
// Generate SHA256 hash of verifier for challenge
|
||||
h := sha256.Sum256([]byte(verifier))
|
||||
challenge = base64.RawURLEncoding.EncodeToString(h[:])
|
||||
|
||||
return verifier, challenge, nil
|
||||
}
|
||||
|
||||
// generateState generates a random state parameter.
|
||||
func generateStateParam() (string, error) {
|
||||
b := make([]byte, 16)
|
||||
if _, err := rand.Read(b); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return base64.RawURLEncoding.EncodeToString(b), nil
|
||||
}
|
||||
|
||||
// buildLoginURL constructs the Kiro OAuth login URL.
|
||||
// The login endpoint expects a GET request with query parameters.
|
||||
// Format: /login?idp=Google&redirect_uri=...&code_challenge=...&code_challenge_method=S256&state=...&prompt=select_account
|
||||
// The prompt=select_account parameter forces the account selection screen even if already logged in.
|
||||
func (c *SocialAuthClient) buildLoginURL(provider, redirectURI, codeChallenge, state string) string {
|
||||
return fmt.Sprintf("%s/login?idp=%s&redirect_uri=%s&code_challenge=%s&code_challenge_method=S256&state=%s&prompt=select_account",
|
||||
kiroAuthServiceEndpoint,
|
||||
provider,
|
||||
url.QueryEscape(redirectURI),
|
||||
codeChallenge,
|
||||
state,
|
||||
)
|
||||
}
|
||||
|
||||
// createToken exchanges the authorization code for tokens.
|
||||
func (c *SocialAuthClient) createToken(ctx context.Context, req *CreateTokenRequest) (*SocialTokenResponse, error) {
|
||||
body, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal token request: %w", err)
|
||||
}
|
||||
|
||||
tokenURL := kiroAuthServiceEndpoint + "/oauth/token"
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenURL, strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create token request: %w", err)
|
||||
}
|
||||
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("User-Agent", "cli-proxy-api/1.0.0")
|
||||
|
||||
resp, err := c.httpClient.Do(httpReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("token request failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read token response: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
log.Debugf("token exchange failed (status %d): %s", resp.StatusCode, string(respBody))
|
||||
return nil, fmt.Errorf("token exchange failed (status %d)", resp.StatusCode)
|
||||
}
|
||||
|
||||
var tokenResp SocialTokenResponse
|
||||
if err := json.Unmarshal(respBody, &tokenResp); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse token response: %w", err)
|
||||
}
|
||||
|
||||
return &tokenResp, nil
|
||||
}
|
||||
|
||||
// RefreshSocialToken refreshes an expired social auth token.
|
||||
func (c *SocialAuthClient) RefreshSocialToken(ctx context.Context, refreshToken string) (*KiroTokenData, error) {
|
||||
body, err := json.Marshal(&RefreshTokenRequest{RefreshToken: refreshToken})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal refresh request: %w", err)
|
||||
}
|
||||
|
||||
refreshURL := kiroAuthServiceEndpoint + "/refreshToken"
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, refreshURL, strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create refresh request: %w", err)
|
||||
}
|
||||
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("User-Agent", "cli-proxy-api/1.0.0")
|
||||
|
||||
resp, err := c.httpClient.Do(httpReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("refresh request failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read refresh response: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
log.Debugf("token refresh failed (status %d): %s", resp.StatusCode, string(respBody))
|
||||
return nil, fmt.Errorf("token refresh failed (status %d)", resp.StatusCode)
|
||||
}
|
||||
|
||||
var tokenResp SocialTokenResponse
|
||||
if err := json.Unmarshal(respBody, &tokenResp); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse refresh response: %w", err)
|
||||
}
|
||||
|
||||
// Validate ExpiresIn - use default 1 hour if invalid
|
||||
expiresIn := tokenResp.ExpiresIn
|
||||
if expiresIn <= 0 {
|
||||
expiresIn = 3600 // Default 1 hour
|
||||
}
|
||||
expiresAt := time.Now().Add(time.Duration(expiresIn) * time.Second)
|
||||
|
||||
return &KiroTokenData{
|
||||
AccessToken: tokenResp.AccessToken,
|
||||
RefreshToken: tokenResp.RefreshToken,
|
||||
ProfileArn: tokenResp.ProfileArn,
|
||||
ExpiresAt: expiresAt.Format(time.RFC3339),
|
||||
AuthMethod: "social",
|
||||
Provider: "", // Caller should preserve original provider
|
||||
}, nil
|
||||
}
|
||||
|
||||
// LoginWithSocial performs OAuth login with Google.
|
||||
func (c *SocialAuthClient) LoginWithSocial(ctx context.Context, provider SocialProvider) (*KiroTokenData, error) {
|
||||
providerName := string(provider)
|
||||
|
||||
fmt.Println("\n╔══════════════════════════════════════════════════════════╗")
|
||||
fmt.Printf("║ Kiro Authentication (%s) ║\n", providerName)
|
||||
fmt.Println("╚══════════════════════════════════════════════════════════╝")
|
||||
|
||||
// Step 1: Setup protocol handler
|
||||
fmt.Println("\nSetting up authentication...")
|
||||
|
||||
// Start the local callback server
|
||||
handlerPort, err := c.protocolHandler.Start(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to start callback server: %w", err)
|
||||
}
|
||||
defer c.protocolHandler.Stop()
|
||||
|
||||
// Ensure protocol handler is installed and set as default
|
||||
if err := SetupProtocolHandlerIfNeeded(handlerPort); err != nil {
|
||||
fmt.Println("\n⚠ Protocol handler setup failed. Trying alternative method...")
|
||||
fmt.Println(" If you see a browser 'Open with' dialog, select your default browser.")
|
||||
fmt.Println(" For manual setup instructions, run: cliproxy kiro --help-protocol")
|
||||
log.Debugf("kiro: protocol handler setup error: %v", err)
|
||||
// Continue anyway - user might have set it up manually or select browser manually
|
||||
} else {
|
||||
// Force set our handler as default (prevents "Open with" dialog)
|
||||
forceDefaultProtocolHandler()
|
||||
}
|
||||
|
||||
// Step 2: Generate PKCE codes
|
||||
codeVerifier, codeChallenge, err := generatePKCE()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to generate PKCE: %w", err)
|
||||
}
|
||||
|
||||
// Step 3: Generate state
|
||||
state, err := generateStateParam()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to generate state: %w", err)
|
||||
}
|
||||
|
||||
// Step 4: Build the login URL (Kiro uses GET request with query params)
|
||||
authURL := c.buildLoginURL(providerName, KiroRedirectURI, codeChallenge, state)
|
||||
|
||||
// Set incognito mode based on config (defaults to true for Kiro, can be overridden with --no-incognito)
|
||||
// Incognito mode enables multi-account support by bypassing cached sessions
|
||||
if c.cfg != nil {
|
||||
browser.SetIncognitoMode(c.cfg.IncognitoBrowser)
|
||||
if !c.cfg.IncognitoBrowser {
|
||||
log.Info("kiro: using normal browser mode (--no-incognito). Note: You may not be able to select a different account.")
|
||||
} else {
|
||||
log.Debug("kiro: using incognito mode for multi-account support")
|
||||
}
|
||||
} else {
|
||||
browser.SetIncognitoMode(true) // Default to incognito if no config
|
||||
log.Debug("kiro: using incognito mode for multi-account support (default)")
|
||||
}
|
||||
|
||||
// Step 5: Open browser for user authentication
|
||||
fmt.Println("\n════════════════════════════════════════════════════════════")
|
||||
fmt.Printf(" Opening browser for %s authentication...\n", providerName)
|
||||
fmt.Println("════════════════════════════════════════════════════════════")
|
||||
fmt.Printf("\n URL: %s\n\n", authURL)
|
||||
|
||||
if err := browser.OpenURL(authURL); err != nil {
|
||||
log.Warnf("Could not open browser automatically: %v", err)
|
||||
fmt.Println(" ⚠ Could not open browser automatically.")
|
||||
fmt.Println(" Please open the URL above in your browser manually.")
|
||||
} else {
|
||||
fmt.Println(" (Browser opened automatically)")
|
||||
}
|
||||
|
||||
fmt.Println("\n Waiting for authentication callback...")
|
||||
|
||||
// Step 6: Wait for callback
|
||||
callback, err := c.protocolHandler.WaitForCallback(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to receive callback: %w", err)
|
||||
}
|
||||
|
||||
if callback.Error != "" {
|
||||
return nil, fmt.Errorf("authentication error: %s", callback.Error)
|
||||
}
|
||||
|
||||
if callback.State != state {
|
||||
// Log state values for debugging, but don't expose in user-facing error
|
||||
log.Debugf("kiro: OAuth state mismatch - expected %s, got %s", state, callback.State)
|
||||
return nil, fmt.Errorf("OAuth state validation failed - please try again")
|
||||
}
|
||||
|
||||
if callback.Code == "" {
|
||||
return nil, fmt.Errorf("no authorization code received")
|
||||
}
|
||||
|
||||
fmt.Println("\n✓ Authorization received!")
|
||||
|
||||
// Step 7: Exchange code for tokens
|
||||
fmt.Println("Exchanging code for tokens...")
|
||||
|
||||
tokenReq := &CreateTokenRequest{
|
||||
Code: callback.Code,
|
||||
CodeVerifier: codeVerifier,
|
||||
RedirectURI: KiroRedirectURI,
|
||||
}
|
||||
|
||||
tokenResp, err := c.createToken(ctx, tokenReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to exchange code for tokens: %w", err)
|
||||
}
|
||||
|
||||
fmt.Println("\n✓ Authentication successful!")
|
||||
|
||||
// Close the browser window
|
||||
if err := browser.CloseBrowser(); err != nil {
|
||||
log.Debugf("Failed to close browser: %v", err)
|
||||
}
|
||||
|
||||
// Validate ExpiresIn - use default 1 hour if invalid
|
||||
expiresIn := tokenResp.ExpiresIn
|
||||
if expiresIn <= 0 {
|
||||
expiresIn = 3600
|
||||
}
|
||||
expiresAt := time.Now().Add(time.Duration(expiresIn) * time.Second)
|
||||
|
||||
// Try to extract email from JWT access token first
|
||||
email := ExtractEmailFromJWT(tokenResp.AccessToken)
|
||||
|
||||
// If no email in JWT, ask user for account label (only in interactive mode)
|
||||
if email == "" && isInteractiveTerminal() {
|
||||
fmt.Print("\n Enter account label for file naming (optional, press Enter to skip): ")
|
||||
reader := bufio.NewReader(os.Stdin)
|
||||
var err error
|
||||
email, err = reader.ReadString('\n')
|
||||
if err != nil {
|
||||
log.Debugf("Failed to read account label: %v", err)
|
||||
}
|
||||
email = strings.TrimSpace(email)
|
||||
}
|
||||
|
||||
return &KiroTokenData{
|
||||
AccessToken: tokenResp.AccessToken,
|
||||
RefreshToken: tokenResp.RefreshToken,
|
||||
ProfileArn: tokenResp.ProfileArn,
|
||||
ExpiresAt: expiresAt.Format(time.RFC3339),
|
||||
AuthMethod: "social",
|
||||
Provider: providerName,
|
||||
Email: email, // JWT email or user-provided label
|
||||
}, nil
|
||||
}
|
||||
|
||||
// LoginWithGoogle performs OAuth login with Google.
|
||||
func (c *SocialAuthClient) LoginWithGoogle(ctx context.Context) (*KiroTokenData, error) {
|
||||
return c.LoginWithSocial(ctx, ProviderGoogle)
|
||||
}
|
||||
|
||||
// LoginWithGitHub performs OAuth login with GitHub.
|
||||
func (c *SocialAuthClient) LoginWithGitHub(ctx context.Context) (*KiroTokenData, error) {
|
||||
return c.LoginWithSocial(ctx, ProviderGitHub)
|
||||
}
|
||||
|
||||
// forceDefaultProtocolHandler sets our protocol handler as the default for kiro:// URLs.
|
||||
// This prevents the "Open with" dialog from appearing on Linux.
|
||||
// On non-Linux platforms, this is a no-op as they use different mechanisms.
|
||||
func forceDefaultProtocolHandler() {
|
||||
if runtime.GOOS != "linux" {
|
||||
return // Non-Linux platforms use different handler mechanisms
|
||||
}
|
||||
|
||||
// Set our handler as default using xdg-mime
|
||||
cmd := exec.Command("xdg-mime", "default", "kiro-oauth-handler.desktop", "x-scheme-handler/kiro")
|
||||
if err := cmd.Run(); err != nil {
|
||||
log.Warnf("Failed to set default protocol handler: %v. You may see a handler selection dialog.", err)
|
||||
}
|
||||
}
|
||||
|
||||
// isInteractiveTerminal checks if stdin is connected to an interactive terminal.
|
||||
// Returns false in CI/automated environments or when stdin is piped.
|
||||
func isInteractiveTerminal() bool {
|
||||
return term.IsTerminal(int(os.Stdin.Fd()))
|
||||
}
|
||||
527
internal/auth/kiro/sso_oidc.go
Normal file
527
internal/auth/kiro/sso_oidc.go
Normal file
@@ -0,0 +1,527 @@
|
||||
// Package kiro provides AWS SSO OIDC authentication for Kiro.
|
||||
package kiro
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/browser"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
const (
|
||||
// AWS SSO OIDC endpoints
|
||||
ssoOIDCEndpoint = "https://oidc.us-east-1.amazonaws.com"
|
||||
|
||||
// Kiro's start URL for Builder ID
|
||||
builderIDStartURL = "https://view.awsapps.com/start"
|
||||
|
||||
// Polling interval
|
||||
pollInterval = 5 * time.Second
|
||||
)
|
||||
|
||||
// SSOOIDCClient handles AWS SSO OIDC authentication.
|
||||
type SSOOIDCClient struct {
|
||||
httpClient *http.Client
|
||||
cfg *config.Config
|
||||
}
|
||||
|
||||
// NewSSOOIDCClient creates a new SSO OIDC client.
|
||||
func NewSSOOIDCClient(cfg *config.Config) *SSOOIDCClient {
|
||||
client := &http.Client{Timeout: 30 * time.Second}
|
||||
if cfg != nil {
|
||||
client = util.SetProxy(&cfg.SDKConfig, client)
|
||||
}
|
||||
return &SSOOIDCClient{
|
||||
httpClient: client,
|
||||
cfg: cfg,
|
||||
}
|
||||
}
|
||||
|
||||
// RegisterClientResponse from AWS SSO OIDC.
|
||||
type RegisterClientResponse struct {
|
||||
ClientID string `json:"clientId"`
|
||||
ClientSecret string `json:"clientSecret"`
|
||||
ClientIDIssuedAt int64 `json:"clientIdIssuedAt"`
|
||||
ClientSecretExpiresAt int64 `json:"clientSecretExpiresAt"`
|
||||
}
|
||||
|
||||
// StartDeviceAuthResponse from AWS SSO OIDC.
|
||||
type StartDeviceAuthResponse struct {
|
||||
DeviceCode string `json:"deviceCode"`
|
||||
UserCode string `json:"userCode"`
|
||||
VerificationURI string `json:"verificationUri"`
|
||||
VerificationURIComplete string `json:"verificationUriComplete"`
|
||||
ExpiresIn int `json:"expiresIn"`
|
||||
Interval int `json:"interval"`
|
||||
}
|
||||
|
||||
// CreateTokenResponse from AWS SSO OIDC.
|
||||
type CreateTokenResponse struct {
|
||||
AccessToken string `json:"accessToken"`
|
||||
TokenType string `json:"tokenType"`
|
||||
ExpiresIn int `json:"expiresIn"`
|
||||
RefreshToken string `json:"refreshToken"`
|
||||
}
|
||||
|
||||
// RegisterClient registers a new OIDC client with AWS.
|
||||
func (c *SSOOIDCClient) RegisterClient(ctx context.Context) (*RegisterClientResponse, error) {
|
||||
// Generate unique client name for each registration to support multiple accounts
|
||||
clientName := fmt.Sprintf("CLI-Proxy-API-%d", time.Now().UnixNano())
|
||||
|
||||
payload := map[string]interface{}{
|
||||
"clientName": clientName,
|
||||
"clientType": "public",
|
||||
"scopes": []string{"codewhisperer:completions", "codewhisperer:analysis", "codewhisperer:conversations"},
|
||||
}
|
||||
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/client/register", strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
log.Debugf("register client failed (status %d): %s", resp.StatusCode, string(respBody))
|
||||
return nil, fmt.Errorf("register client failed (status %d)", resp.StatusCode)
|
||||
}
|
||||
|
||||
var result RegisterClientResponse
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
// StartDeviceAuthorization starts the device authorization flow.
|
||||
func (c *SSOOIDCClient) StartDeviceAuthorization(ctx context.Context, clientID, clientSecret string) (*StartDeviceAuthResponse, error) {
|
||||
payload := map[string]string{
|
||||
"clientId": clientID,
|
||||
"clientSecret": clientSecret,
|
||||
"startUrl": builderIDStartURL,
|
||||
}
|
||||
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/device_authorization", strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
log.Debugf("start device auth failed (status %d): %s", resp.StatusCode, string(respBody))
|
||||
return nil, fmt.Errorf("start device auth failed (status %d)", resp.StatusCode)
|
||||
}
|
||||
|
||||
var result StartDeviceAuthResponse
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
// CreateToken polls for the access token after user authorization.
|
||||
func (c *SSOOIDCClient) CreateToken(ctx context.Context, clientID, clientSecret, deviceCode string) (*CreateTokenResponse, error) {
|
||||
payload := map[string]string{
|
||||
"clientId": clientID,
|
||||
"clientSecret": clientSecret,
|
||||
"deviceCode": deviceCode,
|
||||
"grantType": "urn:ietf:params:oauth:grant-type:device_code",
|
||||
}
|
||||
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/token", strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Check for pending authorization
|
||||
if resp.StatusCode == http.StatusBadRequest {
|
||||
var errResp struct {
|
||||
Error string `json:"error"`
|
||||
}
|
||||
if json.Unmarshal(respBody, &errResp) == nil {
|
||||
if errResp.Error == "authorization_pending" {
|
||||
return nil, fmt.Errorf("authorization_pending")
|
||||
}
|
||||
if errResp.Error == "slow_down" {
|
||||
return nil, fmt.Errorf("slow_down")
|
||||
}
|
||||
}
|
||||
log.Debugf("create token failed: %s", string(respBody))
|
||||
return nil, fmt.Errorf("create token failed")
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
log.Debugf("create token failed (status %d): %s", resp.StatusCode, string(respBody))
|
||||
return nil, fmt.Errorf("create token failed (status %d)", resp.StatusCode)
|
||||
}
|
||||
|
||||
var result CreateTokenResponse
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
// RefreshToken refreshes an access token using the refresh token.
|
||||
func (c *SSOOIDCClient) RefreshToken(ctx context.Context, clientID, clientSecret, refreshToken string) (*KiroTokenData, error) {
|
||||
payload := map[string]string{
|
||||
"clientId": clientID,
|
||||
"clientSecret": clientSecret,
|
||||
"refreshToken": refreshToken,
|
||||
"grantType": "refresh_token",
|
||||
}
|
||||
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/token", strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
log.Debugf("token refresh failed (status %d): %s", resp.StatusCode, string(respBody))
|
||||
return nil, fmt.Errorf("token refresh failed (status %d)", resp.StatusCode)
|
||||
}
|
||||
|
||||
var result CreateTokenResponse
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
expiresAt := time.Now().Add(time.Duration(result.ExpiresIn) * time.Second)
|
||||
|
||||
return &KiroTokenData{
|
||||
AccessToken: result.AccessToken,
|
||||
RefreshToken: result.RefreshToken,
|
||||
ExpiresAt: expiresAt.Format(time.RFC3339),
|
||||
AuthMethod: "builder-id",
|
||||
Provider: "AWS",
|
||||
ClientID: clientID,
|
||||
ClientSecret: clientSecret,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// LoginWithBuilderID performs the full device code flow for AWS Builder ID.
|
||||
func (c *SSOOIDCClient) LoginWithBuilderID(ctx context.Context) (*KiroTokenData, error) {
|
||||
fmt.Println("\n╔══════════════════════════════════════════════════════════╗")
|
||||
fmt.Println("║ Kiro Authentication (AWS Builder ID) ║")
|
||||
fmt.Println("╚══════════════════════════════════════════════════════════╝")
|
||||
|
||||
// Step 1: Register client
|
||||
fmt.Println("\nRegistering client...")
|
||||
regResp, err := c.RegisterClient(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to register client: %w", err)
|
||||
}
|
||||
log.Debugf("Client registered: %s", regResp.ClientID)
|
||||
|
||||
// Step 2: Start device authorization
|
||||
fmt.Println("Starting device authorization...")
|
||||
authResp, err := c.StartDeviceAuthorization(ctx, regResp.ClientID, regResp.ClientSecret)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to start device auth: %w", err)
|
||||
}
|
||||
|
||||
// Step 3: Show user the verification URL
|
||||
fmt.Printf("\n")
|
||||
fmt.Println("════════════════════════════════════════════════════════════")
|
||||
fmt.Printf(" Open this URL in your browser:\n")
|
||||
fmt.Printf(" %s\n", authResp.VerificationURIComplete)
|
||||
fmt.Println("════════════════════════════════════════════════════════════")
|
||||
fmt.Printf("\n Or go to: %s\n", authResp.VerificationURI)
|
||||
fmt.Printf(" And enter code: %s\n\n", authResp.UserCode)
|
||||
|
||||
// Set incognito mode based on config (defaults to true for Kiro, can be overridden with --no-incognito)
|
||||
// Incognito mode enables multi-account support by bypassing cached sessions
|
||||
if c.cfg != nil {
|
||||
browser.SetIncognitoMode(c.cfg.IncognitoBrowser)
|
||||
if !c.cfg.IncognitoBrowser {
|
||||
log.Info("kiro: using normal browser mode (--no-incognito). Note: You may not be able to select a different account.")
|
||||
} else {
|
||||
log.Debug("kiro: using incognito mode for multi-account support")
|
||||
}
|
||||
} else {
|
||||
browser.SetIncognitoMode(true) // Default to incognito if no config
|
||||
log.Debug("kiro: using incognito mode for multi-account support (default)")
|
||||
}
|
||||
|
||||
// Open browser using cross-platform browser package
|
||||
if err := browser.OpenURL(authResp.VerificationURIComplete); err != nil {
|
||||
log.Warnf("Could not open browser automatically: %v", err)
|
||||
fmt.Println(" Please open the URL manually in your browser.")
|
||||
} else {
|
||||
fmt.Println(" (Browser opened automatically)")
|
||||
}
|
||||
|
||||
// Step 4: Poll for token
|
||||
fmt.Println("Waiting for authorization...")
|
||||
|
||||
interval := pollInterval
|
||||
if authResp.Interval > 0 {
|
||||
interval = time.Duration(authResp.Interval) * time.Second
|
||||
}
|
||||
|
||||
deadline := time.Now().Add(time.Duration(authResp.ExpiresIn) * time.Second)
|
||||
|
||||
for time.Now().Before(deadline) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
browser.CloseBrowser() // Cleanup on cancel
|
||||
return nil, ctx.Err()
|
||||
case <-time.After(interval):
|
||||
tokenResp, err := c.CreateToken(ctx, regResp.ClientID, regResp.ClientSecret, authResp.DeviceCode)
|
||||
if err != nil {
|
||||
errStr := err.Error()
|
||||
if strings.Contains(errStr, "authorization_pending") {
|
||||
fmt.Print(".")
|
||||
continue
|
||||
}
|
||||
if strings.Contains(errStr, "slow_down") {
|
||||
interval += 5 * time.Second
|
||||
continue
|
||||
}
|
||||
// Close browser on error before returning
|
||||
browser.CloseBrowser()
|
||||
return nil, fmt.Errorf("token creation failed: %w", err)
|
||||
}
|
||||
|
||||
fmt.Println("\n\n✓ Authorization successful!")
|
||||
|
||||
// Close the browser window
|
||||
if err := browser.CloseBrowser(); err != nil {
|
||||
log.Debugf("Failed to close browser: %v", err)
|
||||
}
|
||||
|
||||
// Step 5: Get profile ARN from CodeWhisperer API
|
||||
fmt.Println("Fetching profile information...")
|
||||
profileArn := c.fetchProfileArn(ctx, tokenResp.AccessToken)
|
||||
|
||||
// Extract email from JWT access token
|
||||
email := ExtractEmailFromJWT(tokenResp.AccessToken)
|
||||
if email != "" {
|
||||
fmt.Printf(" Logged in as: %s\n", email)
|
||||
}
|
||||
|
||||
expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second)
|
||||
|
||||
return &KiroTokenData{
|
||||
AccessToken: tokenResp.AccessToken,
|
||||
RefreshToken: tokenResp.RefreshToken,
|
||||
ProfileArn: profileArn,
|
||||
ExpiresAt: expiresAt.Format(time.RFC3339),
|
||||
AuthMethod: "builder-id",
|
||||
Provider: "AWS",
|
||||
ClientID: regResp.ClientID,
|
||||
ClientSecret: regResp.ClientSecret,
|
||||
Email: email,
|
||||
}, nil
|
||||
}
|
||||
}
|
||||
|
||||
// Close browser on timeout for better UX
|
||||
if err := browser.CloseBrowser(); err != nil {
|
||||
log.Debugf("Failed to close browser on timeout: %v", err)
|
||||
}
|
||||
return nil, fmt.Errorf("authorization timed out")
|
||||
}
|
||||
|
||||
// fetchProfileArn retrieves the profile ARN from CodeWhisperer API.
|
||||
// This is needed for file naming since AWS SSO OIDC doesn't return profile info.
|
||||
func (c *SSOOIDCClient) fetchProfileArn(ctx context.Context, accessToken string) string {
|
||||
// Try ListProfiles API first
|
||||
profileArn := c.tryListProfiles(ctx, accessToken)
|
||||
if profileArn != "" {
|
||||
return profileArn
|
||||
}
|
||||
|
||||
// Fallback: Try ListAvailableCustomizations
|
||||
return c.tryListCustomizations(ctx, accessToken)
|
||||
}
|
||||
|
||||
func (c *SSOOIDCClient) tryListProfiles(ctx context.Context, accessToken string) string {
|
||||
payload := map[string]interface{}{
|
||||
"origin": "AI_EDITOR",
|
||||
}
|
||||
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://codewhisperer.us-east-1.amazonaws.com", strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/x-amz-json-1.0")
|
||||
req.Header.Set("x-amz-target", "AmazonCodeWhispererService.ListProfiles")
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
log.Debugf("ListProfiles failed (status %d): %s", resp.StatusCode, string(respBody))
|
||||
return ""
|
||||
}
|
||||
|
||||
log.Debugf("ListProfiles response: %s", string(respBody))
|
||||
|
||||
var result struct {
|
||||
Profiles []struct {
|
||||
Arn string `json:"arn"`
|
||||
} `json:"profiles"`
|
||||
ProfileArn string `json:"profileArn"`
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
if result.ProfileArn != "" {
|
||||
return result.ProfileArn
|
||||
}
|
||||
|
||||
if len(result.Profiles) > 0 {
|
||||
return result.Profiles[0].Arn
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
func (c *SSOOIDCClient) tryListCustomizations(ctx context.Context, accessToken string) string {
|
||||
payload := map[string]interface{}{
|
||||
"origin": "AI_EDITOR",
|
||||
}
|
||||
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://codewhisperer.us-east-1.amazonaws.com", strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/x-amz-json-1.0")
|
||||
req.Header.Set("x-amz-target", "AmazonCodeWhispererService.ListAvailableCustomizations")
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
log.Debugf("ListAvailableCustomizations failed (status %d): %s", resp.StatusCode, string(respBody))
|
||||
return ""
|
||||
}
|
||||
|
||||
log.Debugf("ListAvailableCustomizations response: %s", string(respBody))
|
||||
|
||||
var result struct {
|
||||
Customizations []struct {
|
||||
Arn string `json:"arn"`
|
||||
} `json:"customizations"`
|
||||
ProfileArn string `json:"profileArn"`
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
if result.ProfileArn != "" {
|
||||
return result.ProfileArn
|
||||
}
|
||||
|
||||
if len(result.Customizations) > 0 {
|
||||
return result.Customizations[0].Arn
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
72
internal/auth/kiro/token.go
Normal file
72
internal/auth/kiro/token.go
Normal file
@@ -0,0 +1,72 @@
|
||||
package kiro
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
)
|
||||
|
||||
// KiroTokenStorage holds the persistent token data for Kiro authentication.
|
||||
type KiroTokenStorage struct {
|
||||
// AccessToken is the OAuth2 access token for API access
|
||||
AccessToken string `json:"access_token"`
|
||||
// RefreshToken is used to obtain new access tokens
|
||||
RefreshToken string `json:"refresh_token"`
|
||||
// ProfileArn is the AWS CodeWhisperer profile ARN
|
||||
ProfileArn string `json:"profile_arn"`
|
||||
// ExpiresAt is the timestamp when the token expires
|
||||
ExpiresAt string `json:"expires_at"`
|
||||
// AuthMethod indicates the authentication method used
|
||||
AuthMethod string `json:"auth_method"`
|
||||
// Provider indicates the OAuth provider
|
||||
Provider string `json:"provider"`
|
||||
// LastRefresh is the timestamp of the last token refresh
|
||||
LastRefresh string `json:"last_refresh"`
|
||||
}
|
||||
|
||||
// SaveTokenToFile persists the token storage to the specified file path.
|
||||
func (s *KiroTokenStorage) SaveTokenToFile(authFilePath string) error {
|
||||
dir := filepath.Dir(authFilePath)
|
||||
if err := os.MkdirAll(dir, 0700); err != nil {
|
||||
return fmt.Errorf("failed to create directory: %w", err)
|
||||
}
|
||||
|
||||
data, err := json.MarshalIndent(s, "", " ")
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal token storage: %w", err)
|
||||
}
|
||||
|
||||
if err := os.WriteFile(authFilePath, data, 0600); err != nil {
|
||||
return fmt.Errorf("failed to write token file: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// LoadFromFile loads token storage from the specified file path.
|
||||
func LoadFromFile(authFilePath string) (*KiroTokenStorage, error) {
|
||||
data, err := os.ReadFile(authFilePath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read token file: %w", err)
|
||||
}
|
||||
|
||||
var storage KiroTokenStorage
|
||||
if err := json.Unmarshal(data, &storage); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse token file: %w", err)
|
||||
}
|
||||
|
||||
return &storage, nil
|
||||
}
|
||||
|
||||
// ToTokenData converts storage to KiroTokenData for API use.
|
||||
func (s *KiroTokenStorage) ToTokenData() *KiroTokenData {
|
||||
return &KiroTokenData{
|
||||
AccessToken: s.AccessToken,
|
||||
RefreshToken: s.RefreshToken,
|
||||
ProfileArn: s.ProfileArn,
|
||||
ExpiresAt: s.ExpiresAt,
|
||||
AuthMethod: s.AuthMethod,
|
||||
Provider: s.Provider,
|
||||
}
|
||||
}
|
||||
@@ -6,14 +6,48 @@ import (
|
||||
"fmt"
|
||||
"os/exec"
|
||||
"runtime"
|
||||
"sync"
|
||||
|
||||
pkgbrowser "github.com/pkg/browser"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/skratchdot/open-golang/open"
|
||||
)
|
||||
|
||||
// incognitoMode controls whether to open URLs in incognito/private mode.
|
||||
// This is useful for OAuth flows where you want to use a different account.
|
||||
var incognitoMode bool
|
||||
|
||||
// lastBrowserProcess stores the last opened browser process for cleanup
|
||||
var lastBrowserProcess *exec.Cmd
|
||||
var browserMutex sync.Mutex
|
||||
|
||||
// SetIncognitoMode enables or disables incognito/private browsing mode.
|
||||
func SetIncognitoMode(enabled bool) {
|
||||
incognitoMode = enabled
|
||||
}
|
||||
|
||||
// IsIncognitoMode returns whether incognito mode is enabled.
|
||||
func IsIncognitoMode() bool {
|
||||
return incognitoMode
|
||||
}
|
||||
|
||||
// CloseBrowser closes the last opened browser process.
|
||||
func CloseBrowser() error {
|
||||
browserMutex.Lock()
|
||||
defer browserMutex.Unlock()
|
||||
|
||||
if lastBrowserProcess == nil || lastBrowserProcess.Process == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
err := lastBrowserProcess.Process.Kill()
|
||||
lastBrowserProcess = nil
|
||||
return err
|
||||
}
|
||||
|
||||
// OpenURL opens the specified URL in the default web browser.
|
||||
// It first attempts to use a platform-agnostic library and falls back to
|
||||
// platform-specific commands if that fails.
|
||||
// It uses the pkg/browser library which provides robust cross-platform support
|
||||
// for Windows, macOS, and Linux.
|
||||
// If incognito mode is enabled, it will open in a private/incognito window.
|
||||
//
|
||||
// Parameters:
|
||||
// - url: The URL to open.
|
||||
@@ -21,16 +55,22 @@ import (
|
||||
// Returns:
|
||||
// - An error if the URL cannot be opened, otherwise nil.
|
||||
func OpenURL(url string) error {
|
||||
fmt.Printf("Attempting to open URL in browser: %s\n", url)
|
||||
log.Debugf("Opening URL in browser: %s (incognito=%v)", url, incognitoMode)
|
||||
|
||||
// Try using the open-golang library first
|
||||
err := open.Run(url)
|
||||
// If incognito mode is enabled, use platform-specific incognito commands
|
||||
if incognitoMode {
|
||||
log.Debug("Using incognito mode")
|
||||
return openURLIncognito(url)
|
||||
}
|
||||
|
||||
// Use pkg/browser for cross-platform support
|
||||
err := pkgbrowser.OpenURL(url)
|
||||
if err == nil {
|
||||
log.Debug("Successfully opened URL using open-golang library")
|
||||
log.Debug("Successfully opened URL using pkg/browser library")
|
||||
return nil
|
||||
}
|
||||
|
||||
log.Debugf("open-golang failed: %v, trying platform-specific commands", err)
|
||||
log.Debugf("pkg/browser failed: %v, trying platform-specific commands", err)
|
||||
|
||||
// Fallback to platform-specific commands
|
||||
return openURLPlatformSpecific(url)
|
||||
@@ -78,18 +118,394 @@ func openURLPlatformSpecific(url string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// openURLIncognito opens a URL in incognito/private browsing mode.
|
||||
// It first tries to detect the default browser and use its incognito flag.
|
||||
// Falls back to a chain of known browsers if detection fails.
|
||||
//
|
||||
// Parameters:
|
||||
// - url: The URL to open.
|
||||
//
|
||||
// Returns:
|
||||
// - An error if the URL cannot be opened, otherwise nil.
|
||||
func openURLIncognito(url string) error {
|
||||
// First, try to detect and use the default browser
|
||||
if cmd := tryDefaultBrowserIncognito(url); cmd != nil {
|
||||
log.Debugf("Using detected default browser: %s %v", cmd.Path, cmd.Args[1:])
|
||||
if err := cmd.Start(); err == nil {
|
||||
storeBrowserProcess(cmd)
|
||||
log.Debug("Successfully opened URL in default browser's incognito mode")
|
||||
return nil
|
||||
}
|
||||
log.Debugf("Failed to start default browser, trying fallback chain")
|
||||
}
|
||||
|
||||
// Fallback to known browser chain
|
||||
cmd := tryFallbackBrowsersIncognito(url)
|
||||
if cmd == nil {
|
||||
log.Warn("No browser with incognito support found, falling back to normal mode")
|
||||
return openURLPlatformSpecific(url)
|
||||
}
|
||||
|
||||
log.Debugf("Running incognito command: %s %v", cmd.Path, cmd.Args[1:])
|
||||
err := cmd.Start()
|
||||
if err != nil {
|
||||
log.Warnf("Failed to open incognito browser: %v, falling back to normal mode", err)
|
||||
return openURLPlatformSpecific(url)
|
||||
}
|
||||
|
||||
storeBrowserProcess(cmd)
|
||||
log.Debug("Successfully opened URL in incognito/private mode")
|
||||
return nil
|
||||
}
|
||||
|
||||
// storeBrowserProcess safely stores the browser process for later cleanup.
|
||||
func storeBrowserProcess(cmd *exec.Cmd) {
|
||||
browserMutex.Lock()
|
||||
lastBrowserProcess = cmd
|
||||
browserMutex.Unlock()
|
||||
}
|
||||
|
||||
// tryDefaultBrowserIncognito attempts to detect the default browser and return
|
||||
// an exec.Cmd configured with the appropriate incognito flag.
|
||||
func tryDefaultBrowserIncognito(url string) *exec.Cmd {
|
||||
switch runtime.GOOS {
|
||||
case "darwin":
|
||||
return tryDefaultBrowserMacOS(url)
|
||||
case "windows":
|
||||
return tryDefaultBrowserWindows(url)
|
||||
case "linux":
|
||||
return tryDefaultBrowserLinux(url)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// tryDefaultBrowserMacOS detects the default browser on macOS.
|
||||
func tryDefaultBrowserMacOS(url string) *exec.Cmd {
|
||||
// Try to get default browser from Launch Services
|
||||
out, err := exec.Command("defaults", "read", "com.apple.LaunchServices/com.apple.launchservices.secure", "LSHandlers").Output()
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
output := string(out)
|
||||
var browserName string
|
||||
|
||||
// Parse the output to find the http/https handler
|
||||
if containsBrowserID(output, "com.google.chrome") {
|
||||
browserName = "chrome"
|
||||
} else if containsBrowserID(output, "org.mozilla.firefox") {
|
||||
browserName = "firefox"
|
||||
} else if containsBrowserID(output, "com.apple.safari") {
|
||||
browserName = "safari"
|
||||
} else if containsBrowserID(output, "com.brave.browser") {
|
||||
browserName = "brave"
|
||||
} else if containsBrowserID(output, "com.microsoft.edgemac") {
|
||||
browserName = "edge"
|
||||
}
|
||||
|
||||
return createMacOSIncognitoCmd(browserName, url)
|
||||
}
|
||||
|
||||
// containsBrowserID checks if the LaunchServices output contains a browser ID.
|
||||
func containsBrowserID(output, bundleID string) bool {
|
||||
return stringContains(output, bundleID)
|
||||
}
|
||||
|
||||
// stringContains is a simple string contains check.
|
||||
func stringContains(s, substr string) bool {
|
||||
return len(s) >= len(substr) && (s == substr || len(substr) == 0 ||
|
||||
(len(s) > 0 && len(substr) > 0 && findSubstring(s, substr)))
|
||||
}
|
||||
|
||||
func findSubstring(s, substr string) bool {
|
||||
for i := 0; i <= len(s)-len(substr); i++ {
|
||||
if s[i:i+len(substr)] == substr {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// createMacOSIncognitoCmd creates the appropriate incognito command for macOS browsers.
|
||||
func createMacOSIncognitoCmd(browserName, url string) *exec.Cmd {
|
||||
switch browserName {
|
||||
case "chrome":
|
||||
// Try direct path first
|
||||
chromePath := "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome"
|
||||
if _, err := exec.LookPath(chromePath); err == nil {
|
||||
return exec.Command(chromePath, "--incognito", url)
|
||||
}
|
||||
return exec.Command("open", "-na", "Google Chrome", "--args", "--incognito", url)
|
||||
case "firefox":
|
||||
return exec.Command("open", "-na", "Firefox", "--args", "--private-window", url)
|
||||
case "safari":
|
||||
// Safari doesn't have CLI incognito, try AppleScript
|
||||
return tryAppleScriptSafariPrivate(url)
|
||||
case "brave":
|
||||
return exec.Command("open", "-na", "Brave Browser", "--args", "--incognito", url)
|
||||
case "edge":
|
||||
return exec.Command("open", "-na", "Microsoft Edge", "--args", "--inprivate", url)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// tryAppleScriptSafariPrivate attempts to open Safari in private browsing mode using AppleScript.
|
||||
func tryAppleScriptSafariPrivate(url string) *exec.Cmd {
|
||||
// AppleScript to open a new private window in Safari
|
||||
script := fmt.Sprintf(`
|
||||
tell application "Safari"
|
||||
activate
|
||||
tell application "System Events"
|
||||
keystroke "n" using {command down, shift down}
|
||||
delay 0.5
|
||||
end tell
|
||||
set URL of document 1 to "%s"
|
||||
end tell
|
||||
`, url)
|
||||
|
||||
cmd := exec.Command("osascript", "-e", script)
|
||||
// Test if this approach works by checking if Safari is available
|
||||
if _, err := exec.LookPath("/Applications/Safari.app/Contents/MacOS/Safari"); err != nil {
|
||||
log.Debug("Safari not found, AppleScript private window not available")
|
||||
return nil
|
||||
}
|
||||
log.Debug("Attempting Safari private window via AppleScript")
|
||||
return cmd
|
||||
}
|
||||
|
||||
// tryDefaultBrowserWindows detects the default browser on Windows via registry.
|
||||
func tryDefaultBrowserWindows(url string) *exec.Cmd {
|
||||
// Query registry for default browser
|
||||
out, err := exec.Command("reg", "query",
|
||||
`HKEY_CURRENT_USER\Software\Microsoft\Windows\Shell\Associations\UrlAssociations\http\UserChoice`,
|
||||
"/v", "ProgId").Output()
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
output := string(out)
|
||||
var browserName string
|
||||
|
||||
// Map ProgId to browser name
|
||||
if stringContains(output, "ChromeHTML") {
|
||||
browserName = "chrome"
|
||||
} else if stringContains(output, "FirefoxURL") {
|
||||
browserName = "firefox"
|
||||
} else if stringContains(output, "MSEdgeHTM") {
|
||||
browserName = "edge"
|
||||
} else if stringContains(output, "BraveHTML") {
|
||||
browserName = "brave"
|
||||
}
|
||||
|
||||
return createWindowsIncognitoCmd(browserName, url)
|
||||
}
|
||||
|
||||
// createWindowsIncognitoCmd creates the appropriate incognito command for Windows browsers.
|
||||
func createWindowsIncognitoCmd(browserName, url string) *exec.Cmd {
|
||||
switch browserName {
|
||||
case "chrome":
|
||||
paths := []string{
|
||||
"chrome",
|
||||
`C:\Program Files\Google\Chrome\Application\chrome.exe`,
|
||||
`C:\Program Files (x86)\Google\Chrome\Application\chrome.exe`,
|
||||
}
|
||||
for _, p := range paths {
|
||||
if _, err := exec.LookPath(p); err == nil {
|
||||
return exec.Command(p, "--incognito", url)
|
||||
}
|
||||
}
|
||||
case "firefox":
|
||||
if path, err := exec.LookPath("firefox"); err == nil {
|
||||
return exec.Command(path, "--private-window", url)
|
||||
}
|
||||
case "edge":
|
||||
paths := []string{
|
||||
"msedge",
|
||||
`C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe`,
|
||||
`C:\Program Files\Microsoft\Edge\Application\msedge.exe`,
|
||||
}
|
||||
for _, p := range paths {
|
||||
if _, err := exec.LookPath(p); err == nil {
|
||||
return exec.Command(p, "--inprivate", url)
|
||||
}
|
||||
}
|
||||
case "brave":
|
||||
paths := []string{
|
||||
`C:\Program Files\BraveSoftware\Brave-Browser\Application\brave.exe`,
|
||||
`C:\Program Files (x86)\BraveSoftware\Brave-Browser\Application\brave.exe`,
|
||||
}
|
||||
for _, p := range paths {
|
||||
if _, err := exec.LookPath(p); err == nil {
|
||||
return exec.Command(p, "--incognito", url)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// tryDefaultBrowserLinux detects the default browser on Linux using xdg-settings.
|
||||
func tryDefaultBrowserLinux(url string) *exec.Cmd {
|
||||
out, err := exec.Command("xdg-settings", "get", "default-web-browser").Output()
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
desktop := string(out)
|
||||
var browserName string
|
||||
|
||||
// Map .desktop file to browser name
|
||||
if stringContains(desktop, "google-chrome") || stringContains(desktop, "chrome") {
|
||||
browserName = "chrome"
|
||||
} else if stringContains(desktop, "firefox") {
|
||||
browserName = "firefox"
|
||||
} else if stringContains(desktop, "chromium") {
|
||||
browserName = "chromium"
|
||||
} else if stringContains(desktop, "brave") {
|
||||
browserName = "brave"
|
||||
} else if stringContains(desktop, "microsoft-edge") || stringContains(desktop, "msedge") {
|
||||
browserName = "edge"
|
||||
}
|
||||
|
||||
return createLinuxIncognitoCmd(browserName, url)
|
||||
}
|
||||
|
||||
// createLinuxIncognitoCmd creates the appropriate incognito command for Linux browsers.
|
||||
func createLinuxIncognitoCmd(browserName, url string) *exec.Cmd {
|
||||
switch browserName {
|
||||
case "chrome":
|
||||
paths := []string{"google-chrome", "google-chrome-stable"}
|
||||
for _, p := range paths {
|
||||
if path, err := exec.LookPath(p); err == nil {
|
||||
return exec.Command(path, "--incognito", url)
|
||||
}
|
||||
}
|
||||
case "firefox":
|
||||
paths := []string{"firefox", "firefox-esr"}
|
||||
for _, p := range paths {
|
||||
if path, err := exec.LookPath(p); err == nil {
|
||||
return exec.Command(path, "--private-window", url)
|
||||
}
|
||||
}
|
||||
case "chromium":
|
||||
paths := []string{"chromium", "chromium-browser"}
|
||||
for _, p := range paths {
|
||||
if path, err := exec.LookPath(p); err == nil {
|
||||
return exec.Command(path, "--incognito", url)
|
||||
}
|
||||
}
|
||||
case "brave":
|
||||
if path, err := exec.LookPath("brave-browser"); err == nil {
|
||||
return exec.Command(path, "--incognito", url)
|
||||
}
|
||||
case "edge":
|
||||
if path, err := exec.LookPath("microsoft-edge"); err == nil {
|
||||
return exec.Command(path, "--inprivate", url)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// tryFallbackBrowsersIncognito tries a chain of known browsers as fallback.
|
||||
func tryFallbackBrowsersIncognito(url string) *exec.Cmd {
|
||||
switch runtime.GOOS {
|
||||
case "darwin":
|
||||
return tryFallbackBrowsersMacOS(url)
|
||||
case "windows":
|
||||
return tryFallbackBrowsersWindows(url)
|
||||
case "linux":
|
||||
return tryFallbackBrowsersLinuxChain(url)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// tryFallbackBrowsersMacOS tries known browsers on macOS.
|
||||
func tryFallbackBrowsersMacOS(url string) *exec.Cmd {
|
||||
// Try Chrome
|
||||
chromePath := "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome"
|
||||
if _, err := exec.LookPath(chromePath); err == nil {
|
||||
return exec.Command(chromePath, "--incognito", url)
|
||||
}
|
||||
// Try Firefox
|
||||
if _, err := exec.LookPath("/Applications/Firefox.app/Contents/MacOS/firefox"); err == nil {
|
||||
return exec.Command("open", "-na", "Firefox", "--args", "--private-window", url)
|
||||
}
|
||||
// Try Brave
|
||||
if _, err := exec.LookPath("/Applications/Brave Browser.app/Contents/MacOS/Brave Browser"); err == nil {
|
||||
return exec.Command("open", "-na", "Brave Browser", "--args", "--incognito", url)
|
||||
}
|
||||
// Try Edge
|
||||
if _, err := exec.LookPath("/Applications/Microsoft Edge.app/Contents/MacOS/Microsoft Edge"); err == nil {
|
||||
return exec.Command("open", "-na", "Microsoft Edge", "--args", "--inprivate", url)
|
||||
}
|
||||
// Last resort: try Safari with AppleScript
|
||||
if cmd := tryAppleScriptSafariPrivate(url); cmd != nil {
|
||||
log.Info("Using Safari with AppleScript for private browsing (may require accessibility permissions)")
|
||||
return cmd
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// tryFallbackBrowsersWindows tries known browsers on Windows.
|
||||
func tryFallbackBrowsersWindows(url string) *exec.Cmd {
|
||||
// Chrome
|
||||
chromePaths := []string{
|
||||
"chrome",
|
||||
`C:\Program Files\Google\Chrome\Application\chrome.exe`,
|
||||
`C:\Program Files (x86)\Google\Chrome\Application\chrome.exe`,
|
||||
}
|
||||
for _, p := range chromePaths {
|
||||
if _, err := exec.LookPath(p); err == nil {
|
||||
return exec.Command(p, "--incognito", url)
|
||||
}
|
||||
}
|
||||
// Firefox
|
||||
if path, err := exec.LookPath("firefox"); err == nil {
|
||||
return exec.Command(path, "--private-window", url)
|
||||
}
|
||||
// Edge (usually available on Windows 10+)
|
||||
edgePaths := []string{
|
||||
"msedge",
|
||||
`C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe`,
|
||||
`C:\Program Files\Microsoft\Edge\Application\msedge.exe`,
|
||||
}
|
||||
for _, p := range edgePaths {
|
||||
if _, err := exec.LookPath(p); err == nil {
|
||||
return exec.Command(p, "--inprivate", url)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// tryFallbackBrowsersLinuxChain tries known browsers on Linux.
|
||||
func tryFallbackBrowsersLinuxChain(url string) *exec.Cmd {
|
||||
type browserConfig struct {
|
||||
name string
|
||||
flag string
|
||||
}
|
||||
browsers := []browserConfig{
|
||||
{"google-chrome", "--incognito"},
|
||||
{"google-chrome-stable", "--incognito"},
|
||||
{"chromium", "--incognito"},
|
||||
{"chromium-browser", "--incognito"},
|
||||
{"firefox", "--private-window"},
|
||||
{"firefox-esr", "--private-window"},
|
||||
{"brave-browser", "--incognito"},
|
||||
{"microsoft-edge", "--inprivate"},
|
||||
}
|
||||
for _, b := range browsers {
|
||||
if path, err := exec.LookPath(b.name); err == nil {
|
||||
return exec.Command(path, b.flag, url)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// IsAvailable checks if the system has a command available to open a web browser.
|
||||
// It verifies the presence of necessary commands for the current operating system.
|
||||
//
|
||||
// Returns:
|
||||
// - true if a browser can be opened, false otherwise.
|
||||
func IsAvailable() bool {
|
||||
// First check if open-golang can work
|
||||
testErr := open.Run("about:blank")
|
||||
if testErr == nil {
|
||||
return true
|
||||
}
|
||||
|
||||
// Check platform-specific commands
|
||||
switch runtime.GOOS {
|
||||
case "darwin":
|
||||
|
||||
@@ -19,6 +19,7 @@ func newAuthManager() *sdkAuth.Manager {
|
||||
sdkAuth.NewQwenAuthenticator(),
|
||||
sdkAuth.NewIFlowAuthenticator(),
|
||||
sdkAuth.NewAntigravityAuthenticator(),
|
||||
sdkAuth.NewKiroAuthenticator(),
|
||||
)
|
||||
return manager
|
||||
}
|
||||
|
||||
160
internal/cmd/kiro_login.go
Normal file
160
internal/cmd/kiro_login.go
Normal file
@@ -0,0 +1,160 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
// DoKiroLogin triggers the Kiro authentication flow with Google OAuth.
|
||||
// This is the default login method (same as --kiro-google-login).
|
||||
//
|
||||
// Parameters:
|
||||
// - cfg: The application configuration
|
||||
// - options: Login options including Prompt field
|
||||
func DoKiroLogin(cfg *config.Config, options *LoginOptions) {
|
||||
// Use Google login as default
|
||||
DoKiroGoogleLogin(cfg, options)
|
||||
}
|
||||
|
||||
// DoKiroGoogleLogin triggers Kiro authentication with Google OAuth.
|
||||
// This uses a custom protocol handler (kiro://) to receive the callback.
|
||||
//
|
||||
// Parameters:
|
||||
// - cfg: The application configuration
|
||||
// - options: Login options including prompts
|
||||
func DoKiroGoogleLogin(cfg *config.Config, options *LoginOptions) {
|
||||
if options == nil {
|
||||
options = &LoginOptions{}
|
||||
}
|
||||
|
||||
// Note: Kiro defaults to incognito mode for multi-account support.
|
||||
// Users can override with --no-incognito if they want to use existing browser sessions.
|
||||
|
||||
manager := newAuthManager()
|
||||
|
||||
// Use KiroAuthenticator with Google login
|
||||
authenticator := sdkAuth.NewKiroAuthenticator()
|
||||
record, err := authenticator.LoginWithGoogle(context.Background(), cfg, &sdkAuth.LoginOptions{
|
||||
NoBrowser: options.NoBrowser,
|
||||
Metadata: map[string]string{},
|
||||
Prompt: options.Prompt,
|
||||
})
|
||||
if err != nil {
|
||||
log.Errorf("Kiro Google authentication failed: %v", err)
|
||||
fmt.Println("\nTroubleshooting:")
|
||||
fmt.Println("1. Make sure the protocol handler is installed")
|
||||
fmt.Println("2. Complete the Google login in the browser")
|
||||
fmt.Println("3. If callback fails, try: --kiro-import (after logging in via Kiro IDE)")
|
||||
return
|
||||
}
|
||||
|
||||
// Save the auth record
|
||||
savedPath, err := manager.SaveAuth(record, cfg)
|
||||
if err != nil {
|
||||
log.Errorf("Failed to save auth: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
if savedPath != "" {
|
||||
fmt.Printf("Authentication saved to %s\n", savedPath)
|
||||
}
|
||||
if record != nil && record.Label != "" {
|
||||
fmt.Printf("Authenticated as %s\n", record.Label)
|
||||
}
|
||||
fmt.Println("Kiro Google authentication successful!")
|
||||
}
|
||||
|
||||
// DoKiroAWSLogin triggers Kiro authentication with AWS Builder ID.
|
||||
// This uses the device code flow for AWS SSO OIDC authentication.
|
||||
//
|
||||
// Parameters:
|
||||
// - cfg: The application configuration
|
||||
// - options: Login options including prompts
|
||||
func DoKiroAWSLogin(cfg *config.Config, options *LoginOptions) {
|
||||
if options == nil {
|
||||
options = &LoginOptions{}
|
||||
}
|
||||
|
||||
// Note: Kiro defaults to incognito mode for multi-account support.
|
||||
// Users can override with --no-incognito if they want to use existing browser sessions.
|
||||
|
||||
manager := newAuthManager()
|
||||
|
||||
// Use KiroAuthenticator with AWS Builder ID login (device code flow)
|
||||
authenticator := sdkAuth.NewKiroAuthenticator()
|
||||
record, err := authenticator.Login(context.Background(), cfg, &sdkAuth.LoginOptions{
|
||||
NoBrowser: options.NoBrowser,
|
||||
Metadata: map[string]string{},
|
||||
Prompt: options.Prompt,
|
||||
})
|
||||
if err != nil {
|
||||
log.Errorf("Kiro AWS authentication failed: %v", err)
|
||||
fmt.Println("\nTroubleshooting:")
|
||||
fmt.Println("1. Make sure you have an AWS Builder ID")
|
||||
fmt.Println("2. Complete the authorization in the browser")
|
||||
fmt.Println("3. If callback fails, try: --kiro-import (after logging in via Kiro IDE)")
|
||||
return
|
||||
}
|
||||
|
||||
// Save the auth record
|
||||
savedPath, err := manager.SaveAuth(record, cfg)
|
||||
if err != nil {
|
||||
log.Errorf("Failed to save auth: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
if savedPath != "" {
|
||||
fmt.Printf("Authentication saved to %s\n", savedPath)
|
||||
}
|
||||
if record != nil && record.Label != "" {
|
||||
fmt.Printf("Authenticated as %s\n", record.Label)
|
||||
}
|
||||
fmt.Println("Kiro AWS authentication successful!")
|
||||
}
|
||||
|
||||
// DoKiroImport imports Kiro token from Kiro IDE's token file.
|
||||
// This is useful for users who have already logged in via Kiro IDE
|
||||
// and want to use the same credentials in CLI Proxy API.
|
||||
//
|
||||
// Parameters:
|
||||
// - cfg: The application configuration
|
||||
// - options: Login options (currently unused for import)
|
||||
func DoKiroImport(cfg *config.Config, options *LoginOptions) {
|
||||
if options == nil {
|
||||
options = &LoginOptions{}
|
||||
}
|
||||
|
||||
manager := newAuthManager()
|
||||
|
||||
// Use ImportFromKiroIDE instead of Login
|
||||
authenticator := sdkAuth.NewKiroAuthenticator()
|
||||
record, err := authenticator.ImportFromKiroIDE(context.Background(), cfg)
|
||||
if err != nil {
|
||||
log.Errorf("Kiro token import failed: %v", err)
|
||||
fmt.Println("\nMake sure you have logged in to Kiro IDE first:")
|
||||
fmt.Println("1. Open Kiro IDE")
|
||||
fmt.Println("2. Click 'Sign in with Google' (or GitHub)")
|
||||
fmt.Println("3. Complete the login process")
|
||||
fmt.Println("4. Run this command again")
|
||||
return
|
||||
}
|
||||
|
||||
// Save the imported auth record
|
||||
savedPath, err := manager.SaveAuth(record, cfg)
|
||||
if err != nil {
|
||||
log.Errorf("Failed to save auth: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
if savedPath != "" {
|
||||
fmt.Printf("Authentication saved to %s\n", savedPath)
|
||||
}
|
||||
if record != nil && record.Label != "" {
|
||||
fmt.Printf("Imported as %s\n", record.Label)
|
||||
}
|
||||
fmt.Println("Kiro token import successful!")
|
||||
}
|
||||
@@ -58,6 +58,9 @@ type Config struct {
|
||||
// GeminiKey defines Gemini API key configurations with optional routing overrides.
|
||||
GeminiKey []GeminiKey `yaml:"gemini-api-key" json:"gemini-api-key"`
|
||||
|
||||
// KiroKey defines a list of Kiro (AWS CodeWhisperer) configurations.
|
||||
KiroKey []KiroKey `yaml:"kiro" json:"kiro"`
|
||||
|
||||
// Codex defines a list of Codex API key configurations as specified in the YAML configuration file.
|
||||
CodexKey []CodexKey `yaml:"codex-api-key" json:"codex-api-key"`
|
||||
|
||||
@@ -80,6 +83,11 @@ type Config struct {
|
||||
// Payload defines default and override rules for provider payload parameters.
|
||||
Payload PayloadConfig `yaml:"payload" json:"payload"`
|
||||
|
||||
// IncognitoBrowser enables opening OAuth URLs in incognito/private browsing mode.
|
||||
// This is useful when you want to login with a different account without logging out
|
||||
// from your current session. Default: false.
|
||||
IncognitoBrowser bool `yaml:"incognito-browser" json:"incognito-browser"`
|
||||
|
||||
legacyMigrationPending bool `yaml:"-" json:"-"`
|
||||
}
|
||||
|
||||
@@ -240,6 +248,31 @@ type GeminiKey struct {
|
||||
ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"`
|
||||
}
|
||||
|
||||
// KiroKey represents the configuration for Kiro (AWS CodeWhisperer) authentication.
|
||||
type KiroKey struct {
|
||||
// TokenFile is the path to the Kiro token file (default: ~/.aws/sso/cache/kiro-auth-token.json)
|
||||
TokenFile string `yaml:"token-file,omitempty" json:"token-file,omitempty"`
|
||||
|
||||
// AccessToken is the OAuth access token for direct configuration.
|
||||
AccessToken string `yaml:"access-token,omitempty" json:"access-token,omitempty"`
|
||||
|
||||
// RefreshToken is the OAuth refresh token for token renewal.
|
||||
RefreshToken string `yaml:"refresh-token,omitempty" json:"refresh-token,omitempty"`
|
||||
|
||||
// ProfileArn is the AWS CodeWhisperer profile ARN.
|
||||
ProfileArn string `yaml:"profile-arn,omitempty" json:"profile-arn,omitempty"`
|
||||
|
||||
// Region is the AWS region (default: us-east-1).
|
||||
Region string `yaml:"region,omitempty" json:"region,omitempty"`
|
||||
|
||||
// ProxyURL optionally overrides the global proxy for this configuration.
|
||||
ProxyURL string `yaml:"proxy-url,omitempty" json:"proxy-url,omitempty"`
|
||||
|
||||
// AgentTaskType sets the Kiro API task type. Known values: "vibe", "dev", "chat".
|
||||
// Leave empty to let API use defaults. Different values may inject different system prompts.
|
||||
AgentTaskType string `yaml:"agent-task-type,omitempty" json:"agent-task-type,omitempty"`
|
||||
}
|
||||
|
||||
// OpenAICompatibility represents the configuration for OpenAI API compatibility
|
||||
// with external providers, allowing model aliases to be routed through OpenAI API format.
|
||||
type OpenAICompatibility struct {
|
||||
@@ -320,6 +353,7 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
|
||||
cfg.UsageStatisticsEnabled = false
|
||||
cfg.DisableCooling = false
|
||||
cfg.AmpCode.RestrictManagementToLocalhost = true // Default to secure: only localhost access
|
||||
cfg.IncognitoBrowser = false // Default to normal browser (AWS uses incognito by force)
|
||||
if err = yaml.Unmarshal(data, &cfg); err != nil {
|
||||
if optional {
|
||||
// In cloud deploy mode, if YAML parsing fails, return empty config instead of error.
|
||||
@@ -370,6 +404,9 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
|
||||
// Sanitize Claude key headers
|
||||
cfg.SanitizeClaudeKeys()
|
||||
|
||||
// Sanitize Kiro keys: trim whitespace from credential fields
|
||||
cfg.SanitizeKiroKeys()
|
||||
|
||||
// Sanitize OpenAI compatibility providers: drop entries without base-url
|
||||
cfg.SanitizeOpenAICompatibility()
|
||||
|
||||
@@ -446,6 +483,22 @@ func (cfg *Config) SanitizeClaudeKeys() {
|
||||
}
|
||||
}
|
||||
|
||||
// SanitizeKiroKeys trims whitespace from Kiro credential fields.
|
||||
func (cfg *Config) SanitizeKiroKeys() {
|
||||
if cfg == nil || len(cfg.KiroKey) == 0 {
|
||||
return
|
||||
}
|
||||
for i := range cfg.KiroKey {
|
||||
entry := &cfg.KiroKey[i]
|
||||
entry.TokenFile = strings.TrimSpace(entry.TokenFile)
|
||||
entry.AccessToken = strings.TrimSpace(entry.AccessToken)
|
||||
entry.RefreshToken = strings.TrimSpace(entry.RefreshToken)
|
||||
entry.ProfileArn = strings.TrimSpace(entry.ProfileArn)
|
||||
entry.Region = strings.TrimSpace(entry.Region)
|
||||
entry.ProxyURL = strings.TrimSpace(entry.ProxyURL)
|
||||
}
|
||||
}
|
||||
|
||||
// SanitizeGeminiKeys deduplicates and normalizes Gemini credentials.
|
||||
func (cfg *Config) SanitizeGeminiKeys() {
|
||||
if cfg == nil {
|
||||
|
||||
@@ -24,4 +24,7 @@ const (
|
||||
|
||||
// Antigravity represents the Antigravity response format identifier.
|
||||
Antigravity = "antigravity"
|
||||
|
||||
// Kiro represents the AWS CodeWhisperer (Kiro) provider identifier.
|
||||
Kiro = "kiro"
|
||||
)
|
||||
|
||||
@@ -38,13 +38,16 @@ func (m *LogFormatter) Format(entry *log.Entry) ([]byte, error) {
|
||||
|
||||
timestamp := entry.Time.Format("2006-01-02 15:04:05")
|
||||
message := strings.TrimRight(entry.Message, "\r\n")
|
||||
|
||||
var formatted string
|
||||
|
||||
// Handle nil Caller (can happen with some log entries)
|
||||
callerFile := "unknown"
|
||||
callerLine := 0
|
||||
if entry.Caller != nil {
|
||||
formatted = fmt.Sprintf("[%s] [%s] [%s:%d] %s\n", timestamp, entry.Level, filepath.Base(entry.Caller.File), entry.Caller.Line, message)
|
||||
} else {
|
||||
formatted = fmt.Sprintf("[%s] [%s] %s\n", timestamp, entry.Level, message)
|
||||
callerFile = filepath.Base(entry.Caller.File)
|
||||
callerLine = entry.Caller.Line
|
||||
}
|
||||
|
||||
formatted := fmt.Sprintf("[%s] [%s] [%s:%d] %s\n", timestamp, entry.Level, callerFile, callerLine, message)
|
||||
buffer.WriteString(formatted)
|
||||
|
||||
return buffer.Bytes(), nil
|
||||
@@ -55,6 +58,7 @@ func (m *LogFormatter) Format(entry *log.Entry) ([]byte, error) {
|
||||
func SetupBaseLogger() {
|
||||
setupOnce.Do(func() {
|
||||
log.SetOutput(os.Stdout)
|
||||
log.SetLevel(log.InfoLevel)
|
||||
log.SetReportCaller(true)
|
||||
log.SetFormatter(&LogFormatter{})
|
||||
|
||||
|
||||
@@ -986,3 +986,161 @@ func GetIFlowModels() []*ModelInfo {
|
||||
}
|
||||
return models
|
||||
}
|
||||
|
||||
// GetKiroModels returns the Kiro (AWS CodeWhisperer) model definitions
|
||||
func GetKiroModels() []*ModelInfo {
|
||||
return []*ModelInfo{
|
||||
{
|
||||
ID: "kiro-auto",
|
||||
Object: "model",
|
||||
Created: 1732752000, // 2024-11-28
|
||||
OwnedBy: "aws",
|
||||
Type: "kiro",
|
||||
DisplayName: "Kiro Auto",
|
||||
Description: "Automatic model selection by AWS CodeWhisperer",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
},
|
||||
{
|
||||
ID: "kiro-claude-opus-4.5",
|
||||
Object: "model",
|
||||
Created: 1732752000,
|
||||
OwnedBy: "aws",
|
||||
Type: "kiro",
|
||||
DisplayName: "Kiro Claude Opus 4.5",
|
||||
Description: "Claude Opus 4.5 via Kiro (2.2x credit)",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
},
|
||||
{
|
||||
ID: "kiro-claude-sonnet-4.5",
|
||||
Object: "model",
|
||||
Created: 1732752000,
|
||||
OwnedBy: "aws",
|
||||
Type: "kiro",
|
||||
DisplayName: "Kiro Claude Sonnet 4.5",
|
||||
Description: "Claude Sonnet 4.5 via Kiro (1.3x credit)",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
},
|
||||
{
|
||||
ID: "kiro-claude-sonnet-4",
|
||||
Object: "model",
|
||||
Created: 1732752000,
|
||||
OwnedBy: "aws",
|
||||
Type: "kiro",
|
||||
DisplayName: "Kiro Claude Sonnet 4",
|
||||
Description: "Claude Sonnet 4 via Kiro (1.3x credit)",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
},
|
||||
{
|
||||
ID: "kiro-claude-haiku-4.5",
|
||||
Object: "model",
|
||||
Created: 1732752000,
|
||||
OwnedBy: "aws",
|
||||
Type: "kiro",
|
||||
DisplayName: "Kiro Claude Haiku 4.5",
|
||||
Description: "Claude Haiku 4.5 via Kiro (0.4x credit)",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
},
|
||||
// --- Chat Variant (No tool calling, for pure conversation) ---
|
||||
{
|
||||
ID: "kiro-claude-opus-4.5-chat",
|
||||
Object: "model",
|
||||
Created: 1732752000,
|
||||
OwnedBy: "aws",
|
||||
Type: "kiro",
|
||||
DisplayName: "Kiro Claude Opus 4.5 (Chat)",
|
||||
Description: "Claude Opus 4.5 for chat only (no tool calling)",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
},
|
||||
// --- Agentic Variants (Optimized for coding agents with chunked writes) ---
|
||||
{
|
||||
ID: "kiro-claude-opus-4.5-agentic",
|
||||
Object: "model",
|
||||
Created: 1732752000,
|
||||
OwnedBy: "aws",
|
||||
Type: "kiro",
|
||||
DisplayName: "Kiro Claude Opus 4.5 (Agentic)",
|
||||
Description: "Claude Opus 4.5 optimized for coding agents (chunked writes)",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
},
|
||||
{
|
||||
ID: "kiro-claude-sonnet-4.5-agentic",
|
||||
Object: "model",
|
||||
Created: 1732752000,
|
||||
OwnedBy: "aws",
|
||||
Type: "kiro",
|
||||
DisplayName: "Kiro Claude Sonnet 4.5 (Agentic)",
|
||||
Description: "Claude Sonnet 4.5 optimized for coding agents (chunked writes)",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// GetAmazonQModels returns the Amazon Q (AWS CodeWhisperer) model definitions.
|
||||
// These models use the same API as Kiro and share the same executor.
|
||||
func GetAmazonQModels() []*ModelInfo {
|
||||
return []*ModelInfo{
|
||||
{
|
||||
ID: "amazonq-auto",
|
||||
Object: "model",
|
||||
Created: 1732752000,
|
||||
OwnedBy: "aws",
|
||||
Type: "kiro", // Uses Kiro executor - same API
|
||||
DisplayName: "Amazon Q Auto",
|
||||
Description: "Automatic model selection by Amazon Q",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
},
|
||||
{
|
||||
ID: "amazonq-claude-opus-4.5",
|
||||
Object: "model",
|
||||
Created: 1732752000,
|
||||
OwnedBy: "aws",
|
||||
Type: "kiro",
|
||||
DisplayName: "Amazon Q Claude Opus 4.5",
|
||||
Description: "Claude Opus 4.5 via Amazon Q (2.2x credit)",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
},
|
||||
{
|
||||
ID: "amazonq-claude-sonnet-4.5",
|
||||
Object: "model",
|
||||
Created: 1732752000,
|
||||
OwnedBy: "aws",
|
||||
Type: "kiro",
|
||||
DisplayName: "Amazon Q Claude Sonnet 4.5",
|
||||
Description: "Claude Sonnet 4.5 via Amazon Q (1.3x credit)",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
},
|
||||
{
|
||||
ID: "amazonq-claude-sonnet-4",
|
||||
Object: "model",
|
||||
Created: 1732752000,
|
||||
OwnedBy: "aws",
|
||||
Type: "kiro",
|
||||
DisplayName: "Amazon Q Claude Sonnet 4",
|
||||
Description: "Claude Sonnet 4 via Amazon Q (1.3x credit)",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
},
|
||||
{
|
||||
ID: "amazonq-claude-haiku-4.5",
|
||||
Object: "model",
|
||||
Created: 1732752000,
|
||||
OwnedBy: "aws",
|
||||
Type: "kiro",
|
||||
DisplayName: "Amazon Q Claude Haiku 4.5",
|
||||
Description: "Claude Haiku 4.5 via Amazon Q (0.4x credit)",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -12,6 +12,7 @@ import (
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
@@ -41,7 +42,10 @@ const (
|
||||
streamScannerBuffer int = 20_971_520
|
||||
)
|
||||
|
||||
var randSource = rand.New(rand.NewSource(time.Now().UnixNano()))
|
||||
var (
|
||||
randSource = rand.New(rand.NewSource(time.Now().UnixNano()))
|
||||
randSourceMutex sync.Mutex
|
||||
)
|
||||
|
||||
// AntigravityExecutor proxies requests to the antigravity upstream.
|
||||
type AntigravityExecutor struct {
|
||||
@@ -754,15 +758,19 @@ func generateRequestID() string {
|
||||
}
|
||||
|
||||
func generateSessionID() string {
|
||||
randSourceMutex.Lock()
|
||||
n := randSource.Int63n(9_000_000_000_000_000_000)
|
||||
randSourceMutex.Unlock()
|
||||
return "-" + strconv.FormatInt(n, 10)
|
||||
}
|
||||
|
||||
func generateProjectID() string {
|
||||
adjectives := []string{"useful", "bright", "swift", "calm", "bold"}
|
||||
nouns := []string{"fuze", "wave", "spark", "flow", "core"}
|
||||
randSourceMutex.Lock()
|
||||
adj := adjectives[randSource.Intn(len(adjectives))]
|
||||
noun := nouns[randSource.Intn(len(nouns))]
|
||||
randSourceMutex.Unlock()
|
||||
randomPart := strings.ToLower(uuid.NewString())[:5]
|
||||
return adj + "-" + noun + "-" + randomPart
|
||||
}
|
||||
|
||||
@@ -1,10 +1,38 @@
|
||||
package executor
|
||||
|
||||
import "time"
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
type codexCache struct {
|
||||
ID string
|
||||
Expire time.Time
|
||||
}
|
||||
|
||||
var codexCacheMap = map[string]codexCache{}
|
||||
var (
|
||||
codexCacheMap = map[string]codexCache{}
|
||||
codexCacheMutex sync.RWMutex
|
||||
)
|
||||
|
||||
// getCodexCache safely retrieves a cache entry
|
||||
func getCodexCache(key string) (codexCache, bool) {
|
||||
codexCacheMutex.RLock()
|
||||
defer codexCacheMutex.RUnlock()
|
||||
cache, ok := codexCacheMap[key]
|
||||
return cache, ok
|
||||
}
|
||||
|
||||
// setCodexCache safely sets a cache entry
|
||||
func setCodexCache(key string, cache codexCache) {
|
||||
codexCacheMutex.Lock()
|
||||
defer codexCacheMutex.Unlock()
|
||||
codexCacheMap[key] = cache
|
||||
}
|
||||
|
||||
// deleteCodexCache safely deletes a cache entry
|
||||
func deleteCodexCache(key string) {
|
||||
codexCacheMutex.Lock()
|
||||
defer codexCacheMutex.Unlock()
|
||||
delete(codexCacheMap, key)
|
||||
}
|
||||
|
||||
@@ -506,12 +506,12 @@ func (e *CodexExecutor) cacheHelper(ctx context.Context, from sdktranslator.Form
|
||||
if userIDResult.Exists() {
|
||||
var hasKey bool
|
||||
key := fmt.Sprintf("%s-%s", req.Model, userIDResult.String())
|
||||
if cache, hasKey = codexCacheMap[key]; !hasKey || cache.Expire.Before(time.Now()) {
|
||||
if cache, hasKey = getCodexCache(key); !hasKey || cache.Expire.Before(time.Now()) {
|
||||
cache = codexCache{
|
||||
ID: uuid.New().String(),
|
||||
Expire: time.Now().Add(1 * time.Hour),
|
||||
}
|
||||
codexCacheMap[key] = cache
|
||||
setCodexCache(key, cache)
|
||||
}
|
||||
}
|
||||
} else if from == "openai-response" {
|
||||
|
||||
2353
internal/runtime/executor/kiro_executor.go
Normal file
2353
internal/runtime/executor/kiro_executor.go
Normal file
File diff suppressed because it is too large
Load Diff
@@ -6,6 +6,7 @@ import (
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
@@ -14,11 +15,19 @@ import (
|
||||
"golang.org/x/net/proxy"
|
||||
)
|
||||
|
||||
// httpClientCache caches HTTP clients by proxy URL to enable connection reuse
|
||||
var (
|
||||
httpClientCache = make(map[string]*http.Client)
|
||||
httpClientCacheMutex sync.RWMutex
|
||||
)
|
||||
|
||||
// newProxyAwareHTTPClient creates an HTTP client with proper proxy configuration priority:
|
||||
// 1. Use auth.ProxyURL if configured (highest priority)
|
||||
// 2. Use cfg.ProxyURL if auth proxy is not configured
|
||||
// 3. Use RoundTripper from context if neither are configured
|
||||
//
|
||||
// This function caches HTTP clients by proxy URL to enable TCP/TLS connection reuse.
|
||||
//
|
||||
// Parameters:
|
||||
// - ctx: The context containing optional RoundTripper
|
||||
// - cfg: The application configuration
|
||||
@@ -28,11 +37,6 @@ import (
|
||||
// Returns:
|
||||
// - *http.Client: An HTTP client with configured proxy or transport
|
||||
func newProxyAwareHTTPClient(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, timeout time.Duration) *http.Client {
|
||||
httpClient := &http.Client{}
|
||||
if timeout > 0 {
|
||||
httpClient.Timeout = timeout
|
||||
}
|
||||
|
||||
// Priority 1: Use auth.ProxyURL if configured
|
||||
var proxyURL string
|
||||
if auth != nil {
|
||||
@@ -44,11 +48,39 @@ func newProxyAwareHTTPClient(ctx context.Context, cfg *config.Config, auth *clip
|
||||
proxyURL = strings.TrimSpace(cfg.ProxyURL)
|
||||
}
|
||||
|
||||
// Build cache key from proxy URL (empty string for no proxy)
|
||||
cacheKey := proxyURL
|
||||
|
||||
// Check cache first
|
||||
httpClientCacheMutex.RLock()
|
||||
if cachedClient, ok := httpClientCache[cacheKey]; ok {
|
||||
httpClientCacheMutex.RUnlock()
|
||||
// Return a wrapper with the requested timeout but shared transport
|
||||
if timeout > 0 {
|
||||
return &http.Client{
|
||||
Transport: cachedClient.Transport,
|
||||
Timeout: timeout,
|
||||
}
|
||||
}
|
||||
return cachedClient
|
||||
}
|
||||
httpClientCacheMutex.RUnlock()
|
||||
|
||||
// Create new client
|
||||
httpClient := &http.Client{}
|
||||
if timeout > 0 {
|
||||
httpClient.Timeout = timeout
|
||||
}
|
||||
|
||||
// If we have a proxy URL configured, set up the transport
|
||||
if proxyURL != "" {
|
||||
transport := buildProxyTransport(proxyURL)
|
||||
if transport != nil {
|
||||
httpClient.Transport = transport
|
||||
// Cache the client
|
||||
httpClientCacheMutex.Lock()
|
||||
httpClientCache[cacheKey] = httpClient
|
||||
httpClientCacheMutex.Unlock()
|
||||
return httpClient
|
||||
}
|
||||
// If proxy setup failed, log and fall through to context RoundTripper
|
||||
@@ -60,6 +92,13 @@ func newProxyAwareHTTPClient(ctx context.Context, cfg *config.Config, auth *clip
|
||||
httpClient.Transport = rt
|
||||
}
|
||||
|
||||
// Cache the client for no-proxy case
|
||||
if proxyURL == "" {
|
||||
httpClientCacheMutex.Lock()
|
||||
httpClientCache[cacheKey] = httpClient
|
||||
httpClientCacheMutex.Unlock()
|
||||
}
|
||||
|
||||
return httpClient
|
||||
}
|
||||
|
||||
|
||||
@@ -331,8 +331,9 @@ func ConvertClaudeResponseToGeminiNonStream(_ context.Context, modelName string,
|
||||
streamingEvents := make([][]byte, 0)
|
||||
|
||||
scanner := bufio.NewScanner(bytes.NewReader(rawJSON))
|
||||
buffer := make([]byte, 20_971_520)
|
||||
scanner.Buffer(buffer, 20_971_520)
|
||||
// Use a smaller initial buffer (64KB) that can grow up to 20MB if needed
|
||||
// This prevents allocating 20MB for every request regardless of size
|
||||
scanner.Buffer(make([]byte, 64*1024), 20_971_520)
|
||||
for scanner.Scan() {
|
||||
line := scanner.Bytes()
|
||||
// log.Debug(string(line))
|
||||
|
||||
@@ -50,6 +50,10 @@ type ToolCallAccumulator struct {
|
||||
// Returns:
|
||||
// - []string: A slice of strings, each containing an OpenAI-compatible JSON response
|
||||
func ConvertClaudeResponseToOpenAI(_ context.Context, modelName string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) []string {
|
||||
var localParam any
|
||||
if param == nil {
|
||||
param = &localParam
|
||||
}
|
||||
if *param == nil {
|
||||
*param = &ConvertAnthropicResponseToOpenAIParams{
|
||||
CreatedAt: 0,
|
||||
|
||||
@@ -60,12 +60,12 @@ func ConvertGeminiCLIResponseToClaude(_ context.Context, _ string, originalReque
|
||||
|
||||
// Track whether tools are being used in this response chunk
|
||||
usedTool := false
|
||||
output := ""
|
||||
var sb strings.Builder
|
||||
|
||||
// Initialize the streaming session with a message_start event
|
||||
// This is only sent for the very first response chunk to establish the streaming session
|
||||
if !(*param).(*Params).HasFirstResponse {
|
||||
output = "event: message_start\n"
|
||||
sb.WriteString("event: message_start\n")
|
||||
|
||||
// Create the initial message structure with default values according to Claude Code API specification
|
||||
// This follows the Claude Code API specification for streaming message initialization
|
||||
@@ -78,7 +78,7 @@ func ConvertGeminiCLIResponseToClaude(_ context.Context, _ string, originalReque
|
||||
if responseIDResult := gjson.GetBytes(rawJSON, "response.responseId"); responseIDResult.Exists() {
|
||||
messageStartTemplate, _ = sjson.Set(messageStartTemplate, "message.id", responseIDResult.String())
|
||||
}
|
||||
output = output + fmt.Sprintf("data: %s\n\n\n", messageStartTemplate)
|
||||
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", messageStartTemplate))
|
||||
|
||||
(*param).(*Params).HasFirstResponse = true
|
||||
}
|
||||
@@ -101,62 +101,52 @@ func ConvertGeminiCLIResponseToClaude(_ context.Context, _ string, originalReque
|
||||
if partResult.Get("thought").Bool() {
|
||||
// Continue existing thinking block if already in thinking state
|
||||
if (*param).(*Params).ResponseType == 2 {
|
||||
output = output + "event: content_block_delta\n"
|
||||
sb.WriteString("event: content_block_delta\n")
|
||||
data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"thinking_delta","thinking":""}}`, (*param).(*Params).ResponseIndex), "delta.thinking", partTextResult.String())
|
||||
output = output + fmt.Sprintf("data: %s\n\n\n", data)
|
||||
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", data))
|
||||
} else {
|
||||
// Transition from another state to thinking
|
||||
// First, close any existing content block
|
||||
if (*param).(*Params).ResponseType != 0 {
|
||||
if (*param).(*Params).ResponseType == 2 {
|
||||
// output = output + "event: content_block_delta\n"
|
||||
// output = output + fmt.Sprintf(`data: {"type":"content_block_delta","index":%d,"delta":{"type":"signature_delta","signature":null}}`, (*param).(*Params).ResponseIndex)
|
||||
// output = output + "\n\n\n"
|
||||
}
|
||||
output = output + "event: content_block_stop\n"
|
||||
output = output + fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex)
|
||||
output = output + "\n\n\n"
|
||||
sb.WriteString("event: content_block_stop\n")
|
||||
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex))
|
||||
sb.WriteString("\n\n\n")
|
||||
(*param).(*Params).ResponseIndex++
|
||||
}
|
||||
|
||||
// Start a new thinking content block
|
||||
output = output + "event: content_block_start\n"
|
||||
output = output + fmt.Sprintf(`data: {"type":"content_block_start","index":%d,"content_block":{"type":"thinking","thinking":""}}`, (*param).(*Params).ResponseIndex)
|
||||
output = output + "\n\n\n"
|
||||
output = output + "event: content_block_delta\n"
|
||||
sb.WriteString("event: content_block_start\n")
|
||||
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_start","index":%d,"content_block":{"type":"thinking","thinking":""}}`, (*param).(*Params).ResponseIndex))
|
||||
sb.WriteString("\n\n\n")
|
||||
sb.WriteString("event: content_block_delta\n")
|
||||
data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"thinking_delta","thinking":""}}`, (*param).(*Params).ResponseIndex), "delta.thinking", partTextResult.String())
|
||||
output = output + fmt.Sprintf("data: %s\n\n\n", data)
|
||||
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", data))
|
||||
(*param).(*Params).ResponseType = 2 // Set state to thinking
|
||||
}
|
||||
} else {
|
||||
// Process regular text content (user-visible output)
|
||||
// Continue existing text block if already in content state
|
||||
if (*param).(*Params).ResponseType == 1 {
|
||||
output = output + "event: content_block_delta\n"
|
||||
sb.WriteString("event: content_block_delta\n")
|
||||
data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"text_delta","text":""}}`, (*param).(*Params).ResponseIndex), "delta.text", partTextResult.String())
|
||||
output = output + fmt.Sprintf("data: %s\n\n\n", data)
|
||||
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", data))
|
||||
} else {
|
||||
// Transition from another state to text content
|
||||
// First, close any existing content block
|
||||
if (*param).(*Params).ResponseType != 0 {
|
||||
if (*param).(*Params).ResponseType == 2 {
|
||||
// output = output + "event: content_block_delta\n"
|
||||
// output = output + fmt.Sprintf(`data: {"type":"content_block_delta","index":%d,"delta":{"type":"signature_delta","signature":null}}`, (*param).(*Params).ResponseIndex)
|
||||
// output = output + "\n\n\n"
|
||||
}
|
||||
output = output + "event: content_block_stop\n"
|
||||
output = output + fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex)
|
||||
output = output + "\n\n\n"
|
||||
sb.WriteString("event: content_block_stop\n")
|
||||
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex))
|
||||
sb.WriteString("\n\n\n")
|
||||
(*param).(*Params).ResponseIndex++
|
||||
}
|
||||
|
||||
// Start a new text content block
|
||||
output = output + "event: content_block_start\n"
|
||||
output = output + fmt.Sprintf(`data: {"type":"content_block_start","index":%d,"content_block":{"type":"text","text":""}}`, (*param).(*Params).ResponseIndex)
|
||||
output = output + "\n\n\n"
|
||||
output = output + "event: content_block_delta\n"
|
||||
sb.WriteString("event: content_block_start\n")
|
||||
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_start","index":%d,"content_block":{"type":"text","text":""}}`, (*param).(*Params).ResponseIndex))
|
||||
sb.WriteString("\n\n\n")
|
||||
sb.WriteString("event: content_block_delta\n")
|
||||
data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"text_delta","text":""}}`, (*param).(*Params).ResponseIndex), "delta.text", partTextResult.String())
|
||||
output = output + fmt.Sprintf("data: %s\n\n\n", data)
|
||||
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", data))
|
||||
(*param).(*Params).ResponseType = 1 // Set state to content
|
||||
}
|
||||
}
|
||||
@@ -169,42 +159,35 @@ func ConvertGeminiCLIResponseToClaude(_ context.Context, _ string, originalReque
|
||||
// Handle state transitions when switching to function calls
|
||||
// Close any existing function call block first
|
||||
if (*param).(*Params).ResponseType == 3 {
|
||||
output = output + "event: content_block_stop\n"
|
||||
output = output + fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex)
|
||||
output = output + "\n\n\n"
|
||||
sb.WriteString("event: content_block_stop\n")
|
||||
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex))
|
||||
sb.WriteString("\n\n\n")
|
||||
(*param).(*Params).ResponseIndex++
|
||||
(*param).(*Params).ResponseType = 0
|
||||
}
|
||||
|
||||
// Special handling for thinking state transition
|
||||
if (*param).(*Params).ResponseType == 2 {
|
||||
// output = output + "event: content_block_delta\n"
|
||||
// output = output + fmt.Sprintf(`data: {"type":"content_block_delta","index":%d,"delta":{"type":"signature_delta","signature":null}}`, (*param).(*Params).ResponseIndex)
|
||||
// output = output + "\n\n\n"
|
||||
}
|
||||
|
||||
// Close any other existing content block
|
||||
if (*param).(*Params).ResponseType != 0 {
|
||||
output = output + "event: content_block_stop\n"
|
||||
output = output + fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex)
|
||||
output = output + "\n\n\n"
|
||||
sb.WriteString("event: content_block_stop\n")
|
||||
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex))
|
||||
sb.WriteString("\n\n\n")
|
||||
(*param).(*Params).ResponseIndex++
|
||||
}
|
||||
|
||||
// Start a new tool use content block
|
||||
// This creates the structure for a function call in Claude Code format
|
||||
output = output + "event: content_block_start\n"
|
||||
sb.WriteString("event: content_block_start\n")
|
||||
|
||||
// Create the tool use block with unique ID and function details
|
||||
data := fmt.Sprintf(`{"type":"content_block_start","index":%d,"content_block":{"type":"tool_use","id":"","name":"","input":{}}}`, (*param).(*Params).ResponseIndex)
|
||||
data, _ = sjson.Set(data, "content_block.id", fmt.Sprintf("%s-%d", fcName, time.Now().UnixNano()))
|
||||
data, _ = sjson.Set(data, "content_block.name", fcName)
|
||||
output = output + fmt.Sprintf("data: %s\n\n\n", data)
|
||||
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", data))
|
||||
|
||||
if fcArgsResult := functionCallResult.Get("args"); fcArgsResult.Exists() {
|
||||
output = output + "event: content_block_delta\n"
|
||||
sb.WriteString("event: content_block_delta\n")
|
||||
data, _ = sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"input_json_delta","partial_json":""}}`, (*param).(*Params).ResponseIndex), "delta.partial_json", fcArgsResult.Raw)
|
||||
output = output + fmt.Sprintf("data: %s\n\n\n", data)
|
||||
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", data))
|
||||
}
|
||||
(*param).(*Params).ResponseType = 3
|
||||
}
|
||||
@@ -216,13 +199,13 @@ func ConvertGeminiCLIResponseToClaude(_ context.Context, _ string, originalReque
|
||||
if usageResult.Exists() && bytes.Contains(rawJSON, []byte(`"finishReason"`)) {
|
||||
if candidatesTokenCountResult := usageResult.Get("candidatesTokenCount"); candidatesTokenCountResult.Exists() {
|
||||
// Close the final content block
|
||||
output = output + "event: content_block_stop\n"
|
||||
output = output + fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex)
|
||||
output = output + "\n\n\n"
|
||||
sb.WriteString("event: content_block_stop\n")
|
||||
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex))
|
||||
sb.WriteString("\n\n\n")
|
||||
|
||||
// Send the final message delta with usage information and stop reason
|
||||
output = output + "event: message_delta\n"
|
||||
output = output + `data: `
|
||||
sb.WriteString("event: message_delta\n")
|
||||
sb.WriteString(`data: `)
|
||||
|
||||
// Create the message delta template with appropriate stop reason
|
||||
template := `{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"input_tokens":0,"output_tokens":0}}`
|
||||
@@ -236,11 +219,11 @@ func ConvertGeminiCLIResponseToClaude(_ context.Context, _ string, originalReque
|
||||
template, _ = sjson.Set(template, "usage.output_tokens", candidatesTokenCountResult.Int()+thoughtsTokenCount)
|
||||
template, _ = sjson.Set(template, "usage.input_tokens", usageResult.Get("promptTokenCount").Int())
|
||||
|
||||
output = output + template + "\n\n\n"
|
||||
sb.WriteString(template + "\n\n\n")
|
||||
}
|
||||
}
|
||||
|
||||
return []string{output}
|
||||
return []string{sb.String()}
|
||||
}
|
||||
|
||||
// ConvertGeminiCLIResponseToClaudeNonStream converts a non-streaming Gemini CLI response to a non-streaming Claude response.
|
||||
|
||||
@@ -33,4 +33,7 @@ import (
|
||||
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/antigravity/gemini"
|
||||
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/antigravity/openai/chat-completions"
|
||||
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/antigravity/openai/responses"
|
||||
|
||||
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/kiro/claude"
|
||||
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/kiro/openai/chat-completions"
|
||||
)
|
||||
|
||||
19
internal/translator/kiro/claude/init.go
Normal file
19
internal/translator/kiro/claude/init.go
Normal file
@@ -0,0 +1,19 @@
|
||||
package claude
|
||||
|
||||
import (
|
||||
. "github.com/router-for-me/CLIProxyAPI/v6/internal/constant"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/translator/translator"
|
||||
)
|
||||
|
||||
func init() {
|
||||
translator.Register(
|
||||
Claude,
|
||||
Kiro,
|
||||
ConvertClaudeRequestToKiro,
|
||||
interfaces.TranslateResponse{
|
||||
Stream: ConvertKiroResponseToClaude,
|
||||
NonStream: ConvertKiroResponseToClaudeNonStream,
|
||||
},
|
||||
)
|
||||
}
|
||||
24
internal/translator/kiro/claude/kiro_claude.go
Normal file
24
internal/translator/kiro/claude/kiro_claude.go
Normal file
@@ -0,0 +1,24 @@
|
||||
// Package claude provides translation between Kiro and Claude formats.
|
||||
// Since Kiro uses Claude-compatible format internally, translations are mostly pass-through.
|
||||
package claude
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
)
|
||||
|
||||
// ConvertClaudeRequestToKiro converts Claude request to Kiro format.
|
||||
// Since Kiro uses Claude format internally, this is mostly a pass-through.
|
||||
func ConvertClaudeRequestToKiro(modelName string, inputRawJSON []byte, stream bool) []byte {
|
||||
return bytes.Clone(inputRawJSON)
|
||||
}
|
||||
|
||||
// ConvertKiroResponseToClaude converts Kiro streaming response to Claude format.
|
||||
func ConvertKiroResponseToClaude(ctx context.Context, model string, originalRequest, request, rawResponse []byte, param *any) []string {
|
||||
return []string{string(rawResponse)}
|
||||
}
|
||||
|
||||
// ConvertKiroResponseToClaudeNonStream converts Kiro non-streaming response to Claude format.
|
||||
func ConvertKiroResponseToClaudeNonStream(ctx context.Context, model string, originalRequest, request, rawResponse []byte, param *any) string {
|
||||
return string(rawResponse)
|
||||
}
|
||||
19
internal/translator/kiro/openai/chat-completions/init.go
Normal file
19
internal/translator/kiro/openai/chat-completions/init.go
Normal file
@@ -0,0 +1,19 @@
|
||||
package chat_completions
|
||||
|
||||
import (
|
||||
. "github.com/router-for-me/CLIProxyAPI/v6/internal/constant"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/translator/translator"
|
||||
)
|
||||
|
||||
func init() {
|
||||
translator.Register(
|
||||
OpenAI,
|
||||
Kiro,
|
||||
ConvertOpenAIRequestToKiro,
|
||||
interfaces.TranslateResponse{
|
||||
Stream: ConvertKiroResponseToOpenAI,
|
||||
NonStream: ConvertKiroResponseToOpenAINonStream,
|
||||
},
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,258 @@
|
||||
// Package chat_completions provides request translation from OpenAI to Kiro format.
|
||||
package chat_completions
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
// ConvertOpenAIRequestToKiro transforms an OpenAI Chat Completions API request into Kiro (Claude) format.
|
||||
// Kiro uses Claude-compatible format internally, so we primarily pass through to Claude format.
|
||||
// Supports tool calling: OpenAI tools -> Claude tools, tool_calls -> tool_use, tool messages -> tool_result.
|
||||
func ConvertOpenAIRequestToKiro(modelName string, inputRawJSON []byte, stream bool) []byte {
|
||||
rawJSON := bytes.Clone(inputRawJSON)
|
||||
root := gjson.ParseBytes(rawJSON)
|
||||
|
||||
// Build Claude-compatible request
|
||||
out := `{"model":"","max_tokens":32000,"messages":[]}`
|
||||
|
||||
// Set model
|
||||
out, _ = sjson.Set(out, "model", modelName)
|
||||
|
||||
// Copy max_tokens if present
|
||||
if v := root.Get("max_tokens"); v.Exists() {
|
||||
out, _ = sjson.Set(out, "max_tokens", v.Int())
|
||||
}
|
||||
|
||||
// Copy temperature if present
|
||||
if v := root.Get("temperature"); v.Exists() {
|
||||
out, _ = sjson.Set(out, "temperature", v.Float())
|
||||
}
|
||||
|
||||
// Copy top_p if present
|
||||
if v := root.Get("top_p"); v.Exists() {
|
||||
out, _ = sjson.Set(out, "top_p", v.Float())
|
||||
}
|
||||
|
||||
// Convert OpenAI tools to Claude tools format
|
||||
if tools := root.Get("tools"); tools.Exists() && tools.IsArray() {
|
||||
claudeTools := make([]interface{}, 0)
|
||||
for _, tool := range tools.Array() {
|
||||
if tool.Get("type").String() == "function" {
|
||||
fn := tool.Get("function")
|
||||
claudeTool := map[string]interface{}{
|
||||
"name": fn.Get("name").String(),
|
||||
"description": fn.Get("description").String(),
|
||||
}
|
||||
// Convert parameters to input_schema
|
||||
if params := fn.Get("parameters"); params.Exists() {
|
||||
claudeTool["input_schema"] = params.Value()
|
||||
} else {
|
||||
claudeTool["input_schema"] = map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{},
|
||||
}
|
||||
}
|
||||
claudeTools = append(claudeTools, claudeTool)
|
||||
}
|
||||
}
|
||||
if len(claudeTools) > 0 {
|
||||
out, _ = sjson.Set(out, "tools", claudeTools)
|
||||
}
|
||||
}
|
||||
|
||||
// Process messages
|
||||
messages := root.Get("messages")
|
||||
if messages.Exists() && messages.IsArray() {
|
||||
claudeMessages := make([]interface{}, 0)
|
||||
var systemPrompt string
|
||||
|
||||
// Track pending tool results to merge with next user message
|
||||
var pendingToolResults []map[string]interface{}
|
||||
|
||||
for _, msg := range messages.Array() {
|
||||
role := msg.Get("role").String()
|
||||
content := msg.Get("content")
|
||||
|
||||
if role == "system" {
|
||||
// Extract system message
|
||||
if content.IsArray() {
|
||||
for _, part := range content.Array() {
|
||||
if part.Get("type").String() == "text" {
|
||||
systemPrompt += part.Get("text").String() + "\n"
|
||||
}
|
||||
}
|
||||
} else {
|
||||
systemPrompt = content.String()
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if role == "tool" {
|
||||
// OpenAI tool message -> Claude tool_result content block
|
||||
toolCallID := msg.Get("tool_call_id").String()
|
||||
toolContent := content.String()
|
||||
|
||||
toolResult := map[string]interface{}{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": toolCallID,
|
||||
}
|
||||
|
||||
// Handle content - can be string or structured
|
||||
if content.IsArray() {
|
||||
contentParts := make([]interface{}, 0)
|
||||
for _, part := range content.Array() {
|
||||
if part.Get("type").String() == "text" {
|
||||
contentParts = append(contentParts, map[string]interface{}{
|
||||
"type": "text",
|
||||
"text": part.Get("text").String(),
|
||||
})
|
||||
}
|
||||
}
|
||||
toolResult["content"] = contentParts
|
||||
} else {
|
||||
toolResult["content"] = toolContent
|
||||
}
|
||||
|
||||
pendingToolResults = append(pendingToolResults, toolResult)
|
||||
continue
|
||||
}
|
||||
|
||||
claudeMsg := map[string]interface{}{
|
||||
"role": role,
|
||||
}
|
||||
|
||||
// Handle assistant messages with tool_calls
|
||||
if role == "assistant" && msg.Get("tool_calls").Exists() {
|
||||
contentParts := make([]interface{}, 0)
|
||||
|
||||
// Add text content if present
|
||||
if content.Exists() && content.String() != "" {
|
||||
contentParts = append(contentParts, map[string]interface{}{
|
||||
"type": "text",
|
||||
"text": content.String(),
|
||||
})
|
||||
}
|
||||
|
||||
// Convert tool_calls to tool_use blocks
|
||||
for _, toolCall := range msg.Get("tool_calls").Array() {
|
||||
toolUseID := toolCall.Get("id").String()
|
||||
fnName := toolCall.Get("function.name").String()
|
||||
fnArgs := toolCall.Get("function.arguments").String()
|
||||
|
||||
// Parse arguments JSON
|
||||
var argsMap map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(fnArgs), &argsMap); err != nil {
|
||||
argsMap = map[string]interface{}{"raw": fnArgs}
|
||||
}
|
||||
|
||||
contentParts = append(contentParts, map[string]interface{}{
|
||||
"type": "tool_use",
|
||||
"id": toolUseID,
|
||||
"name": fnName,
|
||||
"input": argsMap,
|
||||
})
|
||||
}
|
||||
|
||||
claudeMsg["content"] = contentParts
|
||||
claudeMessages = append(claudeMessages, claudeMsg)
|
||||
continue
|
||||
}
|
||||
|
||||
// Handle user messages - may need to include pending tool results
|
||||
if role == "user" && len(pendingToolResults) > 0 {
|
||||
contentParts := make([]interface{}, 0)
|
||||
|
||||
// Add pending tool results first
|
||||
for _, tr := range pendingToolResults {
|
||||
contentParts = append(contentParts, tr)
|
||||
}
|
||||
pendingToolResults = nil
|
||||
|
||||
// Add user content
|
||||
if content.IsArray() {
|
||||
for _, part := range content.Array() {
|
||||
partType := part.Get("type").String()
|
||||
if partType == "text" {
|
||||
contentParts = append(contentParts, map[string]interface{}{
|
||||
"type": "text",
|
||||
"text": part.Get("text").String(),
|
||||
})
|
||||
} else if partType == "image_url" {
|
||||
contentParts = append(contentParts, map[string]interface{}{
|
||||
"type": "image",
|
||||
"source": map[string]interface{}{
|
||||
"type": "url",
|
||||
"url": part.Get("image_url.url").String(),
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
} else if content.String() != "" {
|
||||
contentParts = append(contentParts, map[string]interface{}{
|
||||
"type": "text",
|
||||
"text": content.String(),
|
||||
})
|
||||
}
|
||||
|
||||
claudeMsg["content"] = contentParts
|
||||
claudeMessages = append(claudeMessages, claudeMsg)
|
||||
continue
|
||||
}
|
||||
|
||||
// Handle regular content
|
||||
if content.IsArray() {
|
||||
contentParts := make([]interface{}, 0)
|
||||
for _, part := range content.Array() {
|
||||
partType := part.Get("type").String()
|
||||
if partType == "text" {
|
||||
contentParts = append(contentParts, map[string]interface{}{
|
||||
"type": "text",
|
||||
"text": part.Get("text").String(),
|
||||
})
|
||||
} else if partType == "image_url" {
|
||||
contentParts = append(contentParts, map[string]interface{}{
|
||||
"type": "image",
|
||||
"source": map[string]interface{}{
|
||||
"type": "url",
|
||||
"url": part.Get("image_url.url").String(),
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
claudeMsg["content"] = contentParts
|
||||
} else {
|
||||
claudeMsg["content"] = content.String()
|
||||
}
|
||||
|
||||
claudeMessages = append(claudeMessages, claudeMsg)
|
||||
}
|
||||
|
||||
// If there are pending tool results without a following user message,
|
||||
// create a user message with just the tool results
|
||||
if len(pendingToolResults) > 0 {
|
||||
contentParts := make([]interface{}, 0)
|
||||
for _, tr := range pendingToolResults {
|
||||
contentParts = append(contentParts, tr)
|
||||
}
|
||||
claudeMessages = append(claudeMessages, map[string]interface{}{
|
||||
"role": "user",
|
||||
"content": contentParts,
|
||||
})
|
||||
}
|
||||
|
||||
out, _ = sjson.Set(out, "messages", claudeMessages)
|
||||
|
||||
if systemPrompt != "" {
|
||||
out, _ = sjson.Set(out, "system", systemPrompt)
|
||||
}
|
||||
}
|
||||
|
||||
// Set stream
|
||||
out, _ = sjson.Set(out, "stream", stream)
|
||||
|
||||
return []byte(out)
|
||||
}
|
||||
@@ -0,0 +1,316 @@
|
||||
// Package chat_completions provides response translation from Kiro to OpenAI format.
|
||||
package chat_completions
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
// ConvertKiroResponseToOpenAI converts Kiro streaming response to OpenAI SSE format.
|
||||
// Handles Claude SSE events: content_block_start, content_block_delta, input_json_delta,
|
||||
// content_block_stop, message_delta, and message_stop.
|
||||
func ConvertKiroResponseToOpenAI(ctx context.Context, model string, originalRequest, request, rawResponse []byte, param *any) []string {
|
||||
root := gjson.ParseBytes(rawResponse)
|
||||
var results []string
|
||||
|
||||
eventType := root.Get("type").String()
|
||||
|
||||
switch eventType {
|
||||
case "message_start":
|
||||
// Initial message event - could emit initial chunk if needed
|
||||
return results
|
||||
|
||||
case "content_block_start":
|
||||
// Start of a content block (text or tool_use)
|
||||
blockType := root.Get("content_block.type").String()
|
||||
index := int(root.Get("index").Int())
|
||||
|
||||
if blockType == "tool_use" {
|
||||
// Start of tool_use block
|
||||
toolUseID := root.Get("content_block.id").String()
|
||||
toolName := root.Get("content_block.name").String()
|
||||
|
||||
toolCall := map[string]interface{}{
|
||||
"index": index,
|
||||
"id": toolUseID,
|
||||
"type": "function",
|
||||
"function": map[string]interface{}{
|
||||
"name": toolName,
|
||||
"arguments": "",
|
||||
},
|
||||
}
|
||||
|
||||
response := map[string]interface{}{
|
||||
"id": "chatcmpl-" + uuid.New().String()[:24],
|
||||
"object": "chat.completion.chunk",
|
||||
"created": time.Now().Unix(),
|
||||
"model": model,
|
||||
"choices": []map[string]interface{}{
|
||||
{
|
||||
"index": 0,
|
||||
"delta": map[string]interface{}{
|
||||
"tool_calls": []map[string]interface{}{toolCall},
|
||||
},
|
||||
"finish_reason": nil,
|
||||
},
|
||||
},
|
||||
}
|
||||
result, _ := json.Marshal(response)
|
||||
results = append(results, string(result))
|
||||
}
|
||||
return results
|
||||
|
||||
case "content_block_delta":
|
||||
index := int(root.Get("index").Int())
|
||||
deltaType := root.Get("delta.type").String()
|
||||
|
||||
if deltaType == "text_delta" {
|
||||
// Text content delta
|
||||
contentDelta := root.Get("delta.text").String()
|
||||
if contentDelta != "" {
|
||||
response := map[string]interface{}{
|
||||
"id": "chatcmpl-" + uuid.New().String()[:24],
|
||||
"object": "chat.completion.chunk",
|
||||
"created": time.Now().Unix(),
|
||||
"model": model,
|
||||
"choices": []map[string]interface{}{
|
||||
{
|
||||
"index": 0,
|
||||
"delta": map[string]interface{}{
|
||||
"content": contentDelta,
|
||||
},
|
||||
"finish_reason": nil,
|
||||
},
|
||||
},
|
||||
}
|
||||
result, _ := json.Marshal(response)
|
||||
results = append(results, string(result))
|
||||
}
|
||||
} else if deltaType == "input_json_delta" {
|
||||
// Tool input delta (streaming arguments)
|
||||
partialJSON := root.Get("delta.partial_json").String()
|
||||
if partialJSON != "" {
|
||||
toolCall := map[string]interface{}{
|
||||
"index": index,
|
||||
"function": map[string]interface{}{
|
||||
"arguments": partialJSON,
|
||||
},
|
||||
}
|
||||
|
||||
response := map[string]interface{}{
|
||||
"id": "chatcmpl-" + uuid.New().String()[:24],
|
||||
"object": "chat.completion.chunk",
|
||||
"created": time.Now().Unix(),
|
||||
"model": model,
|
||||
"choices": []map[string]interface{}{
|
||||
{
|
||||
"index": 0,
|
||||
"delta": map[string]interface{}{
|
||||
"tool_calls": []map[string]interface{}{toolCall},
|
||||
},
|
||||
"finish_reason": nil,
|
||||
},
|
||||
},
|
||||
}
|
||||
result, _ := json.Marshal(response)
|
||||
results = append(results, string(result))
|
||||
}
|
||||
}
|
||||
return results
|
||||
|
||||
case "content_block_stop":
|
||||
// End of content block - no output needed for OpenAI format
|
||||
return results
|
||||
|
||||
case "message_delta":
|
||||
// Final message delta with stop_reason
|
||||
stopReason := root.Get("delta.stop_reason").String()
|
||||
if stopReason != "" {
|
||||
finishReason := "stop"
|
||||
if stopReason == "tool_use" {
|
||||
finishReason = "tool_calls"
|
||||
} else if stopReason == "end_turn" {
|
||||
finishReason = "stop"
|
||||
} else if stopReason == "max_tokens" {
|
||||
finishReason = "length"
|
||||
}
|
||||
|
||||
response := map[string]interface{}{
|
||||
"id": "chatcmpl-" + uuid.New().String()[:24],
|
||||
"object": "chat.completion.chunk",
|
||||
"created": time.Now().Unix(),
|
||||
"model": model,
|
||||
"choices": []map[string]interface{}{
|
||||
{
|
||||
"index": 0,
|
||||
"delta": map[string]interface{}{},
|
||||
"finish_reason": finishReason,
|
||||
},
|
||||
},
|
||||
}
|
||||
result, _ := json.Marshal(response)
|
||||
results = append(results, string(result))
|
||||
}
|
||||
return results
|
||||
|
||||
case "message_stop":
|
||||
// End of message - could emit [DONE] marker
|
||||
return results
|
||||
}
|
||||
|
||||
// Fallback: handle raw content for backward compatibility
|
||||
var contentDelta string
|
||||
if delta := root.Get("delta.text"); delta.Exists() {
|
||||
contentDelta = delta.String()
|
||||
} else if content := root.Get("content"); content.Exists() && root.Get("type").String() == "" {
|
||||
contentDelta = content.String()
|
||||
}
|
||||
|
||||
if contentDelta != "" {
|
||||
response := map[string]interface{}{
|
||||
"id": "chatcmpl-" + uuid.New().String()[:24],
|
||||
"object": "chat.completion.chunk",
|
||||
"created": time.Now().Unix(),
|
||||
"model": model,
|
||||
"choices": []map[string]interface{}{
|
||||
{
|
||||
"index": 0,
|
||||
"delta": map[string]interface{}{
|
||||
"content": contentDelta,
|
||||
},
|
||||
"finish_reason": nil,
|
||||
},
|
||||
},
|
||||
}
|
||||
result, _ := json.Marshal(response)
|
||||
results = append(results, string(result))
|
||||
}
|
||||
|
||||
// Handle tool_use content blocks (Claude format) - fallback
|
||||
toolUses := root.Get("delta.tool_use")
|
||||
if !toolUses.Exists() {
|
||||
toolUses = root.Get("tool_use")
|
||||
}
|
||||
if toolUses.Exists() && toolUses.IsObject() {
|
||||
inputJSON := toolUses.Get("input").String()
|
||||
if inputJSON == "" {
|
||||
if inputObj := toolUses.Get("input"); inputObj.Exists() {
|
||||
inputBytes, _ := json.Marshal(inputObj.Value())
|
||||
inputJSON = string(inputBytes)
|
||||
}
|
||||
}
|
||||
|
||||
toolCall := map[string]interface{}{
|
||||
"index": 0,
|
||||
"id": toolUses.Get("id").String(),
|
||||
"type": "function",
|
||||
"function": map[string]interface{}{
|
||||
"name": toolUses.Get("name").String(),
|
||||
"arguments": inputJSON,
|
||||
},
|
||||
}
|
||||
|
||||
response := map[string]interface{}{
|
||||
"id": "chatcmpl-" + uuid.New().String()[:24],
|
||||
"object": "chat.completion.chunk",
|
||||
"created": time.Now().Unix(),
|
||||
"model": model,
|
||||
"choices": []map[string]interface{}{
|
||||
{
|
||||
"index": 0,
|
||||
"delta": map[string]interface{}{
|
||||
"tool_calls": []map[string]interface{}{toolCall},
|
||||
},
|
||||
"finish_reason": nil,
|
||||
},
|
||||
},
|
||||
}
|
||||
result, _ := json.Marshal(response)
|
||||
results = append(results, string(result))
|
||||
}
|
||||
|
||||
return results
|
||||
}
|
||||
|
||||
// ConvertKiroResponseToOpenAINonStream converts Kiro non-streaming response to OpenAI format.
|
||||
func ConvertKiroResponseToOpenAINonStream(ctx context.Context, model string, originalRequest, request, rawResponse []byte, param *any) string {
|
||||
root := gjson.ParseBytes(rawResponse)
|
||||
|
||||
var content string
|
||||
var toolCalls []map[string]interface{}
|
||||
|
||||
contentArray := root.Get("content")
|
||||
if contentArray.IsArray() {
|
||||
for _, item := range contentArray.Array() {
|
||||
itemType := item.Get("type").String()
|
||||
if itemType == "text" {
|
||||
content += item.Get("text").String()
|
||||
} else if itemType == "tool_use" {
|
||||
// Convert Claude tool_use to OpenAI tool_calls format
|
||||
inputJSON := item.Get("input").String()
|
||||
if inputJSON == "" {
|
||||
// If input is an object, marshal it
|
||||
if inputObj := item.Get("input"); inputObj.Exists() {
|
||||
inputBytes, _ := json.Marshal(inputObj.Value())
|
||||
inputJSON = string(inputBytes)
|
||||
}
|
||||
}
|
||||
toolCall := map[string]interface{}{
|
||||
"id": item.Get("id").String(),
|
||||
"type": "function",
|
||||
"function": map[string]interface{}{
|
||||
"name": item.Get("name").String(),
|
||||
"arguments": inputJSON,
|
||||
},
|
||||
}
|
||||
toolCalls = append(toolCalls, toolCall)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
content = root.Get("content").String()
|
||||
}
|
||||
|
||||
inputTokens := root.Get("usage.input_tokens").Int()
|
||||
outputTokens := root.Get("usage.output_tokens").Int()
|
||||
|
||||
message := map[string]interface{}{
|
||||
"role": "assistant",
|
||||
"content": content,
|
||||
}
|
||||
|
||||
// Add tool_calls if present
|
||||
if len(toolCalls) > 0 {
|
||||
message["tool_calls"] = toolCalls
|
||||
}
|
||||
|
||||
finishReason := "stop"
|
||||
if len(toolCalls) > 0 {
|
||||
finishReason = "tool_calls"
|
||||
}
|
||||
|
||||
response := map[string]interface{}{
|
||||
"id": "chatcmpl-" + uuid.New().String()[:24],
|
||||
"object": "chat.completion",
|
||||
"created": time.Now().Unix(),
|
||||
"model": model,
|
||||
"choices": []map[string]interface{}{
|
||||
{
|
||||
"index": 0,
|
||||
"message": message,
|
||||
"finish_reason": finishReason,
|
||||
},
|
||||
},
|
||||
"usage": map[string]interface{}{
|
||||
"prompt_tokens": inputTokens,
|
||||
"completion_tokens": outputTokens,
|
||||
"total_tokens": inputTokens + outputTokens,
|
||||
},
|
||||
}
|
||||
|
||||
result, _ := json.Marshal(response)
|
||||
return string(result)
|
||||
}
|
||||
@@ -21,6 +21,7 @@ import (
|
||||
|
||||
"github.com/fsnotify/fsnotify"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
kiroauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/kiro"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/geminicli"
|
||||
"gopkg.in/yaml.v3"
|
||||
|
||||
@@ -176,6 +177,9 @@ func (w *Watcher) Start(ctx context.Context) error {
|
||||
}
|
||||
log.Debugf("watching auth directory: %s", w.authDir)
|
||||
|
||||
// Watch Kiro IDE token file directory for automatic token updates
|
||||
w.watchKiroIDETokenFile()
|
||||
|
||||
// Start the event processing goroutine
|
||||
go w.processEvents(ctx)
|
||||
|
||||
@@ -184,6 +188,31 @@ func (w *Watcher) Start(ctx context.Context) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// watchKiroIDETokenFile adds the Kiro IDE token file directory to the watcher.
|
||||
// This enables automatic detection of token updates from Kiro IDE.
|
||||
func (w *Watcher) watchKiroIDETokenFile() {
|
||||
homeDir, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
log.Debugf("failed to get home directory for Kiro IDE token watch: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Kiro IDE stores tokens in ~/.aws/sso/cache/
|
||||
kiroTokenDir := filepath.Join(homeDir, ".aws", "sso", "cache")
|
||||
|
||||
// Check if directory exists
|
||||
if _, statErr := os.Stat(kiroTokenDir); os.IsNotExist(statErr) {
|
||||
log.Debugf("Kiro IDE token directory does not exist: %s", kiroTokenDir)
|
||||
return
|
||||
}
|
||||
|
||||
if errAdd := w.watcher.Add(kiroTokenDir); errAdd != nil {
|
||||
log.Debugf("failed to watch Kiro IDE token directory %s: %v", kiroTokenDir, errAdd)
|
||||
return
|
||||
}
|
||||
log.Debugf("watching Kiro IDE token directory: %s", kiroTokenDir)
|
||||
}
|
||||
|
||||
// Stop stops the file watcher
|
||||
func (w *Watcher) Stop() error {
|
||||
w.stopDispatch()
|
||||
@@ -744,10 +773,20 @@ func (w *Watcher) handleEvent(event fsnotify.Event) {
|
||||
isConfigEvent := event.Name == w.configPath && event.Op&configOps != 0
|
||||
authOps := fsnotify.Create | fsnotify.Write | fsnotify.Remove | fsnotify.Rename
|
||||
isAuthJSON := strings.HasPrefix(event.Name, w.authDir) && strings.HasSuffix(event.Name, ".json") && event.Op&authOps != 0
|
||||
if !isConfigEvent && !isAuthJSON {
|
||||
|
||||
// Check for Kiro IDE token file changes
|
||||
isKiroIDEToken := w.isKiroIDETokenFile(event.Name) && event.Op&authOps != 0
|
||||
|
||||
if !isConfigEvent && !isAuthJSON && !isKiroIDEToken {
|
||||
// Ignore unrelated files (e.g., cookie snapshots *.cookie) and other noise.
|
||||
return
|
||||
}
|
||||
|
||||
// Handle Kiro IDE token file changes
|
||||
if isKiroIDEToken {
|
||||
w.handleKiroIDETokenChange(event)
|
||||
return
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
log.Debugf("file system event detected: %s %s", event.Op.String(), event.Name)
|
||||
@@ -805,6 +844,51 @@ func (w *Watcher) scheduleConfigReload() {
|
||||
})
|
||||
}
|
||||
|
||||
// isKiroIDETokenFile checks if the given path is the Kiro IDE token file.
|
||||
func (w *Watcher) isKiroIDETokenFile(path string) bool {
|
||||
// Check if it's the kiro-auth-token.json file in ~/.aws/sso/cache/
|
||||
// Use filepath.ToSlash to ensure consistent separators across platforms (Windows uses backslashes)
|
||||
normalized := filepath.ToSlash(path)
|
||||
return strings.HasSuffix(normalized, "kiro-auth-token.json") && strings.Contains(normalized, ".aws/sso/cache")
|
||||
}
|
||||
|
||||
// handleKiroIDETokenChange processes changes to the Kiro IDE token file.
|
||||
// When the token file is updated by Kiro IDE, this triggers a reload of Kiro auth.
|
||||
func (w *Watcher) handleKiroIDETokenChange(event fsnotify.Event) {
|
||||
log.Debugf("Kiro IDE token file event detected: %s %s", event.Op.String(), event.Name)
|
||||
|
||||
if event.Op&(fsnotify.Remove|fsnotify.Rename) != 0 {
|
||||
// Token file removed - wait briefly for potential atomic replace
|
||||
time.Sleep(replaceCheckDelay)
|
||||
if _, statErr := os.Stat(event.Name); statErr != nil {
|
||||
log.Debugf("Kiro IDE token file removed: %s", event.Name)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Try to load the updated token
|
||||
tokenData, err := kiroauth.LoadKiroIDEToken()
|
||||
if err != nil {
|
||||
log.Debugf("failed to load Kiro IDE token after change: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
log.Infof("Kiro IDE token file updated, access token refreshed (provider: %s)", tokenData.Provider)
|
||||
|
||||
// Trigger auth state refresh to pick up the new token
|
||||
w.refreshAuthState()
|
||||
|
||||
// Notify callback if set
|
||||
w.clientsMutex.RLock()
|
||||
cfg := w.config
|
||||
w.clientsMutex.RUnlock()
|
||||
|
||||
if w.reloadCallback != nil && cfg != nil {
|
||||
log.Debugf("triggering server update callback after Kiro IDE token change")
|
||||
w.reloadCallback(cfg)
|
||||
}
|
||||
}
|
||||
|
||||
func (w *Watcher) reloadConfigIfChanged() {
|
||||
data, err := os.ReadFile(w.configPath)
|
||||
if err != nil {
|
||||
@@ -1181,6 +1265,67 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
|
||||
applyAuthExcludedModelsMeta(a, cfg, ck.ExcludedModels, "apikey")
|
||||
out = append(out, a)
|
||||
}
|
||||
// Kiro (AWS CodeWhisperer) -> synthesize auths
|
||||
var kAuth *kiroauth.KiroAuth
|
||||
if len(cfg.KiroKey) > 0 {
|
||||
kAuth = kiroauth.NewKiroAuth(cfg)
|
||||
}
|
||||
for i := range cfg.KiroKey {
|
||||
kk := cfg.KiroKey[i]
|
||||
var accessToken, profileArn string
|
||||
|
||||
// Try to load from token file first
|
||||
if kk.TokenFile != "" && kAuth != nil {
|
||||
tokenData, err := kAuth.LoadTokenFromFile(kk.TokenFile)
|
||||
if err != nil {
|
||||
log.Warnf("failed to load kiro token file %s: %v", kk.TokenFile, err)
|
||||
} else {
|
||||
accessToken = tokenData.AccessToken
|
||||
profileArn = tokenData.ProfileArn
|
||||
}
|
||||
}
|
||||
|
||||
// Override with direct config values if provided
|
||||
if kk.AccessToken != "" {
|
||||
accessToken = kk.AccessToken
|
||||
}
|
||||
if kk.ProfileArn != "" {
|
||||
profileArn = kk.ProfileArn
|
||||
}
|
||||
|
||||
if accessToken == "" {
|
||||
log.Warnf("kiro config[%d] missing access_token, skipping", i)
|
||||
continue
|
||||
}
|
||||
|
||||
// profileArn is optional for AWS Builder ID users
|
||||
id, token := idGen.next("kiro:token", accessToken, profileArn)
|
||||
attrs := map[string]string{
|
||||
"source": fmt.Sprintf("config:kiro[%s]", token),
|
||||
"access_token": accessToken,
|
||||
}
|
||||
if profileArn != "" {
|
||||
attrs["profile_arn"] = profileArn
|
||||
}
|
||||
if kk.Region != "" {
|
||||
attrs["region"] = kk.Region
|
||||
}
|
||||
if kk.AgentTaskType != "" {
|
||||
attrs["agent_task_type"] = kk.AgentTaskType
|
||||
}
|
||||
proxyURL := strings.TrimSpace(kk.ProxyURL)
|
||||
a := &coreauth.Auth{
|
||||
ID: id,
|
||||
Provider: "kiro",
|
||||
Label: "kiro-token",
|
||||
Status: coreauth.StatusActive,
|
||||
ProxyURL: proxyURL,
|
||||
Attributes: attrs,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
out = append(out, a)
|
||||
}
|
||||
for i := range cfg.OpenAICompatibility {
|
||||
compat := &cfg.OpenAICompatibility[i]
|
||||
providerName := strings.ToLower(strings.TrimSpace(compat.Name))
|
||||
@@ -1287,7 +1432,12 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
|
||||
}
|
||||
|
||||
// Also synthesize auth entries directly from auth files (for OAuth/file-backed providers)
|
||||
entries, _ := os.ReadDir(w.authDir)
|
||||
log.Debugf("SnapshotCoreAuths: scanning auth directory: %s", w.authDir)
|
||||
entries, readErr := os.ReadDir(w.authDir)
|
||||
if readErr != nil {
|
||||
log.Errorf("SnapshotCoreAuths: failed to read auth directory %s: %v", w.authDir, readErr)
|
||||
}
|
||||
log.Debugf("SnapshotCoreAuths: found %d entries in auth directory", len(entries))
|
||||
for _, e := range entries {
|
||||
if e.IsDir() {
|
||||
continue
|
||||
@@ -1306,9 +1456,20 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
|
||||
continue
|
||||
}
|
||||
t, _ := metadata["type"].(string)
|
||||
|
||||
// Detect Kiro auth files by auth_method field (they don't have "type" field)
|
||||
if t == "" {
|
||||
if authMethod, _ := metadata["auth_method"].(string); authMethod == "builder-id" || authMethod == "social" {
|
||||
t = "kiro"
|
||||
log.Debugf("SnapshotCoreAuths: detected Kiro auth by auth_method: %s", name)
|
||||
}
|
||||
}
|
||||
|
||||
if t == "" {
|
||||
log.Debugf("SnapshotCoreAuths: skipping file without type: %s", name)
|
||||
continue
|
||||
}
|
||||
log.Debugf("SnapshotCoreAuths: processing auth file: %s (type=%s)", name, t)
|
||||
provider := strings.ToLower(t)
|
||||
if provider == "gemini" {
|
||||
provider = "gemini-cli"
|
||||
@@ -1317,6 +1478,12 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
|
||||
if email, _ := metadata["email"].(string); email != "" {
|
||||
label = email
|
||||
}
|
||||
// For Kiro, use provider field as label if available
|
||||
if provider == "kiro" {
|
||||
if kiroProvider, _ := metadata["provider"].(string); kiroProvider != "" {
|
||||
label = fmt.Sprintf("kiro-%s", strings.ToLower(kiroProvider))
|
||||
}
|
||||
}
|
||||
// Use relative path under authDir as ID to stay consistent with the file-based token store
|
||||
id := full
|
||||
if rel, errRel := filepath.Rel(w.authDir, full); errRel == nil && rel != "" {
|
||||
@@ -1342,6 +1509,16 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
// Set NextRefreshAfter for Kiro auth based on expires_at
|
||||
if provider == "kiro" {
|
||||
if expiresAtStr, ok := metadata["expires_at"].(string); ok && expiresAtStr != "" {
|
||||
if expiresAt, parseErr := time.Parse(time.RFC3339, expiresAtStr); parseErr == nil {
|
||||
// Refresh 30 minutes before expiry
|
||||
a.NextRefreshAfter = expiresAt.Add(-30 * time.Minute)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
applyAuthExcludedModelsMeta(a, cfg, nil, "oauth")
|
||||
if provider == "gemini-cli" {
|
||||
if virtuals := synthesizeGeminiVirtualAuths(a, metadata, now); len(virtuals) > 0 {
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
|
||||
@@ -50,7 +51,25 @@ type BaseAPIHandler struct {
|
||||
Cfg *config.SDKConfig
|
||||
|
||||
// OpenAICompatProviders is a list of provider names for OpenAI compatibility.
|
||||
OpenAICompatProviders []string
|
||||
openAICompatProviders []string
|
||||
openAICompatMutex sync.RWMutex
|
||||
}
|
||||
|
||||
// GetOpenAICompatProviders safely returns a copy of the provider names
|
||||
func (h *BaseAPIHandler) GetOpenAICompatProviders() []string {
|
||||
h.openAICompatMutex.RLock()
|
||||
defer h.openAICompatMutex.RUnlock()
|
||||
result := make([]string, len(h.openAICompatProviders))
|
||||
copy(result, h.openAICompatProviders)
|
||||
return result
|
||||
}
|
||||
|
||||
// SetOpenAICompatProviders safely sets the provider names
|
||||
func (h *BaseAPIHandler) SetOpenAICompatProviders(providers []string) {
|
||||
h.openAICompatMutex.Lock()
|
||||
defer h.openAICompatMutex.Unlock()
|
||||
h.openAICompatProviders = make([]string, len(providers))
|
||||
copy(h.openAICompatProviders, providers)
|
||||
}
|
||||
|
||||
// NewBaseAPIHandlers creates a new API handlers instance.
|
||||
@@ -63,11 +82,12 @@ type BaseAPIHandler struct {
|
||||
// Returns:
|
||||
// - *BaseAPIHandler: A new API handlers instance
|
||||
func NewBaseAPIHandlers(cfg *config.SDKConfig, authManager *coreauth.Manager, openAICompatProviders []string) *BaseAPIHandler {
|
||||
return &BaseAPIHandler{
|
||||
Cfg: cfg,
|
||||
AuthManager: authManager,
|
||||
OpenAICompatProviders: openAICompatProviders,
|
||||
h := &BaseAPIHandler{
|
||||
Cfg: cfg,
|
||||
AuthManager: authManager,
|
||||
}
|
||||
h.SetOpenAICompatProviders(openAICompatProviders)
|
||||
return h
|
||||
}
|
||||
|
||||
// UpdateClients updates the handlers' client list and configuration.
|
||||
@@ -363,7 +383,7 @@ func (h *BaseAPIHandler) parseDynamicModel(modelName string) (providerName, mode
|
||||
}
|
||||
|
||||
// Check if the provider is a configured openai-compatibility provider
|
||||
for _, pName := range h.OpenAICompatProviders {
|
||||
for _, pName := range h.GetOpenAICompatProviders() {
|
||||
if pName == providerPart {
|
||||
return providerPart, modelPart, true
|
||||
}
|
||||
|
||||
357
sdk/auth/kiro.go
Normal file
357
sdk/auth/kiro.go
Normal file
@@ -0,0 +1,357 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
kiroauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/kiro"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
// extractKiroIdentifier extracts a meaningful identifier for file naming.
|
||||
// Returns account name if provided, otherwise profile ARN ID.
|
||||
// All extracted values are sanitized to prevent path injection attacks.
|
||||
func extractKiroIdentifier(accountName, profileArn string) string {
|
||||
// Priority 1: Use account name if provided
|
||||
if accountName != "" {
|
||||
return kiroauth.SanitizeEmailForFilename(accountName)
|
||||
}
|
||||
|
||||
// Priority 2: Use profile ARN ID part (sanitized to prevent path injection)
|
||||
if profileArn != "" {
|
||||
parts := strings.Split(profileArn, "/")
|
||||
if len(parts) >= 2 {
|
||||
// Sanitize the ARN component to prevent path traversal
|
||||
return kiroauth.SanitizeEmailForFilename(parts[len(parts)-1])
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: timestamp
|
||||
return fmt.Sprintf("%d", time.Now().UnixNano()%100000)
|
||||
}
|
||||
|
||||
// KiroAuthenticator implements OAuth authentication for Kiro with Google login.
|
||||
type KiroAuthenticator struct{}
|
||||
|
||||
// NewKiroAuthenticator constructs a Kiro authenticator.
|
||||
func NewKiroAuthenticator() *KiroAuthenticator {
|
||||
return &KiroAuthenticator{}
|
||||
}
|
||||
|
||||
// Provider returns the provider key for the authenticator.
|
||||
func (a *KiroAuthenticator) Provider() string {
|
||||
return "kiro"
|
||||
}
|
||||
|
||||
// RefreshLead indicates how soon before expiry a refresh should be attempted.
|
||||
func (a *KiroAuthenticator) RefreshLead() *time.Duration {
|
||||
d := 30 * time.Minute
|
||||
return &d
|
||||
}
|
||||
|
||||
// Login performs OAuth login for Kiro with AWS Builder ID.
|
||||
func (a *KiroAuthenticator) Login(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
|
||||
if cfg == nil {
|
||||
return nil, fmt.Errorf("kiro auth: configuration is required")
|
||||
}
|
||||
|
||||
oauth := kiroauth.NewKiroOAuth(cfg)
|
||||
|
||||
// Use AWS Builder ID device code flow
|
||||
tokenData, err := oauth.LoginWithBuilderID(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("login failed: %w", err)
|
||||
}
|
||||
|
||||
// Parse expires_at
|
||||
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
|
||||
if err != nil {
|
||||
expiresAt = time.Now().Add(1 * time.Hour)
|
||||
}
|
||||
|
||||
// Extract identifier for file naming
|
||||
idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
|
||||
|
||||
now := time.Now()
|
||||
fileName := fmt.Sprintf("kiro-aws-%s.json", idPart)
|
||||
|
||||
record := &coreauth.Auth{
|
||||
ID: fileName,
|
||||
Provider: "kiro",
|
||||
FileName: fileName,
|
||||
Label: "kiro-aws",
|
||||
Status: coreauth.StatusActive,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
Metadata: map[string]any{
|
||||
"type": "kiro",
|
||||
"access_token": tokenData.AccessToken,
|
||||
"refresh_token": tokenData.RefreshToken,
|
||||
"profile_arn": tokenData.ProfileArn,
|
||||
"expires_at": tokenData.ExpiresAt,
|
||||
"auth_method": tokenData.AuthMethod,
|
||||
"provider": tokenData.Provider,
|
||||
"client_id": tokenData.ClientID,
|
||||
"client_secret": tokenData.ClientSecret,
|
||||
"email": tokenData.Email,
|
||||
},
|
||||
Attributes: map[string]string{
|
||||
"profile_arn": tokenData.ProfileArn,
|
||||
"source": "aws-builder-id",
|
||||
"email": tokenData.Email,
|
||||
},
|
||||
NextRefreshAfter: expiresAt.Add(-30 * time.Minute),
|
||||
}
|
||||
|
||||
if tokenData.Email != "" {
|
||||
fmt.Printf("\n✓ Kiro authentication completed successfully! (Account: %s)\n", tokenData.Email)
|
||||
} else {
|
||||
fmt.Println("\n✓ Kiro authentication completed successfully!")
|
||||
}
|
||||
|
||||
return record, nil
|
||||
}
|
||||
|
||||
// LoginWithGoogle performs OAuth login for Kiro with Google.
|
||||
// This uses a custom protocol handler (kiro://) to receive the callback.
|
||||
func (a *KiroAuthenticator) LoginWithGoogle(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
|
||||
if cfg == nil {
|
||||
return nil, fmt.Errorf("kiro auth: configuration is required")
|
||||
}
|
||||
|
||||
oauth := kiroauth.NewKiroOAuth(cfg)
|
||||
|
||||
// Use Google OAuth flow with protocol handler
|
||||
tokenData, err := oauth.LoginWithGoogle(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("google login failed: %w", err)
|
||||
}
|
||||
|
||||
// Parse expires_at
|
||||
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
|
||||
if err != nil {
|
||||
expiresAt = time.Now().Add(1 * time.Hour)
|
||||
}
|
||||
|
||||
// Extract identifier for file naming
|
||||
idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
|
||||
|
||||
now := time.Now()
|
||||
fileName := fmt.Sprintf("kiro-google-%s.json", idPart)
|
||||
|
||||
record := &coreauth.Auth{
|
||||
ID: fileName,
|
||||
Provider: "kiro",
|
||||
FileName: fileName,
|
||||
Label: "kiro-google",
|
||||
Status: coreauth.StatusActive,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
Metadata: map[string]any{
|
||||
"type": "kiro",
|
||||
"access_token": tokenData.AccessToken,
|
||||
"refresh_token": tokenData.RefreshToken,
|
||||
"profile_arn": tokenData.ProfileArn,
|
||||
"expires_at": tokenData.ExpiresAt,
|
||||
"auth_method": tokenData.AuthMethod,
|
||||
"provider": tokenData.Provider,
|
||||
"email": tokenData.Email,
|
||||
},
|
||||
Attributes: map[string]string{
|
||||
"profile_arn": tokenData.ProfileArn,
|
||||
"source": "google-oauth",
|
||||
"email": tokenData.Email,
|
||||
},
|
||||
NextRefreshAfter: expiresAt.Add(-30 * time.Minute),
|
||||
}
|
||||
|
||||
if tokenData.Email != "" {
|
||||
fmt.Printf("\n✓ Kiro Google authentication completed successfully! (Account: %s)\n", tokenData.Email)
|
||||
} else {
|
||||
fmt.Println("\n✓ Kiro Google authentication completed successfully!")
|
||||
}
|
||||
|
||||
return record, nil
|
||||
}
|
||||
|
||||
// LoginWithGitHub performs OAuth login for Kiro with GitHub.
|
||||
// This uses a custom protocol handler (kiro://) to receive the callback.
|
||||
func (a *KiroAuthenticator) LoginWithGitHub(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
|
||||
if cfg == nil {
|
||||
return nil, fmt.Errorf("kiro auth: configuration is required")
|
||||
}
|
||||
|
||||
oauth := kiroauth.NewKiroOAuth(cfg)
|
||||
|
||||
// Use GitHub OAuth flow with protocol handler
|
||||
tokenData, err := oauth.LoginWithGitHub(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("github login failed: %w", err)
|
||||
}
|
||||
|
||||
// Parse expires_at
|
||||
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
|
||||
if err != nil {
|
||||
expiresAt = time.Now().Add(1 * time.Hour)
|
||||
}
|
||||
|
||||
// Extract identifier for file naming
|
||||
idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
|
||||
|
||||
now := time.Now()
|
||||
fileName := fmt.Sprintf("kiro-github-%s.json", idPart)
|
||||
|
||||
record := &coreauth.Auth{
|
||||
ID: fileName,
|
||||
Provider: "kiro",
|
||||
FileName: fileName,
|
||||
Label: "kiro-github",
|
||||
Status: coreauth.StatusActive,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
Metadata: map[string]any{
|
||||
"type": "kiro",
|
||||
"access_token": tokenData.AccessToken,
|
||||
"refresh_token": tokenData.RefreshToken,
|
||||
"profile_arn": tokenData.ProfileArn,
|
||||
"expires_at": tokenData.ExpiresAt,
|
||||
"auth_method": tokenData.AuthMethod,
|
||||
"provider": tokenData.Provider,
|
||||
"email": tokenData.Email,
|
||||
},
|
||||
Attributes: map[string]string{
|
||||
"profile_arn": tokenData.ProfileArn,
|
||||
"source": "github-oauth",
|
||||
"email": tokenData.Email,
|
||||
},
|
||||
NextRefreshAfter: expiresAt.Add(-30 * time.Minute),
|
||||
}
|
||||
|
||||
if tokenData.Email != "" {
|
||||
fmt.Printf("\n✓ Kiro GitHub authentication completed successfully! (Account: %s)\n", tokenData.Email)
|
||||
} else {
|
||||
fmt.Println("\n✓ Kiro GitHub authentication completed successfully!")
|
||||
}
|
||||
|
||||
return record, nil
|
||||
}
|
||||
|
||||
// ImportFromKiroIDE imports token from Kiro IDE's token file.
|
||||
func (a *KiroAuthenticator) ImportFromKiroIDE(ctx context.Context, cfg *config.Config) (*coreauth.Auth, error) {
|
||||
tokenData, err := kiroauth.LoadKiroIDEToken()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to load Kiro IDE token: %w", err)
|
||||
}
|
||||
|
||||
// Parse expires_at
|
||||
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
|
||||
if err != nil {
|
||||
expiresAt = time.Now().Add(1 * time.Hour)
|
||||
}
|
||||
|
||||
// Extract email from JWT if not already set (for imported tokens)
|
||||
if tokenData.Email == "" {
|
||||
tokenData.Email = kiroauth.ExtractEmailFromJWT(tokenData.AccessToken)
|
||||
}
|
||||
|
||||
// Extract identifier for file naming
|
||||
idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
|
||||
// Sanitize provider to prevent path traversal (defense-in-depth)
|
||||
provider := kiroauth.SanitizeEmailForFilename(strings.ToLower(strings.TrimSpace(tokenData.Provider)))
|
||||
if provider == "" {
|
||||
provider = "imported" // Fallback for legacy tokens without provider
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
fileName := fmt.Sprintf("kiro-%s-%s.json", provider, idPart)
|
||||
|
||||
record := &coreauth.Auth{
|
||||
ID: fileName,
|
||||
Provider: "kiro",
|
||||
FileName: fileName,
|
||||
Label: fmt.Sprintf("kiro-%s", provider),
|
||||
Status: coreauth.StatusActive,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
Metadata: map[string]any{
|
||||
"type": "kiro",
|
||||
"access_token": tokenData.AccessToken,
|
||||
"refresh_token": tokenData.RefreshToken,
|
||||
"profile_arn": tokenData.ProfileArn,
|
||||
"expires_at": tokenData.ExpiresAt,
|
||||
"auth_method": tokenData.AuthMethod,
|
||||
"provider": tokenData.Provider,
|
||||
"email": tokenData.Email,
|
||||
},
|
||||
Attributes: map[string]string{
|
||||
"profile_arn": tokenData.ProfileArn,
|
||||
"source": "kiro-ide-import",
|
||||
"email": tokenData.Email,
|
||||
},
|
||||
NextRefreshAfter: expiresAt.Add(-30 * time.Minute),
|
||||
}
|
||||
|
||||
// Display the email if extracted
|
||||
if tokenData.Email != "" {
|
||||
fmt.Printf("\n✓ Imported Kiro token from IDE (Provider: %s, Account: %s)\n", tokenData.Provider, tokenData.Email)
|
||||
} else {
|
||||
fmt.Printf("\n✓ Imported Kiro token from IDE (Provider: %s)\n", tokenData.Provider)
|
||||
}
|
||||
|
||||
return record, nil
|
||||
}
|
||||
|
||||
// Refresh refreshes an expired Kiro token using AWS SSO OIDC.
|
||||
func (a *KiroAuthenticator) Refresh(ctx context.Context, cfg *config.Config, auth *coreauth.Auth) (*coreauth.Auth, error) {
|
||||
if auth == nil || auth.Metadata == nil {
|
||||
return nil, fmt.Errorf("invalid auth record")
|
||||
}
|
||||
|
||||
refreshToken, ok := auth.Metadata["refresh_token"].(string)
|
||||
if !ok || refreshToken == "" {
|
||||
return nil, fmt.Errorf("refresh token not found")
|
||||
}
|
||||
|
||||
clientID, _ := auth.Metadata["client_id"].(string)
|
||||
clientSecret, _ := auth.Metadata["client_secret"].(string)
|
||||
authMethod, _ := auth.Metadata["auth_method"].(string)
|
||||
|
||||
var tokenData *kiroauth.KiroTokenData
|
||||
var err error
|
||||
|
||||
// Use SSO OIDC refresh for AWS Builder ID, otherwise use Kiro's OAuth refresh endpoint
|
||||
if clientID != "" && clientSecret != "" && authMethod == "builder-id" {
|
||||
ssoClient := kiroauth.NewSSOOIDCClient(cfg)
|
||||
tokenData, err = ssoClient.RefreshToken(ctx, clientID, clientSecret, refreshToken)
|
||||
} else {
|
||||
// Fallback to Kiro's refresh endpoint (for social auth: Google/GitHub)
|
||||
oauth := kiroauth.NewKiroOAuth(cfg)
|
||||
tokenData, err = oauth.RefreshToken(ctx, refreshToken)
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("token refresh failed: %w", err)
|
||||
}
|
||||
|
||||
// Parse expires_at
|
||||
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
|
||||
if err != nil {
|
||||
expiresAt = time.Now().Add(1 * time.Hour)
|
||||
}
|
||||
|
||||
// Clone auth to avoid mutating the input parameter
|
||||
updated := auth.Clone()
|
||||
now := time.Now()
|
||||
updated.UpdatedAt = now
|
||||
updated.LastRefreshedAt = now
|
||||
updated.Metadata["access_token"] = tokenData.AccessToken
|
||||
updated.Metadata["refresh_token"] = tokenData.RefreshToken
|
||||
updated.Metadata["expires_at"] = tokenData.ExpiresAt
|
||||
updated.Metadata["last_refresh"] = now.Format(time.RFC3339) // For double-check optimization
|
||||
updated.NextRefreshAfter = expiresAt.Add(-30 * time.Minute)
|
||||
|
||||
return updated, nil
|
||||
}
|
||||
@@ -74,3 +74,16 @@ func (m *Manager) Login(ctx context.Context, provider string, cfg *config.Config
|
||||
}
|
||||
return record, savedPath, nil
|
||||
}
|
||||
|
||||
// SaveAuth persists an auth record directly without going through the login flow.
|
||||
func (m *Manager) SaveAuth(record *coreauth.Auth, cfg *config.Config) (string, error) {
|
||||
if m.store == nil {
|
||||
return "", fmt.Errorf("no store configured")
|
||||
}
|
||||
if cfg != nil {
|
||||
if dirSetter, ok := m.store.(interface{ SetBaseDir(string) }); ok {
|
||||
dirSetter.SetBaseDir(cfg.AuthDir)
|
||||
}
|
||||
}
|
||||
return m.store.Save(context.Background(), record)
|
||||
}
|
||||
|
||||
@@ -14,6 +14,7 @@ func init() {
|
||||
registerRefreshLead("gemini", func() Authenticator { return NewGeminiAuthenticator() })
|
||||
registerRefreshLead("gemini-cli", func() Authenticator { return NewGeminiAuthenticator() })
|
||||
registerRefreshLead("antigravity", func() Authenticator { return NewAntigravityAuthenticator() })
|
||||
registerRefreshLead("kiro", func() Authenticator { return NewKiroAuthenticator() })
|
||||
}
|
||||
|
||||
func registerRefreshLead(provider string, factory func() Authenticator) {
|
||||
|
||||
@@ -379,6 +379,8 @@ func (s *Service) ensureExecutorsForAuth(a *coreauth.Auth) {
|
||||
s.coreManager.RegisterExecutor(executor.NewQwenExecutor(s.cfg))
|
||||
case "iflow":
|
||||
s.coreManager.RegisterExecutor(executor.NewIFlowExecutor(s.cfg))
|
||||
case "kiro":
|
||||
s.coreManager.RegisterExecutor(executor.NewKiroExecutor(s.cfg))
|
||||
default:
|
||||
providerKey := strings.ToLower(strings.TrimSpace(a.Provider))
|
||||
if providerKey == "" {
|
||||
@@ -721,6 +723,9 @@ func (s *Service) registerModelsForAuth(a *coreauth.Auth) {
|
||||
case "iflow":
|
||||
models = registry.GetIFlowModels()
|
||||
models = applyExcludedModels(models, excluded)
|
||||
case "kiro":
|
||||
models = registry.GetKiroModels()
|
||||
models = applyExcludedModels(models, excluded)
|
||||
default:
|
||||
// Handle OpenAI-compatibility providers by name using config
|
||||
if s.cfg != nil {
|
||||
|
||||
Reference in New Issue
Block a user