From 23663ed4c20157a4bdadfac4e29495acf8417d80 Mon Sep 17 00:00:00 2001
From: TP Honey <tphoney@gmail.com>
Date: Thu, 31 Oct 2024 14:54:54 +0000
Subject: [PATCH] (feat) Engine CreateClients helper method

---
 cmd.go | 63 ++++++++++++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 52 insertions(+), 11 deletions(-)

diff --git a/cmd.go b/cmd.go
index 1d4ee12..339cfb9 100644
--- a/cmd.go
+++ b/cmd.go
@@ -2,13 +2,21 @@ package discovery
 
 import (
 	"fmt"
+	"net/http"
 	"os"
 	"runtime"
+	"time"
 
+	"github.com/getsentry/sentry-go"
 	"github.com/google/uuid"
 	"github.com/overmindtech/sdp-go"
+	"github.com/overmindtech/sdp-go/auth"
+	"github.com/overmindtech/sdp-go/sdpconnect"
+	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
+	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
+	"golang.org/x/oauth2"
 )
 
 func AddEngineFlags(command *cobra.Command) {
@@ -22,7 +30,6 @@ func AddEngineFlags(command *cobra.Command) {
 	cobra.CheckErr(viper.BindEnv("api-key", "OVM_API_KEY", "API_KEY"))
 	command.PersistentFlags().String("source-access-token", "", "The access token to use to authenticate the source for managed sources")
 	cobra.CheckErr(viper.BindEnv("source-access-token", "SOURCE_ACCESS_TOKEN"))
-	// SOURCE_TOKEN_TYPE
 	command.PersistentFlags().String("source-token-type", "", "The type of token to use to authenticate the source for managed sources")
 	cobra.CheckErr(viper.BindEnv("source-token-type", "SOURCE_TOKEN_TYPE"))
 	command.PersistentFlags().Int("max-parallel", 0, "The maximum number of parallel executions")
@@ -88,21 +95,55 @@ func EngineConfigFromViper(engineType, version string) (*EngineConfig, error) {
 }
 
 // MapFromEngineConfig Returns the config as a map
-func MapFromEngineConfig(c *EngineConfig) map[string]any {
-
+func MapFromEngineConfig(ec *EngineConfig) map[string]any {
 	var apiKeyClientSecret string
-
-	if c.ApiKey != "" {
+	if ec.ApiKey != "" {
 		apiKeyClientSecret = "[REDACTED]"
 	}
+	var sourceAccessToken string
+	if ec.SourceAccessToken != "" {
+		sourceAccessToken = "[REDACTED]"
+	}
 
 	return map[string]interface{}{
-		"engine-type":             c.EngineType,
-		"version":                 c.Version,
-		"source-name":             c.SourceName,
-		"source-uuid":             c.SourceUUID,
-		"app":                     c.App,
+		"engine-type":             ec.EngineType,
+		"version":                 ec.Version,
+		"source-name":             ec.SourceName,
+		"source-uuid":             ec.SourceUUID,
+		"source-access-token":     sourceAccessToken,
+		"source-token-type":       ec.SourceTokenType,
+		"managed-source":          ec.OvermindManagedSource,
+		"app":                     ec.App,
 		"api-key":                 apiKeyClientSecret,
-		"max-parallel-executions": c.MaxParallelExecutions,
+		"max-parallel-executions": ec.MaxParallelExecutions,
+	}
+}
+
+func (ec *EngineConfig) CreateClients(oi sdp.OvermindInstance) (auth.TokenClient, *HeartbeatOptions, error) {
+	if ec.ApiKey != "" {
+		tokenClient, err := auth.NewAPIKeyClient(oi.ApiUrl.String(), ec.ApiKey)
+		if err != nil {
+			sentry.CaptureException(err)
+
+			log.WithError(err).Fatal("Could not create API key client")
+		}
+
+		tokenSource := auth.NewAPIKeyTokenSource(ec.ApiKey, oi.ApiUrl.String())
+		transport := oauth2.Transport{
+			Source: tokenSource,
+			Base:   http.DefaultTransport,
+		}
+		authenticatedClient := http.Client{
+			Transport: otelhttp.NewTransport(&transport),
+		}
+		heartbeatOptions := HeartbeatOptions{
+			ManagementClient: sdpconnect.NewManagementServiceClient(
+				&authenticatedClient,
+				oi.ApiUrl.String(),
+			),
+			Frequency: time.Second * 30,
+		}
+		return tokenClient, &heartbeatOptions, nil
 	}
+	return nil, nil, fmt.Errorf("No clients created")
 }