From 1f55116efa2ebb3a86e462adfb1c20c35f225c56 Mon Sep 17 00:00:00 2001
From: Elie <elie.charra@snyk.io>
Date: Tue, 3 May 2022 11:15:26 +0200
Subject: [PATCH] Add client to telemetry

---
 main.go                         |  3 +--
 pkg/config/config.go            |  4 ++++
 pkg/telemetry/telemetry.go      |  7 +++++++
 pkg/telemetry/telemetry_test.go | 35 +++++++++++++++++++++++++++++++++
 4 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/main.go b/main.go
index 8181eb8a..b68ecd42 100644
--- a/main.go
+++ b/main.go
@@ -17,7 +17,6 @@ import (
 	"github.com/snyk/driftctl/pkg/config"
 	"github.com/snyk/driftctl/pkg/version"
 	"github.com/snyk/driftctl/sentry"
-	"github.com/spf13/viper"
 )
 
 func init() {
@@ -35,7 +34,7 @@ func run() int {
 	logger.Init()
 	build := build.Build{}
 	// Check whether driftCTL is run under Snyk CLI
-	isSnyk := viper.GetBool("IS_SNYK")
+	isSnyk := config.IsSnyk()
 	logrus.WithFields(logrus.Fields{
 		"isRelease":               fmt.Sprintf("%t", build.IsRelease()),
 		"isUsageReportingEnabled": fmt.Sprintf("%t", build.IsUsageReportingEnabled()),
diff --git a/pkg/config/config.go b/pkg/config/config.go
index 34ae48cc..d630e56d 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -7,3 +7,7 @@ func Init() {
 	viper.AutomaticEnv()
 	viper.SetEnvPrefix("dctl")
 }
+
+func IsSnyk() bool {
+	return viper.GetBool("IS_SNYK")
+}
diff --git a/pkg/telemetry/telemetry.go b/pkg/telemetry/telemetry.go
index 1102e86b..fefc4fe9 100644
--- a/pkg/telemetry/telemetry.go
+++ b/pkg/telemetry/telemetry.go
@@ -8,6 +8,7 @@ import (
 
 	"github.com/sirupsen/logrus"
 	"github.com/snyk/driftctl/build"
+	"github.com/snyk/driftctl/pkg/config"
 	"github.com/snyk/driftctl/pkg/memstore"
 	"github.com/snyk/driftctl/pkg/version"
 )
@@ -21,6 +22,7 @@ type telemetry struct {
 	Duration       uint   `json:"duration"`
 	ProviderName   string `json:"provider_name"`
 	IaCSourceCount uint   `json:"iac_source_count"`
+	Client         string `json:"client"`
 }
 
 type Telemetry struct {
@@ -42,6 +44,11 @@ func (te Telemetry) SendTelemetry(store memstore.Bucket) {
 		Version: version.Current(),
 		Os:      runtime.GOOS,
 		Arch:    runtime.GOARCH,
+		Client:  "driftctl",
+	}
+
+	if config.IsSnyk() {
+		t.Client = "snyk-cli"
 	}
 
 	if val, ok := store.Get("total_resources").(int); ok {
diff --git a/pkg/telemetry/telemetry_test.go b/pkg/telemetry/telemetry_test.go
index 51c400eb..d7149e61 100644
--- a/pkg/telemetry/telemetry_test.go
+++ b/pkg/telemetry/telemetry_test.go
@@ -13,6 +13,7 @@ import (
 	"github.com/snyk/driftctl/pkg/resource"
 	"github.com/snyk/driftctl/pkg/version"
 	"github.com/snyk/driftctl/test/mocks"
+	"github.com/spf13/viper"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -45,6 +46,7 @@ func TestSendTelemetry(t *testing.T) {
 				Duration:       123,
 				ProviderName:   "aws",
 				IaCSourceCount: 2,
+				Client:         "driftctl",
 			},
 			setStoreValues: func(s memstore.Bucket, a *analyser.Analysis) {
 				s.Set("total_resources", a.Summary().TotalResources)
@@ -67,6 +69,7 @@ func TestSendTelemetry(t *testing.T) {
 				Arch:         runtime.GOARCH,
 				Duration:     124,
 				ProviderName: "aws",
+				Client:       "driftctl",
 			},
 			setStoreValues: func(s memstore.Bucket, a *analyser.Analysis) {
 				s.Set("total_resources", a.Summary().TotalResources)
@@ -90,6 +93,7 @@ func TestSendTelemetry(t *testing.T) {
 				Version: version.Current(),
 				Os:      runtime.GOOS,
 				Arch:    runtime.GOARCH,
+				Client:  "driftctl",
 			},
 			setStoreValues: func(s memstore.Bucket, a *analyser.Analysis) {},
 		},
@@ -150,3 +154,34 @@ func TestTelemetryNotSend(t *testing.T) {
 
 	assert.Zero(t, httpmock.GetTotalCallCount())
 }
+
+func TestTelemetrySetProperClient(t *testing.T) {
+	httpmock.Activate()
+	defer httpmock.DeactivateAndReset()
+	httpmock.RegisterResponder(
+		"POST",
+		"https://telemetry.driftctl.com/telemetry",
+		func(req *http.Request) (*http.Response, error) {
+			requestTelemetry := &telemetry{}
+			requestBody, err := ioutil.ReadAll(req.Body)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			err = json.Unmarshal(requestBody, requestTelemetry)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			assert.Equal(t, "snyk-cli", requestTelemetry.Client)
+
+			return httpmock.NewBytesResponse(202, []byte{}), nil
+		},
+	)
+
+	viper.Set("IS_SNYK", true)
+	store := memstore.New().Bucket(memstore.TelemetryBucket)
+	tl := NewTelemetry(mocks.MockBuild{UsageReporting: true})
+	tl.SendTelemetry(store)
+	assert.Equal(t, 1, httpmock.GetTotalCallCount())
+}