From ee57eb36555885dada8d0258bf4658af6bcdb5f8 Mon Sep 17 00:00:00 2001
From: Fabio Gollinucci <fbgolly@gmail.com>
Date: Fri, 12 Mar 2021 17:50:20 +0100
Subject: [PATCH] Add update tag step

---
 cmd/deploy/deploy.go      | 17 ++++++++-
 internal/canary/canary.go | 78 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 94 insertions(+), 1 deletion(-)

diff --git a/cmd/deploy/deploy.go b/cmd/deploy/deploy.go
index fa4d1dd..a956daf 100644
--- a/cmd/deploy/deploy.go
+++ b/cmd/deploy/deploy.go
@@ -286,14 +286,29 @@ func deploySingleCanary(ses *session.Session, region *string, accountID *string,
 		}
 	}
 
+	isAlreadyDeployed := canary.IsDeployed()
+
 	// Deploy canary
-	fmt.Println(fmt.Sprintf("[%s] Deploying..", canary.Name))
+	if !isAlreadyDeployed {
+		fmt.Println(fmt.Sprintf("[%s] Creating..", canary.Name))
+	} else {
+		fmt.Println(fmt.Sprintf("[%s] Updating..", canary.Name))
+	}
 	artifactBucketLocation := *artifactBucket.Location + "/canary/" + canary.Name
 	err = canary.Deploy(role, &artifactBucketLocation)
 	if err != nil {
 		return err
 	}
 
+	// Update tags
+	if isAlreadyDeployed {
+		fmt.Println(fmt.Sprintf("[%s] Updating tags..", canary.Name))
+		err = canary.UpdateTags(region, accountID)
+		if err != nil {
+			return err
+		}
+	}
+
 	// Wait until canary is created
 	var status *synthetics.CanaryStatus
 	fmt.Println(fmt.Sprintf("[%s] Waiting..", canary.Name))
diff --git a/internal/canary/canary.go b/internal/canary/canary.go
index 86326d7..9e9c705 100644
--- a/internal/canary/canary.go
+++ b/internal/canary/canary.go
@@ -11,6 +11,7 @@ import (
 	"github.com/aws/aws-sdk-go/service/lambda"
 	"github.com/aws/aws-sdk-go/service/s3"
 	"github.com/aws/aws-sdk-go/service/s3/s3manager"
+	"github.com/aws/aws-sdk-go/service/sts"
 	"github.com/aws/aws-sdk-go/service/synthetics"
 	"github.com/daaru00/aws-canary-cli/internal/iam"
 )
@@ -20,6 +21,7 @@ type clients struct {
 	s3         *s3.S3
 	s3uploader *s3manager.Uploader
 	lambda     *lambda.Lambda
+	sts        *sts.STS
 }
 
 // Schedule configuration
@@ -67,6 +69,7 @@ func New(ses *session.Session, name string) *Canary {
 		s3:         s3.New(ses),
 		s3uploader: s3manager.NewUploader(ses),
 		lambda:     lambda.New(ses),
+		sts:        sts.New(ses),
 	}
 
 	return &Canary{
@@ -218,6 +221,81 @@ func (c *Canary) Deploy(role *iam.Role, artifactBucketLocation *string) error {
 	return err
 }
 
+// UpdateTags update canary tags
+func (c *Canary) UpdateTags(region *string, account *string) error {
+	// Build ARN
+	arn := fmt.Sprintf("arn:aws:synthetics:%s:%s:canary:%s", *region, *account, c.Name)
+
+	// Get current tags
+	resTags, err := c.clients.synthetics.ListTagsForResource(&synthetics.ListTagsForResourceInput{
+		ResourceArn: &arn,
+	})
+	if err != nil {
+		return err
+	}
+
+	// Skip if not tags are set
+	if len(c.Tags) == 0 && len(resTags.Tags) == 0 {
+		return nil
+	}
+
+	// Check tags to add
+	tagsToAdd := map[string]string{}
+	for key, value := range c.Tags {
+		var foundTagKey *string
+		for currentTagKey, currentTagValue := range resTags.Tags {
+			if key == currentTagKey && value == *currentTagValue {
+				foundTagKey = &key
+				break
+			}
+		}
+
+		if foundTagKey == nil {
+			tagsToAdd[key] = value
+		}
+	}
+
+	// Add missing tags
+	if len(tagsToAdd) > 0 {
+		_, err := c.clients.synthetics.TagResource(&synthetics.TagResourceInput{
+			ResourceArn: &arn,
+			Tags:        aws.StringMap(tagsToAdd),
+		})
+		if err != nil {
+			return err
+		}
+	}
+
+	// Check accounts ids to remove
+	tagsKeysToRemove := []string{}
+	for currentTagKey := range resTags.Tags {
+		var foundTagKey *string
+		for key := range c.Tags {
+			if key == currentTagKey {
+				foundTagKey = &key
+				break
+			}
+		}
+
+		if foundTagKey == nil {
+			tagsKeysToRemove = append(tagsKeysToRemove, currentTagKey)
+		}
+	}
+
+	// Remove unused tags
+	if len(tagsKeysToRemove) > 0 {
+		_, err = c.clients.synthetics.UntagResource(&synthetics.UntagResourceInput{
+			ResourceArn: &arn,
+			TagKeys:     aws.StringSlice(tagsKeysToRemove),
+		})
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 // Start canary
 func (c *Canary) Start() error {
 	_, err := c.clients.synthetics.StartCanary(&synthetics.StartCanaryInput{