Skip to content

Commit

Permalink
Add support partitions in policy data sources
Browse files Browse the repository at this point in the history
  • Loading branch information
@ashenm
ashenm committed Nov 8, 2024
1 parent 80514f5 commit 66d3d10
Show file tree
Hide file tree
Showing 16 changed files with 330 additions and 44 deletions.
17 changes: 17 additions & 0 deletions aws/constants.go
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
package aws

var AwsConfig = map[string]map[string]string{
"aws": {
"accountId": "414351767826",
"logDeliveryIamArn": "arn:aws:iam::414351767826:role/SaasUsageDeliveryRole-prod-IAMRole-3PLHICCRR1TK",
"unityCatalogueIamArn": "arn:aws:iam::414351767826:role/unity-catalog-prod-UCMasterRole-14S5ZJVKOTYTL",
},
"aws-us-gov": {
"accountId": "044793339203",
"logDeliveryIamArn": "arn:aws-us-gov:iam::044793339203:role/SaasUsageDeliveryRole-prod-aws-gov-IAMRole-L4QM0RCHYQ1G",
"unityCatalogueIamArn": "arn:aws-us-gov:iam::044793339203:role/unity-catalog-prod-UCMasterRole-1QRFA8SGY15OJ",
},
}

var AwsPartitions = []string{"aws", "aws-us-gov"}
var AwsPartitionsValidationError = "aws_partition must be either 'aws' or 'aws-us-gov'"
26 changes: 19 additions & 7 deletions aws/data_aws_assume_role_policy.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (

"github.com/databricks/terraform-provider-databricks/common"
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
)

type awsIamPolicy struct {
Expand All @@ -31,6 +32,13 @@ func DataAwsAssumeRolePolicy() common.Resource {
return common.Resource{
Read: func(ctx context.Context, d *schema.ResourceData, m *common.DatabricksClient) error {
externalID := d.Get("external_id").(string)
awsPartition := d.Get("aws_partition").(string)
databricksAwsAccountId := d.Get("databricks_account_id").(string)

if databricksAwsAccountId == "" {
databricksAwsAccountId = AwsConfig[awsPartition]["accountId"]
}

policy := awsIamPolicy{
Version: "2012-10-17",
Statements: []*awsIamPolicyStatement{
Expand All @@ -43,16 +51,14 @@ func DataAwsAssumeRolePolicy() common.Resource {
},
},
Principal: map[string]string{
"AWS": fmt.Sprintf("arn:aws:iam::%s:root", d.Get("databricks_account_id").(string)),
"AWS": fmt.Sprintf("arn:%s:iam::%s:root", awsPartition, databricksAwsAccountId),
},
},
},
}
if v, ok := d.GetOk("for_log_delivery"); ok {
if v.(bool) {
// this is production UsageDelivery IAM role, that is considered a constant
logDeliveryARN := "arn:aws:iam::414351767826:role/SaasUsageDeliveryRole-prod-IAMRole-3PLHICCRR1TK"
policy.Statements[0].Principal["AWS"] = logDeliveryARN
policy.Statements[0].Principal["AWS"] = AwsConfig[awsPartition]["logDeliveryIamArn"]
}
}
policyJSON, err := json.MarshalIndent(policy, "", " ")
Expand All @@ -65,10 +71,16 @@ func DataAwsAssumeRolePolicy() common.Resource {
return nil
},
Schema: map[string]*schema.Schema{
"aws_partition": {
Type: schema.TypeString,
Optional: true,
ValidateFunc: validation.StringInSlice(AwsPartitions, false),
Default: "aws",
},
"databricks_account_id": {
Type: schema.TypeString,
Default: "414351767826",
Optional: true,
Type: schema.TypeString,
Optional: true,
Deprecated: "databricks_account_id will be will be removed in the next major release.",
},
"for_log_delivery": {
Type: schema.TypeBool,
Expand Down
49 changes: 49 additions & 0 deletions aws/data_aws_assume_role_policy_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,3 +19,52 @@ func TestDataAwsAssumeRolePolicy(t *testing.T) {
j := d.Get("json")
assert.Lenf(t, j, 299, "Strange length for policy: %s", j)
}

func TestDataAwsAssumeRolePolicyGov(t *testing.T) {
d, err := qa.ResourceFixture{
Read: true,
Resource: DataAwsAssumeRolePolicy(),
NonWritable: true,
ID: ".",
HCL: `
aws_partition = "aws-us-gov"
external_id = "abc"
`,
}.Apply(t)
assert.NoError(t, err)
j := d.Get("json")
assert.Lenf(t, j, 306, "Strange length for policy: %s", j)
}

func TestDataAwsAssumeRolePolicyLogDelivery(t *testing.T) {
d, err := qa.ResourceFixture{
Read: true,
Resource: DataAwsAssumeRolePolicy(),
NonWritable: true,
ID: ".",
HCL: `
external_id = "abc"
for_log_delivery = true
`,
}.Apply(t)
assert.NoError(t, err)
j := d.Get("json")
assert.Lenf(t, j, 347, "Strange length for policy: %s", j)
}

func TestDataAwsAssumeRolePolicyLogDeliveryGov(t *testing.T) {
d, err := qa.ResourceFixture{
Read: true,
Resource: DataAwsAssumeRolePolicy(),
NonWritable: true,
ID: ".",
HCL: `
aws_partition = "aws-us-gov"
external_id = "abc"
for_log_delivery = true
`,
}.Apply(t)
assert.NoError(t, err)
j := d.Get("json")
assert.Lenf(t, j, 362, "Strange length for policy: %s", j)
}
25 changes: 19 additions & 6 deletions aws/data_aws_bucket_policy.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,13 @@ func DataAwsBucketPolicy() common.Resource {
return common.Resource{
Read: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
bucket := d.Get("bucket").(string)
awsPartition := d.Get("aws_partition").(string)
databricksAwsAccountId := AwsConfig[awsPartition]["accountId"]

if databricksAwsAccountId == "" {
databricksAwsAccountId = AwsConfig[awsPartition]["accountId"]
}

policy := awsIamPolicy{
Version: "2012-10-17",
Statements: []*awsIamPolicyStatement{
Expand All @@ -30,11 +37,11 @@ func DataAwsBucketPolicy() common.Resource {
"s3:GetBucketLocation",
},
Resources: []string{
fmt.Sprintf("arn:aws:s3:::%s/*", bucket),
fmt.Sprintf("arn:aws:s3:::%s", bucket),
fmt.Sprintf("arn:%s:s3:::%s/*", awsPartition, bucket),
fmt.Sprintf("arn:%s:s3:::%s", awsPartition, bucket),
},
Principal: map[string]string{
"AWS": fmt.Sprintf("arn:aws:iam::%s:root", d.Get("databricks_account_id").(string)),
"AWS": fmt.Sprintf("arn:%s:iam::%s:root", awsPartition, databricksAwsAccountId),
},
},
},
Expand All @@ -60,10 +67,16 @@ func DataAwsBucketPolicy() common.Resource {
return nil
},
Schema: map[string]*schema.Schema{
"aws_partition": {
Type: schema.TypeString,
Optional: true,
ValidateFunc: validation.StringInSlice(AwsPartitions, false),
Default: "aws",
},
"databricks_account_id": {
Type: schema.TypeString,
Default: "414351767826",
Optional: true,
Type: schema.TypeString,
Optional: true,
Deprecated: "databricks_account_id will be will be removed in the next major release.",
},
"databricks_e2_account_id": {
Type: schema.TypeString,
Expand Down
16 changes: 16 additions & 0 deletions aws/data_aws_bucket_policy_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,3 +53,19 @@ func TestDataAwsBucketPolicyConfusedDeputyProblem(t *testing.T) {
j := d.Get("json")
assert.Lenf(t, j, 575, "Strange length for policy: %s", j)
}

func TestDataAwsBucketPolicyPartitionGov(t *testing.T) {
d, err := qa.ResourceFixture{
Read: true,
Resource: DataAwsBucketPolicy(),
NonWritable: true,
ID: ".",
HCL: `
bucket = "abc"
aws_partition = "aws-us-gov"
`,
}.Apply(t)
assert.NoError(t, err)
j := d.Get("json")
assert.Lenf(t, j, 461, "Strange length for policy: %s", j)
}
51 changes: 29 additions & 22 deletions aws/data_aws_crossaccount_policy.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package aws
import (
"context"
"encoding/json"
"errors"
"fmt"
"regexp"
"slices"
Expand All @@ -17,11 +18,16 @@ func DataAwsCrossaccountPolicy() common.Resource {
PassRole []string `json:"pass_roles,omitempty"`
JSON string `json:"json" tf:"computed"`
AwsAccountId string `json:"aws_account_id,omitempty"`
AwsPartition string `json:"aws_partition,omitempty" tf:"default:aws"`
VpcId string `json:"vpc_id,omitempty"`
Region string `json:"region,omitempty"`
SecurityGroupId string `json:"security_group_id,omitempty"`
}
return common.NoClientData(func(ctx context.Context, data *AwsCrossAccountPolicy) error {
if !slices.Contains(AwsPartitions, data.AwsPartition) {
return errors.New(AwsPartitionsValidationError)
}

if !slices.Contains([]string{"managed", "customer", "restricted"}, data.PolicyType) {
return fmt.Errorf("policy_type must be either 'managed', 'customer' or 'restricted'")
}
Expand Down Expand Up @@ -145,7 +151,7 @@ func DataAwsCrossaccountPolicy() common.Resource {
"iam:CreateServiceLinkedRole",
"iam:PutRolePolicy",
},
Resources: "arn:aws:iam::*:role/aws-service-role/spot.amazonaws.com/AWSServiceRoleForEC2Spot",
Resources: fmt.Sprintf("arn:%s:iam::*:role/aws-service-role/spot.amazonaws.com/AWSServiceRoleForEC2Spot", data.AwsPartition),
Condition: map[string]map[string]string{
"StringLike": {
"iam:AWSServiceName": "spot.amazonaws.com",
Expand All @@ -168,6 +174,7 @@ func DataAwsCrossaccountPolicy() common.Resource {
if data.PolicyType == "restricted" {
region := data.Region
aws_account_id := data.AwsAccountId
awsPartition := data.AwsPartition
vpc_id := data.VpcId
security_group_id := data.SecurityGroupId
policy.Statements = append(policy.Statements,
Expand All @@ -179,7 +186,7 @@ func DataAwsCrossaccountPolicy() common.Resource {
"ec2:DisassociateIamInstanceProfile",
"ec2:ReplaceIamInstanceProfileAssociation",
},
Resources: fmt.Sprintf("arn:aws:ec2:%s:%s:instance/*", region, aws_account_id),
Resources: fmt.Sprintf("arn:%s:ec2:%s:%s:instance/*", awsPartition, region, aws_account_id),
Condition: map[string]map[string]string{
"StringEquals": {
"ec2:ResourceTag/Vendor": "Databricks",
Expand All @@ -191,8 +198,8 @@ func DataAwsCrossaccountPolicy() common.Resource {
Effect: "Allow",
Actions: "ec2:RunInstances",
Resources: []string{
fmt.Sprintf("arn:aws:ec2:%s:%s:volume/*", region, aws_account_id),
fmt.Sprintf("arn:aws:ec2:%s:%s:instance/*", region, aws_account_id),
fmt.Sprintf("arn:%s:ec2:%s:%s:volume/*", awsPartition, region, aws_account_id),
fmt.Sprintf("arn:%s:ec2:%s:%s:instance/*", awsPartition, region, aws_account_id),
},
Condition: map[string]map[string]string{
"StringEquals": {
Expand All @@ -204,7 +211,7 @@ func DataAwsCrossaccountPolicy() common.Resource {
Sid: "AllowEc2RunInstanceImagePerTag",
Effect: "Allow",
Actions: "ec2:RunInstances",
Resources: fmt.Sprintf("arn:aws:ec2:%s:%s:image/*", region, aws_account_id),
Resources: fmt.Sprintf("arn:%s:ec2:%s:%s:image/*", awsPartition, region, aws_account_id),
Condition: map[string]map[string]string{
"StringEquals": {
"aws:ResourceTag/Vendor": "Databricks",
Expand All @@ -216,13 +223,13 @@ func DataAwsCrossaccountPolicy() common.Resource {
Effect: "Allow",
Actions: "ec2:RunInstances",
Resources: []string{
fmt.Sprintf("arn:aws:ec2:%s:%s:network-interface/*", region, aws_account_id),
fmt.Sprintf("arn:aws:ec2:%s:%s:subnet/*", region, aws_account_id),
fmt.Sprintf("arn:aws:ec2:%s:%s:security-group/*", region, aws_account_id),
fmt.Sprintf("arn:%s:ec2:%s:%s:network-interface/*", awsPartition, region, aws_account_id),
fmt.Sprintf("arn:%s:ec2:%s:%s:subnet/*", awsPartition, region, aws_account_id),
fmt.Sprintf("arn:%s:ec2:%s:%s:security-group/*", awsPartition, region, aws_account_id),
},
Condition: map[string]map[string]string{
"StringEquals": {
"ec2:vpc": fmt.Sprintf("arn:aws:ec2:%s:%s:vpc/%s", region, aws_account_id, vpc_id),
"ec2:vpc": fmt.Sprintf("arn:%s:ec2:%s:%s:vpc/%s", awsPartition, region, aws_account_id, vpc_id),
},
},
},
Expand All @@ -231,19 +238,19 @@ func DataAwsCrossaccountPolicy() common.Resource {
Effect: "Allow",
Actions: "ec2:RunInstances",
NotResources: []string{
fmt.Sprintf("arn:aws:ec2:%s:%s:image/*", region, aws_account_id),
fmt.Sprintf("arn:aws:ec2:%s:%s:network-interface/*", region, aws_account_id),
fmt.Sprintf("arn:aws:ec2:%s:%s:subnet/*", region, aws_account_id),
fmt.Sprintf("arn:aws:ec2:%s:%s:security-group/*", region, aws_account_id),
fmt.Sprintf("arn:aws:ec2:%s:%s:volume/*", region, aws_account_id),
fmt.Sprintf("arn:aws:ec2:%s:%s:instance/*", region, aws_account_id),
fmt.Sprintf("arn:%s:ec2:%s:%s:image/*", awsPartition, region, aws_account_id),
fmt.Sprintf("arn:%s:ec2:%s:%s:network-interface/*", awsPartition, region, aws_account_id),
fmt.Sprintf("arn:%s:ec2:%s:%s:subnet/*", awsPartition, region, aws_account_id),
fmt.Sprintf("arn:%s:ec2:%s:%s:security-group/*", awsPartition, region, aws_account_id),
fmt.Sprintf("arn:%s:ec2:%s:%s:volume/*", awsPartition, region, aws_account_id),
fmt.Sprintf("arn:%s:ec2:%s:%s:instance/*", awsPartition, region, aws_account_id),
},
},
&awsIamPolicyStatement{
Sid: "EC2TerminateInstancesTag",
Effect: "Allow",
Actions: "ec2:TerminateInstances",
Resources: fmt.Sprintf("arn:aws:ec2:%s:%s:instance/*", region, aws_account_id),
Resources: fmt.Sprintf("arn:%s:ec2:%s:%s:instance/*", awsPartition, region, aws_account_id),
Condition: map[string]map[string]string{
"StringEquals": {
"ec2:ResourceTag/Vendor": "Databricks",
Expand All @@ -258,8 +265,8 @@ func DataAwsCrossaccountPolicy() common.Resource {
"ec2:DetachVolume",
},
Resources: []string{
fmt.Sprintf("arn:aws:ec2:%s:%s:instance/*", region, aws_account_id),
fmt.Sprintf("arn:aws:ec2:%s:%s:volume/*", region, aws_account_id),
fmt.Sprintf("arn:%s:ec2:%s:%s:instance/*", awsPartition, region, aws_account_id),
fmt.Sprintf("arn:%s:ec2:%s:%s:volume/*", awsPartition, region, aws_account_id),
},
Condition: map[string]map[string]string{
"StringEquals": {
Expand All @@ -271,7 +278,7 @@ func DataAwsCrossaccountPolicy() common.Resource {
Sid: "EC2CreateVolumeByTag",
Effect: "Allow",
Actions: "ec2:CreateVolume",
Resources: fmt.Sprintf("arn:aws:ec2:%s:%s:volume/*", region, aws_account_id),
Resources: fmt.Sprintf("arn:%s:ec2:%s:%s:volume/*", awsPartition, region, aws_account_id),
Condition: map[string]map[string]string{
"StringEquals": {
"aws:RequestTag/Vendor": "Databricks",
Expand All @@ -283,7 +290,7 @@ func DataAwsCrossaccountPolicy() common.Resource {
Effect: "Allow",
Actions: "ec2:DeleteVolume",
Resources: []string{
fmt.Sprintf("arn:aws:ec2:%s:%s:volume/*", region, aws_account_id),
fmt.Sprintf("arn:%s:ec2:%s:%s:volume/*", awsPartition, region, aws_account_id),
},
Condition: map[string]map[string]string{
"StringEquals": {
Expand All @@ -300,10 +307,10 @@ func DataAwsCrossaccountPolicy() common.Resource {
"ec2:RevokeSecurityGroupEgress",
"ec2:RevokeSecurityGroupIngress",
},
Resources: fmt.Sprintf("arn:aws:ec2:%s:%s:security-group/%s", region, aws_account_id, security_group_id),
Resources: fmt.Sprintf("arn:%s:ec2:%s:%s:security-group/%s", awsPartition, region, aws_account_id, security_group_id),
Condition: map[string]map[string]string{
"StringEquals": {
"ec2:vpc": fmt.Sprintf("arn:aws:ec2:%s:%s:vpc/%s", region, aws_account_id, vpc_id),
"ec2:vpc": fmt.Sprintf("arn:%s:ec2:%s:%s:vpc/%s", awsPartition, region, aws_account_id, vpc_id),
},
},
},
Expand Down
Loading

0 comments on commit 66d3d10

Please sign in to comment.