diff --git a/internal/auth/kiro/aws.go b/internal/auth/kiro/aws.go index 91f7f3c1..ef775d05 100644 --- a/internal/auth/kiro/aws.go +++ b/internal/auth/kiro/aws.go @@ -360,7 +360,7 @@ func SanitizeEmailForFilename(email string) string { } 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) @@ -378,7 +378,7 @@ func SanitizeEmailForFilename(email string) string { 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, "_") @@ -389,6 +389,65 @@ func SanitizeEmailForFilename(email string) string { parts[i] = part } result = strings.Join(parts, "_") - + return result } + +// ExtractIDCIdentifier extracts a unique identifier from IDC startUrl. +// Examples: +// - "https://d-1234567890.awsapps.com/start" -> "d-1234567890" +// - "https://my-company.awsapps.com/start" -> "my-company" +// - "https://acme-corp.awsapps.com/start" -> "acme-corp" +func ExtractIDCIdentifier(startURL string) string { + if startURL == "" { + return "" + } + + // Remove protocol prefix + url := strings.TrimPrefix(startURL, "https://") + url = strings.TrimPrefix(url, "http://") + + // Extract subdomain (first part before the first dot) + // Format: {identifier}.awsapps.com/start + parts := strings.Split(url, ".") + if len(parts) > 0 && parts[0] != "" { + identifier := parts[0] + // Sanitize for filename safety + identifier = strings.ReplaceAll(identifier, "/", "_") + identifier = strings.ReplaceAll(identifier, "\\", "_") + identifier = strings.ReplaceAll(identifier, ":", "_") + return identifier + } + + return "" +} + +// GenerateTokenFileName generates a unique filename for token storage. +// Priority: email > startUrl identifier (for IDC) > authMethod only +// Format: kiro-{authMethod}-{identifier}.json +func GenerateTokenFileName(tokenData *KiroTokenData) string { + authMethod := tokenData.AuthMethod + if authMethod == "" { + authMethod = "unknown" + } + + // Priority 1: Use email if available + if tokenData.Email != "" { + // Sanitize email for filename (replace @ and . with -) + sanitizedEmail := tokenData.Email + sanitizedEmail = strings.ReplaceAll(sanitizedEmail, "@", "-") + sanitizedEmail = strings.ReplaceAll(sanitizedEmail, ".", "-") + return fmt.Sprintf("kiro-%s-%s.json", authMethod, sanitizedEmail) + } + + // Priority 2: For IDC, use startUrl identifier + if authMethod == "idc" && tokenData.StartURL != "" { + identifier := ExtractIDCIdentifier(tokenData.StartURL) + if identifier != "" { + return fmt.Sprintf("kiro-%s-%s.json", authMethod, identifier) + } + } + + // Priority 3: Fallback to authMethod only + return fmt.Sprintf("kiro-%s.json", authMethod) +} diff --git a/internal/auth/kiro/aws_test.go b/internal/auth/kiro/aws_test.go index 5f60294c..194ad59e 100644 --- a/internal/auth/kiro/aws_test.go +++ b/internal/auth/kiro/aws_test.go @@ -151,11 +151,161 @@ func TestSanitizeEmailForFilename(t *testing.T) { // 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 } + +func TestExtractIDCIdentifier(t *testing.T) { + tests := []struct { + name string + startURL string + expected string + }{ + { + name: "Empty URL", + startURL: "", + expected: "", + }, + { + name: "Standard IDC URL with d- prefix", + startURL: "https://d-1234567890.awsapps.com/start", + expected: "d-1234567890", + }, + { + name: "IDC URL with company name", + startURL: "https://my-company.awsapps.com/start", + expected: "my-company", + }, + { + name: "IDC URL with simple name", + startURL: "https://acme-corp.awsapps.com/start", + expected: "acme-corp", + }, + { + name: "IDC URL without https", + startURL: "http://d-9876543210.awsapps.com/start", + expected: "d-9876543210", + }, + { + name: "IDC URL with subdomain only", + startURL: "https://test.awsapps.com/start", + expected: "test", + }, + { + name: "Builder ID URL", + startURL: "https://view.awsapps.com/start", + expected: "view", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := ExtractIDCIdentifier(tt.startURL) + if result != tt.expected { + t.Errorf("ExtractIDCIdentifier() = %q, want %q", result, tt.expected) + } + }) + } +} + +func TestGenerateTokenFileName(t *testing.T) { + tests := []struct { + name string + tokenData *KiroTokenData + expected string + }{ + { + name: "IDC with email", + tokenData: &KiroTokenData{ + AuthMethod: "idc", + Email: "user@example.com", + StartURL: "https://d-1234567890.awsapps.com/start", + }, + expected: "kiro-idc-user-example-com.json", + }, + { + name: "IDC without email but with startUrl", + tokenData: &KiroTokenData{ + AuthMethod: "idc", + Email: "", + StartURL: "https://d-1234567890.awsapps.com/start", + }, + expected: "kiro-idc-d-1234567890.json", + }, + { + name: "IDC with company name in startUrl", + tokenData: &KiroTokenData{ + AuthMethod: "idc", + Email: "", + StartURL: "https://my-company.awsapps.com/start", + }, + expected: "kiro-idc-my-company.json", + }, + { + name: "IDC without email and without startUrl", + tokenData: &KiroTokenData{ + AuthMethod: "idc", + Email: "", + StartURL: "", + }, + expected: "kiro-idc.json", + }, + { + name: "Builder ID with email", + tokenData: &KiroTokenData{ + AuthMethod: "builder-id", + Email: "user@gmail.com", + StartURL: "https://view.awsapps.com/start", + }, + expected: "kiro-builder-id-user-gmail-com.json", + }, + { + name: "Builder ID without email", + tokenData: &KiroTokenData{ + AuthMethod: "builder-id", + Email: "", + StartURL: "https://view.awsapps.com/start", + }, + expected: "kiro-builder-id.json", + }, + { + name: "Social auth with email", + tokenData: &KiroTokenData{ + AuthMethod: "google", + Email: "user@gmail.com", + }, + expected: "kiro-google-user-gmail-com.json", + }, + { + name: "Empty auth method", + tokenData: &KiroTokenData{ + AuthMethod: "", + Email: "", + }, + expected: "kiro-unknown.json", + }, + { + name: "Email with special characters", + tokenData: &KiroTokenData{ + AuthMethod: "idc", + Email: "user.name+tag@sub.example.com", + StartURL: "https://d-1234567890.awsapps.com/start", + }, + expected: "kiro-idc-user-name+tag-sub-example-com.json", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := GenerateTokenFileName(tt.tokenData) + if result != tt.expected { + t.Errorf("GenerateTokenFileName() = %q, want %q", result, tt.expected) + } + }) + } +} diff --git a/internal/auth/kiro/oauth_web.go b/internal/auth/kiro/oauth_web.go index 6e4269c5..88fba672 100644 --- a/internal/auth/kiro/oauth_web.go +++ b/internal/auth/kiro/oauth_web.go @@ -421,7 +421,7 @@ func (h *OAuthWebHandler) saveTokenToFile(tokenData *KiroTokenData) { log.Errorf("OAuth Web: failed to resolve auth directory: %v", err) } } - + // Fall back to default location if authDir == "" { home, err := os.UserHomeDir() @@ -431,24 +431,16 @@ func (h *OAuthWebHandler) saveTokenToFile(tokenData *KiroTokenData) { } authDir = filepath.Join(home, ".cli-proxy-api") } - + // Create directory if not exists if err := os.MkdirAll(authDir, 0700); err != nil { log.Errorf("OAuth Web: failed to create auth directory: %v", err) return } - - // Generate filename based on auth method - // Format: kiro-{authMethod}.json or kiro-{authMethod}-{email}.json - fileName := fmt.Sprintf("kiro-%s.json", tokenData.AuthMethod) - if tokenData.Email != "" { - // Sanitize email for filename (replace @ and . with -) - sanitizedEmail := tokenData.Email - sanitizedEmail = strings.ReplaceAll(sanitizedEmail, "@", "-") - sanitizedEmail = strings.ReplaceAll(sanitizedEmail, ".", "-") - fileName = fmt.Sprintf("kiro-%s-%s.json", tokenData.AuthMethod, sanitizedEmail) - } - + + // Generate filename using the unified function + fileName := GenerateTokenFileName(tokenData) + authFilePath := filepath.Join(authDir, fileName) // Convert to storage format and save @@ -811,13 +803,8 @@ func (h *OAuthWebHandler) handleImportToken(c *gin.Context) { // Save token to file h.saveTokenToFile(tokenData) - // Generate filename for response - fileName := fmt.Sprintf("kiro-%s.json", tokenData.AuthMethod) - if tokenData.Email != "" { - sanitizedEmail := strings.ReplaceAll(tokenData.Email, "@", "-") - sanitizedEmail = strings.ReplaceAll(sanitizedEmail, ".", "-") - fileName = fmt.Sprintf("kiro-%s-%s.json", tokenData.AuthMethod, sanitizedEmail) - } + // Generate filename for response using the unified function + fileName := GenerateTokenFileName(tokenData) log.Infof("OAuth Web: token imported successfully") c.JSON(http.StatusOK, gin.H{