From cd74df4698bf73b92c47d9d92b804c7d9342c616 Mon Sep 17 00:00:00 2001
From: Oscar Reyes <oscar-rreyes1@hotmail.com>
Date: Tue, 21 May 2024 12:27:17 -0600
Subject: [PATCH] fix(cli): agent connection (#3861)

* fix(cli): agent connection

* fix(cli): agent connection
---
 agent/runner/agent.go       | 56 +++++++++++++++++++++++++++++++++++++
 agent/runner/environment.go | 48 -------------------------------
 agent/runner/runner.go      |  6 ++--
 3 files changed, 59 insertions(+), 51 deletions(-)
 create mode 100644 agent/runner/agent.go
 delete mode 100644 agent/runner/environment.go

diff --git a/agent/runner/agent.go b/agent/runner/agent.go
new file mode 100644
index 0000000000..325687f70d
--- /dev/null
+++ b/agent/runner/agent.go
@@ -0,0 +1,56 @@
+package runner
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+
+	"go.uber.org/zap"
+)
+
+type agentResponse struct {
+	Key string `json:"key"`
+}
+
+func (s *Runner) getAgentToken(ctx context.Context, endpoint, orgID, envID, token string) (string, error) {
+	s.logger.Debug("Getting agent token", zap.String("endpoint", endpoint), zap.String("token", token))
+	req, err := http.NewRequestWithContext(ctx, "GET", fmt.Sprintf("%s/organizations/%s/environments/%s/agent", endpoint, orgID, envID), nil)
+	if err != nil {
+		s.logger.Debug("Failed to create request", zap.Error(err))
+		return "", fmt.Errorf("failed to create request: %w", err)
+	}
+
+	headers := req.Header
+	headers.Set("Authorization", fmt.Sprintf("Bearer %s", token))
+
+	res, err := http.DefaultClient.Do(req)
+	if err != nil {
+		s.logger.Debug("Failed to get agent token", zap.Error(err))
+		return "", fmt.Errorf("failed to get agent token: %w", err)
+	}
+
+	if res.StatusCode != http.StatusOK {
+		s.logger.Debug("Failed to get agent token", zap.String("status", res.Status))
+		return "", fmt.Errorf("failed to get agent token: %s", res.Status)
+	}
+
+	defer res.Body.Close()
+	body, err := io.ReadAll(res.Body)
+	if err != nil {
+		s.logger.Debug("Failed to read response body", zap.Error(err))
+		return "", fmt.Errorf("failed to read response body: %w", err)
+	}
+
+	var agentResponse agentResponse
+	err = json.Unmarshal(body, &agentResponse)
+	if err != nil {
+		s.logger.Debug("Failed to unmarshal response body", zap.Error(err))
+		return "", fmt.Errorf("failed to unmarshal response body: %w", err)
+	}
+
+	s.logger.Debug("Agent token received")
+
+	return agentResponse.Key, nil
+}
diff --git a/agent/runner/environment.go b/agent/runner/environment.go
deleted file mode 100644
index 9b7857f7cd..0000000000
--- a/agent/runner/environment.go
+++ /dev/null
@@ -1,48 +0,0 @@
-package runner
-
-import (
-	"context"
-	"encoding/json"
-	"fmt"
-
-	"github.com/kubeshop/tracetest/cli/config"
-	"github.com/kubeshop/tracetest/cli/pkg/resourcemanager"
-)
-
-type environment struct {
-	ID             string `json:"id"`
-	Name           string `json:"name"`
-	AgentApiKey    string `json:"agentApiKey"`
-	OrganizationID string `json:"organizationID"`
-}
-
-func (s *Runner) getEnvironment(ctx context.Context, cfg config.Config) (environment, error) {
-	resource, err := s.resources.Get("env")
-	if err != nil {
-		return environment{}, err
-	}
-
-	resource = resource.
-		WithHttpClient(config.SetupHttpClient(cfg)).
-		WithOptions(resourcemanager.WithPrefixGetter(func() string {
-			return fmt.Sprintf("/organizations/%s/", cfg.OrganizationID)
-		}))
-
-	resultFormat, err := resourcemanager.Formats.GetWithFallback("json", "json")
-	if err != nil {
-		return environment{}, err
-	}
-
-	raw, err := resource.Get(ctx, cfg.EnvironmentID, resultFormat)
-	if err != nil {
-		return environment{}, err
-	}
-
-	var env environment
-	err = json.Unmarshal([]byte(raw), &env)
-	if err != nil {
-		return environment{}, err
-	}
-
-	return env, nil
-}
diff --git a/agent/runner/runner.go b/agent/runner/runner.go
index 65dfd1beff..a1d3184f12 100644
--- a/agent/runner/runner.go
+++ b/agent/runner/runner.go
@@ -96,17 +96,17 @@ func (s *Runner) onStartAgent(ctx context.Context, cfg config.Config) {
 		return
 	}
 
-	env, err := s.getEnvironment(ctx, cfg)
+	agentToken, err := s.getAgentToken(ctx, cfg.FullURL(), cfg.OrganizationID, cfg.EnvironmentID, cfg.Jwt)
 	if err != nil {
 		s.ui.Error(err.Error())
 	}
 
-	if env.AgentApiKey == "" {
+	if agentToken == "" {
 		s.ui.Error("You are attempting to start the agent in an environment you do not have admin rights to. Please ask the administrator of this environment to grant you admin rights.")
 		return
 	}
 
-	err = s.StartAgent(ctx, cfg.AgentEndpoint, env.AgentApiKey, cfg.UIEndpoint, cfg.EnvironmentID)
+	err = s.StartAgent(ctx, cfg.AgentEndpoint, agentToken, cfg.UIEndpoint, "")
 	if err != nil {
 		s.ui.Error(err.Error())
 	}