feat(kiro): 新增 AWS Builder ID 授权码流程认证及用户邮箱获取增强

Amp-Thread-ID: https://ampcode.com/threads/T-019b2ecc-fb2d-713f-b30d-1196c7dce3e2
Co-authored-by: Amp <amp@ampcode.com>
This commit is contained in:
Ravens2121
2025-12-18 08:16:36 +08:00
parent 54acd69e9d
commit cf9a246d53
6 changed files with 695 additions and 7 deletions

View File

@@ -78,6 +78,7 @@ func main() {
var kiroLogin bool
var kiroGoogleLogin bool
var kiroAWSLogin bool
var kiroAWSAuthCode bool
var kiroImport bool
var githubCopilotLogin bool
var projectID string
@@ -101,6 +102,7 @@ func main() {
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(&kiroAWSAuthCode, "kiro-aws-authcode", false, "Login to Kiro using AWS Builder ID (authorization code flow, better UX)")
flag.BoolVar(&kiroImport, "kiro-import", false, "Import Kiro token from Kiro IDE (~/.aws/sso/cache/kiro-auth-token.json)")
flag.BoolVar(&githubCopilotLogin, "github-copilot-login", false, "Login to GitHub Copilot using device flow")
flag.StringVar(&projectID, "project_id", "", "Project ID (Gemini only, not required)")
@@ -513,6 +515,10 @@ func main() {
// Users can explicitly override with --no-incognito
setKiroIncognitoMode(cfg, useIncognito, noIncognito)
cmd.DoKiroAWSLogin(cfg, options)
} else if kiroAWSAuthCode {
// For Kiro auth with authorization code flow (better UX)
setKiroIncognitoMode(cfg, useIncognito, noIncognito)
cmd.DoKiroAWSAuthCodeLogin(cfg, options)
} else if kiroImport {
cmd.DoKiroImport(cfg, options)
} else {

View File

@@ -0,0 +1,166 @@
// Package kiro provides CodeWhisperer API client for fetching user info.
package kiro
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"time"
"github.com/google/uuid"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
log "github.com/sirupsen/logrus"
)
const (
codeWhispererAPI = "https://codewhisperer.us-east-1.amazonaws.com"
kiroVersion = "0.6.18"
)
// CodeWhispererClient handles CodeWhisperer API calls.
type CodeWhispererClient struct {
httpClient *http.Client
machineID string
}
// UsageLimitsResponse represents the getUsageLimits API response.
type UsageLimitsResponse struct {
DaysUntilReset *int `json:"daysUntilReset,omitempty"`
NextDateReset *float64 `json:"nextDateReset,omitempty"`
UserInfo *UserInfo `json:"userInfo,omitempty"`
SubscriptionInfo *SubscriptionInfo `json:"subscriptionInfo,omitempty"`
UsageBreakdownList []UsageBreakdown `json:"usageBreakdownList,omitempty"`
}
// UserInfo contains user information from the API.
type UserInfo struct {
Email string `json:"email,omitempty"`
UserID string `json:"userId,omitempty"`
}
// SubscriptionInfo contains subscription details.
type SubscriptionInfo struct {
SubscriptionTitle string `json:"subscriptionTitle,omitempty"`
Type string `json:"type,omitempty"`
}
// UsageBreakdown contains usage details.
type UsageBreakdown struct {
UsageLimit *int `json:"usageLimit,omitempty"`
CurrentUsage *int `json:"currentUsage,omitempty"`
UsageLimitWithPrecision *float64 `json:"usageLimitWithPrecision,omitempty"`
CurrentUsageWithPrecision *float64 `json:"currentUsageWithPrecision,omitempty"`
NextDateReset *float64 `json:"nextDateReset,omitempty"`
DisplayName string `json:"displayName,omitempty"`
ResourceType string `json:"resourceType,omitempty"`
}
// NewCodeWhispererClient creates a new CodeWhisperer client.
func NewCodeWhispererClient(cfg *config.Config, machineID string) *CodeWhispererClient {
client := &http.Client{Timeout: 30 * time.Second}
if cfg != nil {
client = util.SetProxy(&cfg.SDKConfig, client)
}
if machineID == "" {
machineID = uuid.New().String()
}
return &CodeWhispererClient{
httpClient: client,
machineID: machineID,
}
}
// generateInvocationID generates a unique invocation ID.
func generateInvocationID() string {
return uuid.New().String()
}
// GetUsageLimits fetches usage limits and user info from CodeWhisperer API.
// This is the recommended way to get user email after login.
func (c *CodeWhispererClient) GetUsageLimits(ctx context.Context, accessToken string) (*UsageLimitsResponse, error) {
url := fmt.Sprintf("%s/getUsageLimits?isEmailRequired=true&origin=AI_EDITOR&resourceType=AGENTIC_REQUEST", codeWhispererAPI)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
// Set headers to match Kiro IDE
xAmzUserAgent := fmt.Sprintf("aws-sdk-js/1.0.0 KiroIDE-%s-%s", kiroVersion, c.machineID)
userAgent := fmt.Sprintf("aws-sdk-js/1.0.0 ua/2.1 os/windows lang/js md/nodejs#20.16.0 api/codewhispererruntime#1.0.0 m/E KiroIDE-%s-%s", kiroVersion, c.machineID)
req.Header.Set("Authorization", "Bearer "+accessToken)
req.Header.Set("x-amz-user-agent", xAmzUserAgent)
req.Header.Set("User-Agent", userAgent)
req.Header.Set("amz-sdk-invocation-id", generateInvocationID())
req.Header.Set("amz-sdk-request", "attempt=1; max=1")
req.Header.Set("Connection", "close")
log.Debugf("codewhisperer: GET %s", url)
resp, err := c.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("request failed: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
log.Debugf("codewhisperer: status=%d, body=%s", resp.StatusCode, string(body))
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API returned status %d: %s", resp.StatusCode, string(body))
}
var result UsageLimitsResponse
if err := json.Unmarshal(body, &result); err != nil {
return nil, fmt.Errorf("failed to parse response: %w", err)
}
return &result, nil
}
// FetchUserEmailFromAPI fetches user email using CodeWhisperer getUsageLimits API.
// This is more reliable than JWT parsing as it uses the official API.
func (c *CodeWhispererClient) FetchUserEmailFromAPI(ctx context.Context, accessToken string) string {
resp, err := c.GetUsageLimits(ctx, accessToken)
if err != nil {
log.Debugf("codewhisperer: failed to get usage limits: %v", err)
return ""
}
if resp.UserInfo != nil && resp.UserInfo.Email != "" {
log.Debugf("codewhisperer: got email from API: %s", resp.UserInfo.Email)
return resp.UserInfo.Email
}
log.Debugf("codewhisperer: no email in response")
return ""
}
// FetchUserEmailWithFallback fetches user email with multiple fallback methods.
// Priority: 1. CodeWhisperer API 2. userinfo endpoint 3. JWT parsing
func FetchUserEmailWithFallback(ctx context.Context, cfg *config.Config, accessToken string) string {
// Method 1: Try CodeWhisperer API (most reliable)
cwClient := NewCodeWhispererClient(cfg, "")
email := cwClient.FetchUserEmailFromAPI(ctx, accessToken)
if email != "" {
return email
}
// Method 2: Try SSO OIDC userinfo endpoint
ssoClient := NewSSOOIDCClient(cfg)
email = ssoClient.FetchUserEmail(ctx, accessToken)
if email != "" {
return email
}
// Method 3: Fallback to JWT parsing
return ExtractEmailFromJWT(accessToken)
}

View File

@@ -163,6 +163,13 @@ func (o *KiroOAuth) LoginWithBuilderID(ctx context.Context) (*KiroTokenData, err
return ssoClient.LoginWithBuilderID(ctx)
}
// LoginWithBuilderIDAuthCode performs OAuth login with AWS Builder ID using authorization code flow.
// This provides a better UX than device code flow as it uses automatic browser callback.
func (o *KiroOAuth) LoginWithBuilderIDAuthCode(ctx context.Context) (*KiroTokenData, error) {
ssoClient := NewSSOOIDCClient(o.cfg)
return ssoClient.LoginWithBuilderIDAuthCode(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{

View File

@@ -3,9 +3,14 @@ package kiro
import (
"context"
"crypto/rand"
"crypto/sha256"
"encoding/base64"
"encoding/json"
"fmt"
"html"
"io"
"net"
"net/http"
"strings"
"time"
@@ -25,6 +30,13 @@ const (
// Polling interval
pollInterval = 5 * time.Second
// Authorization code flow callback
authCodeCallbackPath = "/oauth/callback"
authCodeCallbackPort = 19877
// User-Agent to match official Kiro IDE
kiroUserAgent = "KiroIDE"
)
// SSOOIDCClient handles AWS SSO OIDC authentication.
@@ -73,13 +85,11 @@ type CreateTokenResponse struct {
// 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,
"clientName": "Kiro IDE",
"clientType": "public",
"scopes": []string{"codewhisperer:completions", "codewhisperer:analysis", "codewhisperer:conversations"},
"scopes": []string{"codewhisperer:completions", "codewhisperer:analysis", "codewhisperer:conversations", "codewhisperer:transformations", "codewhisperer:taskassist"},
"grantTypes": []string{"urn:ietf:params:oauth:grant-type:device_code", "refresh_token"},
}
body, err := json.Marshal(payload)
@@ -92,6 +102,7 @@ func (c *SSOOIDCClient) RegisterClient(ctx context.Context) (*RegisterClientResp
return nil, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", kiroUserAgent)
resp, err := c.httpClient.Do(req)
if err != nil {
@@ -135,6 +146,7 @@ func (c *SSOOIDCClient) StartDeviceAuthorization(ctx context.Context, clientID,
return nil, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", kiroUserAgent)
resp, err := c.httpClient.Do(req)
if err != nil {
@@ -179,6 +191,7 @@ func (c *SSOOIDCClient) CreateToken(ctx context.Context, clientID, clientSecret,
return nil, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", kiroUserAgent)
resp, err := c.httpClient.Do(req)
if err != nil {
@@ -240,6 +253,7 @@ func (c *SSOOIDCClient) RefreshToken(ctx context.Context, clientID, clientSecret
return nil, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", kiroUserAgent)
resp, err := c.httpClient.Do(req)
if err != nil {
@@ -370,8 +384,8 @@ func (c *SSOOIDCClient) LoginWithBuilderID(ctx context.Context) (*KiroTokenData,
fmt.Println("Fetching profile information...")
profileArn := c.fetchProfileArn(ctx, tokenResp.AccessToken)
// Extract email from JWT access token
email := ExtractEmailFromJWT(tokenResp.AccessToken)
// Fetch user email (tries CodeWhisperer API first, then userinfo endpoint, then JWT parsing)
email := FetchUserEmailWithFallback(ctx, c.cfg, tokenResp.AccessToken)
if email != "" {
fmt.Printf(" Logged in as: %s\n", email)
}
@@ -399,6 +413,68 @@ func (c *SSOOIDCClient) LoginWithBuilderID(ctx context.Context) (*KiroTokenData,
return nil, fmt.Errorf("authorization timed out")
}
// FetchUserEmail retrieves the user's email from AWS SSO OIDC userinfo endpoint.
// Falls back to JWT parsing if userinfo fails.
func (c *SSOOIDCClient) FetchUserEmail(ctx context.Context, accessToken string) string {
// Method 1: Try userinfo endpoint (standard OIDC)
email := c.tryUserInfoEndpoint(ctx, accessToken)
if email != "" {
return email
}
// Method 2: Fallback to JWT parsing
return ExtractEmailFromJWT(accessToken)
}
// tryUserInfoEndpoint attempts to get user info from AWS SSO OIDC userinfo endpoint.
func (c *SSOOIDCClient) tryUserInfoEndpoint(ctx context.Context, accessToken string) string {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, ssoOIDCEndpoint+"/userinfo", nil)
if err != nil {
return ""
}
req.Header.Set("Authorization", "Bearer "+accessToken)
req.Header.Set("Accept", "application/json")
resp, err := c.httpClient.Do(req)
if err != nil {
log.Debugf("userinfo request failed: %v", err)
return ""
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
respBody, _ := io.ReadAll(resp.Body)
log.Debugf("userinfo endpoint returned status %d: %s", resp.StatusCode, string(respBody))
return ""
}
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return ""
}
log.Debugf("userinfo response: %s", string(respBody))
var userInfo struct {
Email string `json:"email"`
Sub string `json:"sub"`
PreferredUsername string `json:"preferred_username"`
Name string `json:"name"`
}
if err := json.Unmarshal(respBody, &userInfo); err != nil {
return ""
}
if userInfo.Email != "" {
return userInfo.Email
}
if userInfo.PreferredUsername != "" && strings.Contains(userInfo.PreferredUsername, "@") {
return userInfo.PreferredUsername
}
return ""
}
// 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 {
@@ -525,3 +601,323 @@ func (c *SSOOIDCClient) tryListCustomizations(ctx context.Context, accessToken s
return ""
}
// RegisterClientForAuthCode registers a new OIDC client for authorization code flow.
func (c *SSOOIDCClient) RegisterClientForAuthCode(ctx context.Context, redirectURI string) (*RegisterClientResponse, error) {
payload := map[string]interface{}{
"clientName": "Kiro IDE",
"clientType": "public",
"scopes": []string{"codewhisperer:completions", "codewhisperer:analysis", "codewhisperer:conversations", "codewhisperer:transformations", "codewhisperer:taskassist"},
"grantTypes": []string{"authorization_code", "refresh_token"},
"redirectUris": []string{redirectURI},
"issuerUrl": builderIDStartURL,
}
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")
req.Header.Set("User-Agent", kiroUserAgent)
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 for auth code 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
}
// AuthCodeCallbackResult contains the result from authorization code callback.
type AuthCodeCallbackResult struct {
Code string
State string
Error string
}
// startAuthCodeCallbackServer starts a local HTTP server to receive the authorization code callback.
func (c *SSOOIDCClient) startAuthCodeCallbackServer(ctx context.Context, expectedState string) (string, <-chan AuthCodeCallbackResult, error) {
// Try to find an available port
listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", authCodeCallbackPort))
if err != nil {
// Try with dynamic port
log.Warnf("sso oidc: default port %d is busy, falling back to dynamic port", authCodeCallbackPort)
listener, err = net.Listen("tcp", "127.0.0.1:0")
if err != nil {
return "", nil, fmt.Errorf("failed to start callback server: %w", err)
}
}
port := listener.Addr().(*net.TCPAddr).Port
redirectURI := fmt.Sprintf("http://127.0.0.1:%d%s", port, authCodeCallbackPath)
resultChan := make(chan AuthCodeCallbackResult, 1)
server := &http.Server{
ReadHeaderTimeout: 10 * time.Second,
}
mux := http.NewServeMux()
mux.HandleFunc(authCodeCallbackPath, 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")
// Send response to browser
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))
resultChan <- AuthCodeCallbackResult{Error: errParam}
return
}
if state != expectedState {
w.WriteHeader(http.StatusBadRequest)
fmt.Fprint(w, `<!DOCTYPE html>
<html><head><title>Login Failed</title></head>
<body><h1>Login Failed</h1><p>Invalid state parameter</p><p>You can close this window.</p></body></html>`)
resultChan <- AuthCodeCallbackResult{Error: "state mismatch"}
return
}
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>`)
resultChan <- AuthCodeCallbackResult{Code: code, State: state}
})
server.Handler = mux
go func() {
if err := server.Serve(listener); err != nil && err != http.ErrServerClosed {
log.Debugf("auth code callback server error: %v", err)
}
}()
go func() {
select {
case <-ctx.Done():
case <-time.After(10 * time.Minute):
case <-resultChan:
}
_ = server.Shutdown(context.Background())
}()
return redirectURI, resultChan, nil
}
// generatePKCEForAuthCode generates PKCE code verifier and challenge for authorization code flow.
func generatePKCEForAuthCode() (verifier, challenge string, err error) {
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)
h := sha256.Sum256([]byte(verifier))
challenge = base64.RawURLEncoding.EncodeToString(h[:])
return verifier, challenge, nil
}
// generateStateForAuthCode generates a random state parameter.
func generateStateForAuthCode() (string, error) {
b := make([]byte, 16)
if _, err := rand.Read(b); err != nil {
return "", err
}
return base64.RawURLEncoding.EncodeToString(b), nil
}
// CreateTokenWithAuthCode exchanges authorization code for tokens.
func (c *SSOOIDCClient) CreateTokenWithAuthCode(ctx context.Context, clientID, clientSecret, code, codeVerifier, redirectURI string) (*CreateTokenResponse, error) {
payload := map[string]string{
"clientId": clientID,
"clientSecret": clientSecret,
"code": code,
"codeVerifier": codeVerifier,
"redirectUri": redirectURI,
"grantType": "authorization_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")
req.Header.Set("User-Agent", kiroUserAgent)
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("create token with auth code 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
}
// LoginWithBuilderIDAuthCode performs the authorization code flow for AWS Builder ID.
// This provides a better UX than device code flow as it uses automatic browser callback.
func (c *SSOOIDCClient) LoginWithBuilderIDAuthCode(ctx context.Context) (*KiroTokenData, error) {
fmt.Println("\n╔══════════════════════════════════════════════════════════╗")
fmt.Println("║ Kiro Authentication (AWS Builder ID - Auth Code) ║")
fmt.Println("╚══════════════════════════════════════════════════════════╝")
// Step 1: Generate PKCE and state
codeVerifier, codeChallenge, err := generatePKCEForAuthCode()
if err != nil {
return nil, fmt.Errorf("failed to generate PKCE: %w", err)
}
state, err := generateStateForAuthCode()
if err != nil {
return nil, fmt.Errorf("failed to generate state: %w", err)
}
// Step 2: Start callback server
fmt.Println("\nStarting callback server...")
redirectURI, resultChan, err := c.startAuthCodeCallbackServer(ctx, state)
if err != nil {
return nil, fmt.Errorf("failed to start callback server: %w", err)
}
log.Debugf("Callback server started, redirect URI: %s", redirectURI)
// Step 3: Register client with auth code grant type
fmt.Println("Registering client...")
regResp, err := c.RegisterClientForAuthCode(ctx, redirectURI)
if err != nil {
return nil, fmt.Errorf("failed to register client: %w", err)
}
log.Debugf("Client registered: %s", regResp.ClientID)
// Step 4: Build authorization URL
scopes := "codewhisperer:completions,codewhisperer:analysis,codewhisperer:conversations"
authURL := fmt.Sprintf("%s/authorize?response_type=code&client_id=%s&redirect_uri=%s&scopes=%s&state=%s&code_challenge=%s&code_challenge_method=S256",
ssoOIDCEndpoint,
regResp.ClientID,
redirectURI,
scopes,
state,
codeChallenge,
)
// Step 5: Open browser
fmt.Println("\n════════════════════════════════════════════════════════════")
fmt.Println(" Opening browser for authentication...")
fmt.Println("════════════════════════════════════════════════════════════")
fmt.Printf("\n URL: %s\n\n", authURL)
// Set incognito mode
if c.cfg != nil {
browser.SetIncognitoMode(c.cfg.IncognitoBrowser)
} else {
browser.SetIncognitoMode(true)
}
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 authorization callback...")
// Step 6: Wait for callback
select {
case <-ctx.Done():
browser.CloseBrowser()
return nil, ctx.Err()
case <-time.After(10 * time.Minute):
browser.CloseBrowser()
return nil, fmt.Errorf("authorization timed out")
case result := <-resultChan:
if result.Error != "" {
browser.CloseBrowser()
return nil, fmt.Errorf("authorization failed: %s", result.Error)
}
fmt.Println("\n✓ Authorization received!")
// Close browser
if err := browser.CloseBrowser(); err != nil {
log.Debugf("Failed to close browser: %v", err)
}
// Step 7: Exchange code for tokens
fmt.Println("Exchanging code for tokens...")
tokenResp, err := c.CreateTokenWithAuthCode(ctx, regResp.ClientID, regResp.ClientSecret, result.Code, codeVerifier, redirectURI)
if err != nil {
return nil, fmt.Errorf("failed to exchange code for tokens: %w", err)
}
fmt.Println("\n✓ Authentication successful!")
// Step 8: Get profile ARN
fmt.Println("Fetching profile information...")
profileArn := c.fetchProfileArn(ctx, tokenResp.AccessToken)
// Fetch user email (tries CodeWhisperer API first, then userinfo endpoint, then JWT parsing)
email := FetchUserEmailWithFallback(ctx, c.cfg, 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
}
}

View File

@@ -116,6 +116,54 @@ func DoKiroAWSLogin(cfg *config.Config, options *LoginOptions) {
fmt.Println("Kiro AWS authentication successful!")
}
// DoKiroAWSAuthCodeLogin triggers Kiro authentication with AWS Builder ID using authorization code flow.
// This provides a better UX than device code flow as it uses automatic browser callback.
//
// Parameters:
// - cfg: The application configuration
// - options: Login options including prompts
func DoKiroAWSAuthCodeLogin(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 (authorization code flow)
authenticator := sdkAuth.NewKiroAuthenticator()
record, err := authenticator.LoginWithAuthCode(context.Background(), cfg, &sdkAuth.LoginOptions{
NoBrowser: options.NoBrowser,
Metadata: map[string]string{},
Prompt: options.Prompt,
})
if err != nil {
log.Errorf("Kiro AWS authentication (auth code) 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-aws-login (device code flow)")
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.

View File

@@ -117,6 +117,71 @@ func (a *KiroAuthenticator) Login(ctx context.Context, cfg *config.Config, opts
return record, nil
}
// LoginWithAuthCode performs OAuth login for Kiro with AWS Builder ID using authorization code flow.
// This provides a better UX than device code flow as it uses automatic browser callback.
func (a *KiroAuthenticator) LoginWithAuthCode(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 authorization code flow
tokenData, err := oauth.LoginWithBuilderIDAuthCode(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-authcode",
"email": tokenData.Email,
},
// NextRefreshAfter is aligned with RefreshLead (5min)
NextRefreshAfter: expiresAt.Add(-5 * 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) {