Inject a unique detonation ID in the user-agent of K8s and AWS requests

go.mod

@@ -19,6 +19,7 @@ require (
 	github.com/aws/smithy-go v1.10.0
 	github.com/fatih/color v1.13.0
 	github.com/golang-jwt/jwt v3.2.2+incompatible
+	github.com/google/uuid v1.3.0
 	github.com/hashicorp/terraform-exec v0.15.0
 	github.com/jedib0t/go-pretty/v6 v6.2.4
 	github.com/spf13/cobra v1.3.0

go.sum

@@ -324,6 +324,8 @@ github.com/google/pprof v0.0.0-20210609004039-a478d1d731e9/go.mod h1:kpwsk12EmLe
 github.com/google/pprof v0.0.0-20210720184732-4bb14d4b1be1/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE=
 github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI=
 github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
+github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg=
 github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk=
 github.com/googleapis/gax-go/v2 v2.1.0/go.mod h1:Q3nei7sK6ybPYH7twZdmQpAd1MKb7pfu6SK+H1/DsU0=

internal/providers/aws.go

@@ -8,22 +8,26 @@ import (
 	"github.com/aws/aws-sdk-go-v2/service/ec2"
 	"github.com/aws/smithy-go/middleware"
 	smithyhttp "github.com/aws/smithy-go/transport/http"
+	"github.com/google/uuid"
 	"log"
 )
 
-var awsProvider = AWSProvider{}
+var awsProvider = AWSProvider{
+	UniqueCorrelationId: uuid.New(),
+}
 
 func AWS() *AWSProvider {
 	return &awsProvider
 }
 
 type AWSProvider struct {
-	awsConfig *aws.Config
+	awsConfig           *aws.Config
+	UniqueCorrelationId uuid.UUID // unique value injected in the user-agent, to differentiate Stratus Red Team executions
 }
 
 func (m *AWSProvider) GetConnection() aws.Config {
 	if m.awsConfig == nil {
-		cfg, err := config.LoadDefaultConfig(context.Background(), customUserAgentApiOptions)
+		cfg, err := config.LoadDefaultConfig(context.Background(), customUserAgentApiOptions(m.UniqueCorrelationId))
 		if err != nil {
 			log.Fatalf("unable to load AWS configuration, %v", err)
 		}
@@ -45,25 +49,27 @@ func (m *AWSProvider) IsAuthenticatedAgainstAWS() bool {
 }
 
 // Functions below are related to customization of the user-agent header
-// Code taken from https://github.com/aws/aws-sdk-go-v2/issues/1432
-
-var customUserAgentApiOptions = config.WithAPIOptions(func() (v []func(stack *middleware.Stack) error) {
-	v = append(v, attachCustomMiddleware)
-	return v
-}())
+// Code mostly taken from https://github.com/aws/aws-sdk-go-v2/issues/1432
 
-var customerUAMiddleware = middleware.BuildMiddlewareFunc("CustomerUserAgent", func(
-	ctx context.Context, input middleware.BuildInput, next middleware.BuildHandler,
-) (out middleware.BuildOutput, metadata middleware.Metadata, err error) {
-	request, ok := input.Request.(*smithyhttp.Request)
-	if !ok {
-		return out, metadata, fmt.Errorf("unknown transport type %T", input.Request)
-	}
-	request.Header.Set("User-Agent", StratusUserAgent)
+func customUserAgentApiOptions(uniqueCorrelationId uuid.UUID) config.LoadOptionsFunc {
+	return config.WithAPIOptions(func() (v []func(stack *middleware.Stack) error) {
+		v = append(v, func(stack *middleware.Stack) error {
+			return stack.Build.Add(customUserAgentMiddleware(uniqueCorrelationId), middleware.After)
+		})
+		return v
+	}())
+}
 
-	return next.HandleBuild(ctx, input)
-})
+func customUserAgentMiddleware(uniqueId uuid.UUID) middleware.BuildMiddleware {
+	return middleware.BuildMiddlewareFunc("CustomerUserAgent", func(
+		ctx context.Context, input middleware.BuildInput, next middleware.BuildHandler,
+	) (out middleware.BuildOutput, metadata middleware.Metadata, err error) {
+		request, ok := input.Request.(*smithyhttp.Request)
+		if !ok {
+			return out, metadata, fmt.Errorf("unknown transport type %T", input.Request)
+		}
+		request.Header.Set("User-Agent", StratusUserAgent+"_"+uniqueId.String())
 
-func attachCustomMiddleware(stack *middleware.Stack) error {
-	return stack.Build.Add(customerUAMiddleware, middleware.After)
+		return next.HandleBuild(ctx, input)
+	})
 }

internal/providers/kubernetes.go

@@ -3,29 +3,30 @@ package providers
 import (
 	"context"
 	"github.com/datadog/stratus-red-team/internal/utils"
-	"k8s.io/client-go/rest"
-	"log"
-	"os"
-	"path/filepath"
-
+	"github.com/google/uuid"
 	authv1 "k8s.io/api/authorization/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/rest"
 	"k8s.io/client-go/tools/clientcmd"
 	"k8s.io/client-go/util/homedir"
+	"log"
+	"os"
+	"path/filepath"
 )
 
 const (
 	KubeconfigDefaultPath = ".kube/config"
 )
 
 type K8sProvider struct {
-	k8sClient  *kubernetes.Clientset
-	RestConfig *rest.Config
+	k8sClient           *kubernetes.Clientset
+	RestConfig          *rest.Config
+	UniqueCorrelationId uuid.UUID // unique value injected in the user-agent, to differentiate Stratus Red Team executions
 }
 
 var (
-	k8sProvider               K8sProvider
+	k8sProvider               = K8sProvider{UniqueCorrelationId: uuid.New()}
 	kubeConfigPath            string
 	kubeConfigPathWasResolved bool
 )
@@ -75,7 +76,7 @@ func (m *K8sProvider) GetClient() *kubernetes.Clientset {
 		log.Fatalf("unable to build kube config: %v", err)
 	}
 	m.RestConfig = config
-	m.RestConfig.UserAgent = StratusUserAgent
+	m.RestConfig.UserAgent = StratusUserAgent + "_" + m.UniqueCorrelationId.String()
 	m.k8sClient, err = kubernetes.NewForConfig(m.RestConfig)
 	if err != nil {
 		log.Fatalf("unable to create kube client: %v", err)