Add -time-to-protect flag

TomTucka · Jun 23, 2020 · e890152 · e890152
1 parent 85e69d7
commit e890152
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 6 deletions.
diff --git a/README.md b/README.md
@@ -25,7 +25,7 @@ module "workspace-cleanup" {
 ## Using this tool
 
 ```
-Usage: tf-workspace-cleanup -register-workspace=<workspace> -aws-account-id=12345678 -aws-iam-role=operator
+Usage: tf-workspace-cleanup -register-workspace=<workspace> -time-to-protect=2 -aws-account-id=12345678 -aws-iam-role=operator
 Usage: tf-workspace-cleanup -protected-workspaces=true -aws-account-id=12345678 -aws-iam-role=operator
   -aws-account-id string
     	Account ID for IAM Role
@@ -35,4 +35,6 @@ Usage: tf-workspace-cleanup -protected-workspaces=true -aws-account-id=12345678
     	get list of protected workspaces for deletion
   -register-workspace string
     	Register a workspace to be deleted at a later point
+  -time-to-protect=2
+        Time in hours to protect workspace for
 ```
diff --git a/cmd/register_workspace.go b/cmd/register_workspace.go
@@ -18,7 +18,7 @@ type Session struct {
 	AwsSession *session.Session
 }
 
-func RegisterWorkspace(workspace *string, accountId *string, iamRoleName *string) {
+func RegisterWorkspace(workspace *string, accountId *string, iamRoleName *string, timeToProtect *int64) {
 
 	sess, err := session.NewSession()
 	if err != nil {
@@ -37,9 +37,11 @@ func RegisterWorkspace(workspace *string, accountId *string, iamRoleName *string
 		ExpiresTTL    int64
 	}
 
+
 	item := Workspace{
 		WorkspaceName: *workspace,
-		ExpiresTTL:    time.Now().AddDate(0, 0, 1).Unix(),
+		ExpiresTTL:    time.Now().Add(time.Hour * time.Duration(*timeToProtect)).Unix(),
+
 	}
 
 	WorkspaceToPut, err := dynamodbattribute.MarshalMap(item)

diff --git a/main.go b/main.go
@@ -8,18 +8,20 @@ import (
 
 func main()  {
 	flag.Usage = func() {
-		fmt.Println("Usage: tf-workspace-cleanup -register-workspace=<workspace> -aws-account-id=12345678 -aws-iam-role=sirius-ci")
+		fmt.Println("Usage: tf-workspace-cleanup -register-workspace=<workspace> -time-to-protect=2 -aws-account-id=12345678 -aws-iam-role=sirius-ci")
 		fmt.Println("Usage: tf-workspace-cleanup -protected-workspaces=true -aws-account-id=12345678 -aws-iam-role=sirius-ci")
 		flag.PrintDefaults()
 	}
 	var workspaceName string
 	var protectedWorkspaces bool
 	var awsAccountId string
 	var awsIAMRoleName string
+	var timeToProtect int64
 
 	flag.StringVar(&workspaceName, "register-workspace", "", "Register a workspace to be deleted at a later point")
 	flag.StringVar(&awsAccountId, "aws-account-id", "", "Account ID for IAM Role")
-	flag.StringVar(&awsIAMRoleName, "aws-iam-role", "", "AWS IAM Role Name ")
+	flag.StringVar(&awsIAMRoleName, "aws-iam-role", "", "AWS IAM Role Name")
+	flag.Int64Var(&timeToProtect, "time-to-protect", 1 , "Time in hours to protect workspace for")
 	flag.BoolVar(&protectedWorkspaces, "protected-workspaces", false, "get list of protected workspaces for deletion")
 	flag.Parse()
 
@@ -38,7 +40,7 @@ func main()  {
 	}
 
 	if workspaceName != "" {
-		cmd.RegisterWorkspace(&workspaceName, &awsAccountId, &awsIAMRoleName)
+		cmd.RegisterWorkspace(&workspaceName, &awsAccountId, &awsIAMRoleName, &timeToProtect)
 	} else if protectedWorkspaces != true {
 		fmt.Println("Error: Workspace not passed")
 		flag.Usage()