From f8b93219e72bf8d1e464e56ebfb41b9712e7323b Mon Sep 17 00:00:00 2001
From: Christophe Tafani-Dereeper <christophe.tafanidereeper@datadoghq.com>
Date: Wed, 26 Jan 2022 16:08:21 +0100
Subject: [PATCH] Add reversion function to aws.persistence.backdoor-iam-role
 (closes #65)

---
 .../iam-role-backdoor-existing/main.go        | 29 +++++++++++++++----
 .../iam-role-backdoor-existing/main.tf        |  4 +++
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/internal/attacktechniques/aws/persistence/iam-role-backdoor-existing/main.go b/internal/attacktechniques/aws/persistence/iam-role-backdoor-existing/main.go
index e62dc540..6097e5da 100644
--- a/internal/attacktechniques/aws/persistence/iam-role-backdoor-existing/main.go
+++ b/internal/attacktechniques/aws/persistence/iam-role-backdoor-existing/main.go
@@ -9,6 +9,7 @@ import (
 	"github.com/datadog/stratus-red-team/pkg/stratus"
 	"github.com/datadog/stratus-red-team/pkg/stratus/mitreattack"
 	"log"
+	"strings"
 )
 
 //go:embed main.tf
@@ -43,21 +44,39 @@ Detonation:
 		MitreAttackTactics:         []mitreattack.Tactic{mitreattack.Persistence},
 		PrerequisitesTerraformCode: tf,
 		Detonate:                   detonate,
+		Revert:                     revert,
 	})
 }
 
 func detonate(params map[string]string) error {
-	iamClient := iam.NewFromConfig(providers.AWS().GetConnection())
 	roleName := params["role_name"]
 
 	log.Println("Backdooring IAM role " + roleName + " by allowing sts:AssumeRole from an external AWS account")
-	_, err := iamClient.UpdateAssumeRolePolicy(context.Background(), &iam.UpdateAssumeRolePolicyInput{
-		RoleName:       &roleName,
-		PolicyDocument: &maliciousIamPolicy,
-	})
+	err := updateAssumeRolePolicy(roleName, maliciousIamPolicy)
+	if err != nil {
+		return errors.New("unable to backdoor IAM role: " + err.Error())
+	}
+	return nil
+}
+
+func revert(params map[string]string) error {
+	roleName := params["role_name"]
+	roleTrustPolicy := strings.ReplaceAll(params["role_trust_policy"], "\\", "") // Terraform output adds backslashes for some reason
+
+	log.Println("Reverting trust policy of IAM role " + roleName + " to its original state")
+	err := updateAssumeRolePolicy(roleName, roleTrustPolicy)
 
 	if err != nil {
 		return errors.New("unable to backdoor IAM role: " + err.Error())
 	}
 	return nil
 }
+
+func updateAssumeRolePolicy(roleName string, roleTrustPolicy string) error {
+	iamClient := iam.NewFromConfig(providers.AWS().GetConnection())
+	_, err := iamClient.UpdateAssumeRolePolicy(context.Background(), &iam.UpdateAssumeRolePolicyInput{
+		RoleName:       &roleName,
+		PolicyDocument: &roleTrustPolicy,
+	})
+	return err
+}
diff --git a/internal/attacktechniques/aws/persistence/iam-role-backdoor-existing/main.tf b/internal/attacktechniques/aws/persistence/iam-role-backdoor-existing/main.tf
index 35039562..bc791c8c 100644
--- a/internal/attacktechniques/aws/persistence/iam-role-backdoor-existing/main.tf
+++ b/internal/attacktechniques/aws/persistence/iam-role-backdoor-existing/main.tf
@@ -39,6 +39,10 @@ output "role_name" {
   value = aws_iam_role.legit-role.name
 }
 
+output "role_trust_policy" {
+  value = aws_iam_role.legit-role.assume_role_policy
+}
+
 output "display" {
   value = format("IAM role %s ready", aws_iam_role.legit-role.name)
 }
\ No newline at end of file