Add support for tagging EC2 Snapshots (#234)

README.MD

@@ -72,6 +72,18 @@ Read csv and tag ec2 - `awstaghelper ec2 tag-sg`
 Example:
  `awstaghelper ec2 tag-sg --filename sgTag.csv --profile main`  
 
+#### Get snapshot tags
+
+Get list of SGs with required tags - `awstaghelper ec2 get-snapshot-tags`  
+Example:
+`awstaghelper ec2 get-snapshot-tags --filename snapshotTag.csv --tags Name,Owner --profile main`
+
+#### Tag snapshot
+
+Read csv and tag ec2 - `awstaghelper ec2 tag-snapshot`  
+Example:
+`awstaghelper ec2 tag-snapshot --filename snapshotTag.csv --profile main`
+
 ### Rds
 
 #### Get rds tags

cmd/ec2.go

@@ -27,8 +27,8 @@ import (
 // ec2Cmd represents the ec2 command
 var ec2Cmd = &cobra.Command{
 	Use:   "ec2",
-	Short: "Root command for interaction with AWS ec2 services",
-	Long:  `Root command for interaction with AWS ec2 services.`,
+	Short: "Root command for interaction with AWS ec2 services (ec2, security group, snapshot)",
+	Long:  `Root command for interaction with AWS ec2 services (ec2, security group, snapshot).`,
 	//Run: func(cmd *cobra.Command, args []string) {
 	//	fmt.Println("ec2 called")
 	//},

cmd/ec2_snapshot.go

@@ -0,0 +1,66 @@
+/*
+Copyright © 2024 Jaemok Hong [email protected]
+Copyright © 2020 Maksym Postument [email protected]
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+// Package cmd is the package for the CLI of awstaghelper
+package cmd
+
+import (
+	"github.com/mpostument/awstaghelper/pkg"
+
+	"github.com/aws/aws-sdk-go/service/ec2"
+	"github.com/aws/aws-sdk-go/service/sts"
+	"github.com/spf13/cobra"
+)
+
+var getSnapshotCmd = &cobra.Command{
+	Use:   "get-snapshot-tags",
+	Short: "Write snapshot id and required tags to csv",
+	Long: `Write to csv data with snapshot id and required tags to csv.
+This csv can be used with tag-snapshot command to tag aws environment.
+Specify list of tags which should be read using tags flag: --tags Name,Env,Project.
+Csv filename can be specified with flag filename.`,
+	Run: func(cmd *cobra.Command, args []string) {
+		tags, _ := cmd.Flags().GetString("tags")
+		filename, _ := cmd.Flags().GetString("filename")
+		profile, _ := cmd.Flags().GetString("profile")
+		region, _ := cmd.Flags().GetString("region")
+		sess := pkg.GetSession(region, profile)
+		client := ec2.New(sess)
+		stsClient := sts.New(sess)
+		pkg.WriteCsv(pkg.ParseSnapshotTags(tags, client, stsClient), filename)
+	},
+}
+
+var tagSnapshotCmd = &cobra.Command{
+	Use:   "tag-snapshot",
+	Short: "Read csv and tag snapshot with csv data",
+	Long:  `Read csv generated with get-snapshot-tags command and tag snapshot resources with tags from csv.`,
+	Run: func(cmd *cobra.Command, args []string) {
+		filename, _ := cmd.Flags().GetString("filename")
+		profile, _ := cmd.Flags().GetString("profile")
+		region, _ := cmd.Flags().GetString("region")
+		sess := pkg.GetSession(region, profile)
+		csvData := pkg.ReadCsv(filename)
+		client := ec2.New(sess)
+		pkg.TagSnapshot(csvData, client)
+	},
+}
+
+func init() {
+	ec2Cmd.AddCommand(getSnapshotCmd)
+	ec2Cmd.AddCommand(tagSnapshotCmd)
+}

pkg/ec2_snapshot.go

@@ -0,0 +1,91 @@
+/*
+Copyright © 2024 Jaemok Hong [email protected]
+Copyright © 2020 Maksym Postument [email protected]
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package pkg
+
+import (
+	"log"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/ec2"
+	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
+	"github.com/aws/aws-sdk-go/service/sts"
+	"github.com/aws/aws-sdk-go/service/sts/stsiface"
+)
+
+// getSecurityGroups return all security groups from specified region
+func getSnapshot(callerIdentity string, client ec2iface.EC2API) []*ec2.Snapshot {
+	input := &ec2.DescribeSnapshotsInput{
+		OwnerIds: []*string{aws.String(callerIdentity)},
+	}
+
+	var result []*ec2.Snapshot
+
+	err := client.DescribeSnapshotsPages(input,
+		func(page *ec2.DescribeSnapshotsOutput, lastPage bool) bool {
+			result = append(result, page.Snapshots...)
+			return !lastPage
+		})
+	if err != nil {
+		log.Fatal("Not able to get snapshots ", err)
+	}
+	return result
+}
+
+// ParseSecurityGroupTags parse output from getSecurityGroups and return SG ids and specified tags.
+func ParseSnapshotTags(tagsToRead string, client ec2iface.EC2API, stsClient stsiface.STSAPI) [][]string {
+	callerIdentity, err := stsClient.GetCallerIdentity(&sts.GetCallerIdentityInput{})
+	if err != nil {
+		log.Fatal("Not able to get account id", err)
+	}
+	snapshotsOutput := getSnapshot(*callerIdentity.Account, client)
+
+	rows := addHeadersToCsv(tagsToRead, "Id")
+	for _, snapshot := range snapshotsOutput {
+		tags := map[string]string{}
+		for _, tag := range snapshot.Tags {
+			tags[*tag.Key] = *tag.Value
+		}
+		rows = addTagsToCsv(tagsToRead, tags, rows, *snapshot.SnapshotId)
+	}
+	return rows
+}
+
+// TagSecurityGroups tag security groups. Take as input data from csv file. Where first column id
+func TagSnapshot(csvData [][]string, client ec2iface.EC2API) {
+	for r := 1; r < len(csvData); r++ {
+		var tags []*ec2.Tag
+		for c := 1; c < len(csvData[0]); c++ {
+			tags = append(tags, &ec2.Tag{
+				Key:   &csvData[0][c],
+				Value: &csvData[r][c],
+			})
+		}
+
+		input := &ec2.CreateTagsInput{
+			Resources: []*string{
+				aws.String(csvData[r][0]),
+			},
+			Tags: tags,
+		}
+
+		_, err := client.CreateTags(input)
+		if awsErrorHandle(err) {
+			return
+		}
+	}
+}

pkg/ec2_snapshot_test.go

@@ -0,0 +1,112 @@
+/*
+Copyright © 2024 Jaemok Hong [email protected]
+Copyright © 2020 Maksym Postument [email protected]
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package pkg
+
+import (
+	"testing"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/ec2"
+	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
+	"github.com/aws/aws-sdk-go/service/sts"
+	"github.com/aws/aws-sdk-go/service/sts/stsiface"
+	"github.com/stretchr/testify/assert"
+)
+
+type mockedSnapshot struct {
+	ec2iface.EC2API
+	stsiface.STSAPI
+	respGetCallerIdentity sts.GetCallerIdentityOutput
+	respDescribeSnapshots ec2.DescribeSnapshotsOutput
+}
+
+// AWS Mocks
+func (m *mockedSnapshot) DescribeSnapshotsPages(
+	input *ec2.DescribeSnapshotsInput,
+	pageFunc func(*ec2.DescribeSnapshotsOutput, bool) bool) error {
+	pageFunc(&m.respDescribeSnapshots, true)
+	return nil
+}
+
+func (m *mockedSnapshot) GetCallerIdentity(*sts.GetCallerIdentityInput) (*sts.GetCallerIdentityOutput, error) {
+	return &m.respGetCallerIdentity, nil
+}
+
+// Tests
+func TestGetSnapshots(t *testing.T) {
+	cases := []*mockedSnapshot{
+		{
+			respDescribeSnapshots: parseSnapshotTagsResponse,
+		},
+	}
+
+	expectedResult := parseSnapshotTagsResponse.Snapshots
+
+	for _, c := range cases {
+		t.Run("getSnapshot", func(t *testing.T) {
+			result := getSnapshot("callerIdentity", c)
+			assertions := assert.New(t)
+			assertions.EqualValues(expectedResult, result)
+		})
+
+	}
+}
+
+func TestParseSnapshotTags(t *testing.T) {
+	cases := []*mockedSnapshot{
+		{
+			respGetCallerIdentity: getSnapshotCallerIdentityResponse,
+			respDescribeSnapshots: parseSnapshotTagsResponse,
+		},
+	}
+	expectedResult := [][]string{
+		{"Id", "Name", "Environment", "Owner"},
+		{"snapshot-test", "TestSnapshot1", "Test", ""},
+	}
+	for _, c := range cases {
+		t.Run("ParseSnapshotTags", func(t *testing.T) {
+			result := ParseSnapshotTags("Name,Environment,Owner", c, c)
+			assertions := assert.New(t)
+			assertions.EqualValues(expectedResult, result)
+		})
+	}
+}
+
+var getSnapshotCallerIdentityResponse = sts.GetCallerIdentityOutput{
+	Account: aws.String("666666666"),
+}
+
+var parseSnapshotTagsResponse = ec2.DescribeSnapshotsOutput{
+	Snapshots: []*ec2.Snapshot{
+		{
+			Description: aws.String("testSg"),
+			SnapshotId:  aws.String("snapshot-test"),
+			VolumeId:    aws.String("testVolumId"),
+			Tags: []*ec2.Tag{
+				{
+					Key:   aws.String("Name"),
+					Value: aws.String("TestSnapshot1"),
+				},
+				{
+					Key:   aws.String("Environment"),
+					Value: aws.String("Test"),
+				},
+			},
+		},
+	},
+}