From 39fbcdc1af6322dda9312cf4a1d640cca331ac70 Mon Sep 17 00:00:00 2001
From: Dawid Rusnak <dawid@drcode.pl>
Date: Wed, 13 Nov 2024 10:52:44 +0100
Subject: [PATCH] fix: CLI login with custom auth URL (#6024)

* fix: CLI login with custom auth URL
* fix: CLI login with custom agent URL
* fix: typo in README
---
 cmd/kubectl-testkube/commands/common/client.go |  5 ++++-
 cmd/kubectl-testkube/commands/common/helper.go | 14 ++++++++++++--
 cmd/kubectl-testkube/config/data.go            |  1 +
 cmd/tcl/kubectl-testkube/devbox/README.md      |  2 +-
 4 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/cmd/kubectl-testkube/commands/common/client.go b/cmd/kubectl-testkube/commands/common/client.go
index 06bf99ea3a..bfc11858e9 100644
--- a/cmd/kubectl-testkube/commands/common/client.go
+++ b/cmd/kubectl-testkube/commands/common/client.go
@@ -69,7 +69,10 @@ func GetClient(cmd *cobra.Command) (client.Client, string, error) {
 
 		if cfg.CloudContext.ApiKey != "" && cfg.CloudContext.RefreshToken != "" {
 			var refreshToken string
-			authURI := fmt.Sprintf("%s/idp", cfg.CloudContext.ApiUri)
+			authURI := cfg.CloudContext.AuthUri
+			if cfg.CloudContext.AuthUri == "" {
+				authURI = fmt.Sprintf("%s/idp", cfg.CloudContext.ApiUri)
+			}
 			token, refreshToken, err = cloudlogin.CheckAndRefreshToken(context.Background(), authURI, cfg.CloudContext.ApiKey, cfg.CloudContext.RefreshToken)
 			if err != nil {
 				// Error: failed refreshing, go thru login flow
diff --git a/cmd/kubectl-testkube/commands/common/helper.go b/cmd/kubectl-testkube/commands/common/helper.go
index dc773ef9ff..4e9e8e58f7 100644
--- a/cmd/kubectl-testkube/commands/common/helper.go
+++ b/cmd/kubectl-testkube/commands/common/helper.go
@@ -302,14 +302,20 @@ func PopulateLoginDataToContext(orgID, envID, token, refreshToken, dockerContain
 	if options.Master.AgentToken != "" {
 		cfg.CloudContext.AgentKey = options.Master.AgentToken
 	}
-	if options.Master.URIs.Api != "" {
-		cfg.CloudContext.AgentUri = options.Master.URIs.Api
+	if options.Master.URIs.Agent != "" {
+		cfg.CloudContext.AgentUri = options.Master.URIs.Agent
 	}
 	if options.Master.URIs.Ui != "" {
 		cfg.CloudContext.UiUri = options.Master.URIs.Ui
 	}
 	if options.Master.URIs.Api != "" {
 		cfg.CloudContext.ApiUri = options.Master.URIs.Api
+		if options.Master.URIs.Agent == "" {
+			cfg.CloudContext.AgentUri = options.Master.URIs.Api
+		}
+	}
+	if options.Master.URIs.Auth != "" {
+		cfg.CloudContext.AuthUri = options.Master.URIs.Auth
 	}
 	cfg.ContextType = config.ContextTypeCloud
 	cfg.CloudContext.OrganizationId = orgID
@@ -349,6 +355,10 @@ func PopulateAgentDataToContext(options HelmOptions, cfg config.Data) error {
 		cfg.CloudContext.ApiUri = options.Master.URIs.Api
 		updated = true
 	}
+	if options.Master.URIs.Auth != "" {
+		cfg.CloudContext.AuthUri = options.Master.URIs.Auth
+		updated = true
+	}
 	if options.Master.IdToken != "" {
 		cfg.CloudContext.ApiKey = options.Master.IdToken
 		updated = true
diff --git a/cmd/kubectl-testkube/config/data.go b/cmd/kubectl-testkube/config/data.go
index 7f58b7e908..d4dbc6801f 100644
--- a/cmd/kubectl-testkube/config/data.go
+++ b/cmd/kubectl-testkube/config/data.go
@@ -22,6 +22,7 @@ type CloudContext struct {
 	AgentUri            string `json:"agentUri,omitempty"`
 	RootDomain          string `json:"rootDomain,omitempty"`
 	UiUri               string `json:"uiUri,omitempty"`
+	AuthUri             string `json:"authUri,omitempty"`
 	TokenType           string `json:"tokenType,omitempty"`
 	DockerContainerName string `json:"dockerContainerName,omitempty"`
 	CustomAuth          bool   `json:"customConnector,omitempty"`
diff --git a/cmd/tcl/kubectl-testkube/devbox/README.md b/cmd/tcl/kubectl-testkube/devbox/README.md
index 75d7556f18..5502198018 100644
--- a/cmd/tcl/kubectl-testkube/devbox/README.md
+++ b/cmd/tcl/kubectl-testkube/devbox/README.md
@@ -18,7 +18,7 @@ This utility is used to help with development of the Agent features (like Test W
 ## Usage
 
 * Login to Testkube CLI, like `testkube login`
-  * For local development Testkube Enterprise (Skaffold), consider `testkube login --api-uri-override=http://localhost:8099 --agent-uri-override=http://testkube-enterprise-api.tk-dev.svc.local:8089 --auth-uri-override=http://localhost:5556 --custom-auth`
+  * For local development Testkube Enterprise (Skaffold), consider `testkube login --api-uri-override=http://localhost:8099 --agent-uri-override=http://testkube-enterprise-api.tk-dev.svc.cluster.local:8089 --auth-uri-override=http://localhost:5556 --custom-auth`
   * It's worth to create alias for that in own `.bashrc` or `.bash_profile`
   * It's worth to pass a devbox name, like `-n dawid`, so it's not using random name
 * For OSS version - run with `--oss` parameter