chore: Implement --aws_assume_role flag for CLI aws integration

cli/cmd/generate_aws.go

@@ -67,8 +67,9 @@ var (
 	// AwsArnRegex original source: https://regex101.com/r/pOfxYN/1
 	AwsArnRegex = `^arn:(?P<Partition>[^:\n]*):(?P<Service>[^:\n]*):(?P<Region>[^:\n]*):(?P<AccountID>[^:\n]*):(?P<Ignore>(?P<ResourceType>[^:\/\n]*)[:\/])?(?P<Resource>.*)$` //nolint
 	// AwsRegionRegex regex used for validating region input; note intentionally does not match gov cloud
-	AwsRegionRegex  = `(af|ap|ca|eu|me|sa|us)-(central|(north|south)?(east|west)?)-\d`
-	AwsProfileRegex = `([A-Za-z_0-9-]+)`
+	AwsRegionRegex     = `(af|ap|ca|eu|me|sa|us)-(central|(north|south)?(east|west)?)-\d`
+	AwsProfileRegex    = `([A-Za-z_0-9-]+)`
+	AwsAssumeRoleRegex = `^arn:aws:iam::\d{12}:role\/.*$`
 
 	GenerateAwsCommandState      = &aws.GenerateAwsTfConfigurationArgs{}
 	GenerateAwsExistingRoleState = &aws.ExistingIamRoleDetails{}
@@ -112,6 +113,7 @@ See help output for more details on the parameter value(s) required for Terrafor
 			// Setup modifiers for NewTerraform constructor
 			mods := []aws.AwsTerraformModifier{
 				aws.WithAwsProfile(GenerateAwsCommandState.AwsProfile),
+				aws.WithAwsAssumeRole(GenerateAwsCommandState.AwsAssumeRole),
 				aws.WithLaceworkProfile(GenerateAwsCommandState.LaceworkProfile),
 				aws.WithLaceworkAccountID(GenerateAwsCommandState.LaceworkAccountID),
 				aws.ExistingCloudtrailBucketArn(GenerateAwsCommandState.ExistingCloudtrailBucketArn),
@@ -194,6 +196,15 @@ See help output for more details on the parameter value(s) required for Terrafor
 				return err
 			}
 
+			// Validate aws assume role, if passed
+			assumeRole, err := cmd.Flags().GetString("aws_assume_role")
+			if err != nil {
+				return errors.Wrap(err, "failed to load command flags")
+			}
+			if err := validateAwsAssumeRole(assumeRole); assumeRole != "" && err != nil {
+				return err
+			}
+
 			// Validate aws profile, if passed
 			profile, err := cmd.Flags().GetString("aws_profile")
 			if err != nil {
@@ -353,6 +364,11 @@ func initGenerateAwsTfCommandFlags() {
 		"aws_profile",
 		"",
 		"specify aws profile")
+	generateAwsTfCommand.PersistentFlags().StringVar(
+		&GenerateAwsCommandState.AwsAssumeRole,
+		"aws_assume_role",
+		"",
+		"specify aws assume role")
 	generateAwsTfCommand.PersistentFlags().BoolVar(
 		&GenerateAwsCommandState.BucketEncryptionEnabled,
 		"bucket_encryption_enabled",
@@ -494,6 +510,11 @@ func validateAwsProfile(val interface{}) error {
 	return validateStringWithRegex(val, fmt.Sprintf(`^%s$`, AwsProfileRegex), "invalid profile name supplied")
 }
 
+// survey.Validator for aws profile
+func validateAwsAssumeRole(val interface{}) error {
+	return validateStringWithRegex(val, AwsAssumeRoleRegex, "invalid assume name supplied")
+}
+
 func promptAgentlessQuestions(config *aws.GenerateAwsTfConfigurationArgs) error {
 	return nil
 }

cli/docs/lacework_generate_cloud-account_aws.md

@@ -39,6 +39,7 @@ lacework generate cloud-account aws [flags]
 ```
       --agentless                             enable agentless integration
       --apply                                 run terraform apply without executing plan or prompting
+      --aws_assume_role string                specify aws assume role
       --aws_profile string                    specify aws profile
       --aws_region string                     specify aws region
       --aws_subaccount strings                configure an additional aws account; value format must be <aws profile>:<region>

integration/test_resources/help/generate_cloud-account_aws

@@ -27,6 +27,7 @@ Available Commands:
 Flags:
       --agentless                             enable agentless integration
       --apply                                 run terraform apply without executing plan or prompting
+      --aws_assume_role string                specify aws assume role
       --aws_profile string                    specify aws profile
       --aws_region string                     specify aws region
       --aws_subaccount strings                configure an additional aws account; value format must be <aws profile>:<region>

integration/test_resources/help/generate_cloud-account_aws_controltower

@@ -43,6 +43,7 @@ Global Flags:
   -k, --api_key string                        access key id
   -s, --api_secret string                     secret access key
       --api_token string                      access token (replaces the use of api_key and api_secret)
+      --aws_assume_role string                specify aws assume role
       --aws_profile string                    specify aws profile
       --aws_region string                     specify aws region
       --aws_subaccount strings                configure an additional aws account; value format must be <aws profile>:<region>

lwgenerate/aws/aws.go

@@ -95,6 +95,9 @@ type GenerateAwsTfConfigurationArgs struct {
 	// Supply an AWS Profile name for the main account, only asked if configuring multiple
 	AwsProfile string
 
+	// Supply an AWS Assume Role for the main account
+	AwsAssumeRole string
+
 	// Existing S3 Bucket ARN (Required when using existing cloudtrail)
 	ExistingCloudtrailBucketArn string
 
@@ -225,6 +228,13 @@ func WithAwsProfile(name string) AwsTerraformModifier {
 	}
 }
 
+// WithAwsAssumeRole Set the AWS Assume Role to utilize for the main AWS provider
+func WithAwsAssumeRole(assumeRole string) AwsTerraformModifier {
+	return func(c *GenerateAwsTfConfigurationArgs) {
+		c.AwsAssumeRole = assumeRole
+	}
+}
+
 // WithLaceworkProfile Set the Lacework Profile to utilize when integrating
 func WithLaceworkProfile(name string) AwsTerraformModifier {
 	return func(c *GenerateAwsTfConfigurationArgs) {
@@ -434,10 +444,23 @@ func createAwsProvider(args *GenerateAwsTfConfigurationArgs) ([]*hclwrite.Block,
 			attrs["alias"] = "main"
 		}
 
-		provider, err := lwgenerate.NewProvider(
-			"aws",
+		modifiers := []lwgenerate.HclProviderModifier{
 			lwgenerate.HclProviderWithAttributes(attrs),
-		).ToBlock()
+		}
+
+		if args.AwsAssumeRole != "" {
+			assumeRoleBlock, err := lwgenerate.HclCreateGenericBlock(
+				"assume_role",
+				nil,
+				map[string]interface{}{"role_arn": args.AwsAssumeRole},
+			)
+			if err != nil {
+				return nil, err
+			}
+			modifiers = append(modifiers, lwgenerate.HclProviderWithGenericBlocks(assumeRoleBlock))
+		}
+
+		provider, err := lwgenerate.NewProvider("aws", modifiers...).ToBlock()
 		if err != nil {
 			return nil, err
 		}

lwgenerate/hcl.go

@@ -58,10 +58,25 @@ type HclProvider struct {
 
 	// Optional. Extra properties for this module.  Can supply string, bool, int, or map[string]interface{} as values
 	attributes map[string]interface{}
+
+	// optional. Generic blocks
+	blocks []*hclwrite.Block
 }
 
 func (p *HclProvider) ToBlock() (*hclwrite.Block, error) {
-	return HclCreateGenericBlock("provider", []string{p.name}, p.attributes)
+	block, err := HclCreateGenericBlock("provider", []string{p.name}, p.attributes)
+	if err != nil {
+		return nil, err
+	}
+
+	if p.blocks != nil {
+		for _, b := range p.blocks {
+			block.Body().AppendNewline()
+			block.Body().AppendBlock(b)
+		}
+	}
+
+	return block, nil
 }
 
 type HclProviderModifier func(p *HclProvider)
@@ -81,6 +96,13 @@ func HclProviderWithAttributes(attrs map[string]interface{}) HclProviderModifier
 	}
 }
 
+// HclProviderWithGenericBlocks sets the generic blocks within the provider
+func HclProviderWithGenericBlocks(blocks ...*hclwrite.Block) HclProviderModifier {
+	return func(p *HclProvider) {
+		p.blocks = blocks
+	}
+}
+
 type ForEach struct {
 	key   string
 	value map[string]string