From e890152e3e18876b0444cd743ea84e5c4b9102b6 Mon Sep 17 00:00:00 2001
From: Tom Withers <thomas.withers@digital.justice.gov.uk>
Date: Tue, 23 Jun 2020 15:07:18 +0100
Subject: [PATCH] Add -time-to-protect flag

---
 README.md                 | 4 +++-
 cmd/register_workspace.go | 6 ++++--
 main.go                   | 8 +++++---
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 5c97ed7..6df147f 100644
--- 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
index 829c4be..543387e 100644
--- 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
index 077db84..9a897ef 100644
--- a/main.go
+++ b/main.go
@@ -8,7 +8,7 @@ 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()
 	}
@@ -16,10 +16,12 @@ func main()  {
 	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()