Skip to content

Commit

Permalink
Add testing
Browse files Browse the repository at this point in the history
Signed-off-by: Kevin DeJong <[email protected]>
  • Loading branch information
kddejong committed Sep 4, 2019
1 parent 94786cc commit 9d78187
Show file tree
Hide file tree
Showing 5 changed files with 380 additions and 68 deletions.
5 changes: 5 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -1,3 +1,8 @@
## vNext

- Added variable for specifying an IAM policy template (GO Template)
- Update IAM Policy for the principal every time the account is unlocked

## v0.14.0

- Added rds backup delete to nuke
Expand Down
125 changes: 69 additions & 56 deletions cmd/lambda/update_redbox_principal_policy/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import (
"strings"

"github.com/Optum/Redbox/pkg/api/response"
"github.com/Optum/Redbox/pkg/awsiface"
"github.com/Optum/Redbox/pkg/common"
"github.com/Optum/Redbox/pkg/db"
"github.com/Optum/Redbox/pkg/rolemanager"
Expand All @@ -27,12 +28,13 @@ func main() {

func handler(ctx context.Context, snsEvent events.SNSEvent) error {
dbSvc, err := db.NewFromEnv()
awsSession := newAWSSession()
awsSession := session.Must(session.NewSession())
tokenSvc := common.STS{Client: sts.New(awsSession)}
s3Svc := common.S3{
Client: s3.New(awsSession),
Manager: s3manager.NewDownloader(awsSession),
}
roleManagerSvc := &rolemanager.IAMPolicyManager{}

if err != nil {
log.Printf("Unable to setup DB Service: %s", err.Error())
Expand All @@ -48,71 +50,82 @@ func handler(ctx context.Context, snsEvent events.SNSEvent) error {
log.Printf("Failed to read SNS message %s: %s", snsRecord.Message, err.Error())
return err
}
account, err := dbSvc.GetAccount(lease.AccountID)
if err != nil {
log.Printf("Failed to get the redbox account %s: %s",
lease.AccountID, err.Error())
return err
}

accountRes := response.AccountResponse(*account)
fmt.Printf("AccountID = %s \nRoleToAssume = %s\n", lease.AccountID, accountRes.AdminRoleArn)
principalIAMDenyTags := strings.Split(common.RequireEnv("PRINCIPAL_IAM_DENY_TAGS"), ",")
principalRoleName := common.RequireEnv("PRINCIPAL_ROLE_NAME")
policyName := common.RequireEnv("PRINCIPAL_POLICY_NAME")
principalPolicyArn, err := arn.Parse(fmt.Sprintf("arn:aws:iam::%s:policy/%s", lease.AccountID, policyName))
if err != nil {
log.Printf("Failed to parse ARN 'arn:aws:iam::%s:policy/%s': %s", lease.AccountID, policyName, err.Error())
return err
}

policy, err := getPolicy(s3Svc, getPolicyInput{
PrincipalPolicyArn: principalPolicyArn.String(),
PrincipalRoleArn: fmt.Sprintf("arn:aws:iam::%s:role/%s", lease.AccountID, principalRoleName),
PrincipalIAMDenyTags: principalIAMDenyTags,
AdminRoleArn: accountRes.AdminRoleArn,
})

// Assume role into the new Redbox account
accountSession, err := tokenSvc.NewSession(awsSession, accountRes.AdminRoleArn)
if err != nil {
log.Printf("Failed to assume role '%s': %s", accountRes.AdminRoleArn, err.Error())
return err
}
iamSvc := iam.New(accountSession)

// Create the Role + Policy
roleManager := &rolemanager.IAMPolicyManager{}
roleManager.SetIAMClient(iamSvc)
return roleManager.MergePolicy(&rolemanager.MergePolicyInput{
PolicyName: policyName,
PolicyArn: principalPolicyArn,
PolicyDocument: policy,
PolicyDescription: "", // Policy should already exist so this will be ignored
err = processRecord(processRecordInput{
AccountID: lease.AccountID,
DbSvc: dbSvc,
StoragerSvc: s3Svc,
TokenSvc: tokenSvc,
AwsSession: awsSession,
RoleManager: roleManagerSvc,
PrincipalRoleName: common.RequireEnv("PRINCIPAL_ROLE_NAME"),
PrincipalPolicyName: common.RequireEnv("PRINCIPAL_POLICY_NAME"),
PrincipalIAMDenyTags: strings.Split(common.RequireEnv("PRINCIPAL_IAM_DENY_TAGS"), ","),
PolicyBucket: common.RequireEnv("ARTIFACTS_BUCKET"),
PolicyBucketKey: common.RequireEnv("PRINCIPAL_POLICY_S3_KEY"),
})
}
return nil
}

type getPolicyInput struct {
PrincipalPolicyArn string
PrincipalRoleArn string
type processRecordInput struct {
AccountID string
DbSvc db.DBer
StoragerSvc common.Storager
TokenSvc common.TokenService
AwsSession awsiface.AwsSession
RoleManager rolemanager.PolicyManager
PrincipalRoleName string
PrincipalPolicyName string
PrincipalIAMDenyTags []string
AdminRoleArn string
PolicyBucket string
PolicyBucketKey string
}

func getPolicy(storage common.S3, input getPolicyInput) (string, error) {
bucket := common.RequireEnv("ARTIFACTS_BUCKET")
key := common.RequireEnv("PRINCIPAL_POLICY_S3_KEY")
policy, err := storage.GetTemplateObject(bucket, key, input)
return policy, err
}
func processRecord(input processRecordInput) error {

account, err := input.DbSvc.GetAccount(input.AccountID)
if err != nil {
log.Printf("Failed to get the redbox account %s: %s",
input.AccountID, err.Error())
return err
}

accountRes := response.AccountResponse(*account)

principalPolicyArn, err := arn.Parse(fmt.Sprintf("arn:aws:iam::%s:policy/%s", input.AccountID, input.PrincipalPolicyName))
if err != nil {
log.Printf("Failed to parse ARN 'arn:aws:iam::%s:policy/%s': %s", input.AccountID, input.PrincipalPolicyName, err.Error())
return err
}

policy, err := input.StoragerSvc.GetTemplateObject(input.PolicyBucket, input.PolicyBucketKey, getPolicyInput{
PrincipalPolicyArn: principalPolicyArn.String(),
PrincipalRoleArn: fmt.Sprintf("arn:aws:iam::%s:role/%s", input.AccountID, input.PrincipalRoleName),
PrincipalIAMDenyTags: input.PrincipalIAMDenyTags,
AdminRoleArn: accountRes.AdminRoleArn,
})

func newAWSSession() *session.Session {
awsSession, err := session.NewSession()
// Assume role into the new Redbox account
accountSession, err := input.TokenSvc.NewSession(input.AwsSession, accountRes.AdminRoleArn)
if err != nil {
errorMessage := fmt.Sprintf("Failed to create AWS session: %s", err)
log.Fatal(errorMessage)
log.Printf("Failed to assume role '%s': %s", accountRes.AdminRoleArn, err.Error())
return err
}
return awsSession
iamSvc := iam.New(accountSession)

// Create the Role + Policy
input.RoleManager.SetIAMClient(iamSvc)
return input.RoleManager.MergePolicy(&rolemanager.MergePolicyInput{
PolicyName: input.PrincipalPolicyName,
PolicyArn: principalPolicyArn,
PolicyDocument: policy,
})
}

type getPolicyInput struct {
PrincipalPolicyArn string
PrincipalRoleArn string
PrincipalIAMDenyTags []string
AdminRoleArn string
}
107 changes: 107 additions & 0 deletions cmd/lambda/update_redbox_principal_policy/main_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,107 @@
package main

import (
"fmt"
"testing"

"github.com/Optum/Redbox/pkg/rolemanager"

awsMocks "github.com/Optum/Redbox/pkg/awsiface/mocks"
commonmock "github.com/Optum/Redbox/pkg/common/mocks"
"github.com/Optum/Redbox/pkg/db"
dbmock "github.com/Optum/Redbox/pkg/db/mocks"
roleMock "github.com/Optum/Redbox/pkg/rolemanager/mocks"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/aws/arn"
"github.com/aws/aws-sdk-go/aws/client"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
)

// testTransitionFinanceLockInput is the structured input for testing the helper
// function transitionFinanceLock
type testUpdateRedboxPrincipalPolicy struct {
ExpectedError error
GetAccountResult *db.RedboxAccount
GetAccountError error
TransitionLeaseStatusError error
PrincipalPolicyName string
PrincipalRoleName string
PrincipalIAMDenyTags []string
StoragerPolicy string
StoragerError error
}

func TestUpdateRedboxPrincipalPolicy(t *testing.T) {

tests := []testUpdateRedboxPrincipalPolicy{
// Happy Path FinanceLock
{
GetAccountResult: &db.RedboxAccount{
ID: "123456789012",
AdminRoleArn: "arn:aws:iam::123456789012:role/RedBoxAdminRole",
},
PrincipalPolicyName: "RedboxPrincipalPolicy",
PrincipalRoleName: "RedboxPrincipalRole",
PrincipalIAMDenyTags: []string{"Redbox"},
StoragerPolicy: "{\"Test\" : \"Policy\"}",
},
}

// Iterate through each test in the list
for _, test := range tests {
// Setup mocks
mockDB := dbmock.DBer{}
mockDB.On("GetAccount", mock.Anything).Return(
test.GetAccountResult,
test.GetAccountError)
mockS3 := &commonmock.Storager{}
mockS3.On("GetTemplateObject", mock.Anything, mock.Anything, getPolicyInput{
PrincipalPolicyArn: fmt.Sprintf("arn:aws:iam::%s:policy/%s", test.GetAccountResult.ID, test.PrincipalPolicyName),
PrincipalRoleArn: fmt.Sprintf("arn:aws:iam::%s:role/%s", test.GetAccountResult.ID, test.PrincipalRoleName),
PrincipalIAMDenyTags: test.PrincipalIAMDenyTags,
AdminRoleArn: test.GetAccountResult.AdminRoleArn,
}).Return(
test.StoragerPolicy,
test.StoragerError,
)
mockAdminRoleSession := &awsMocks.AwsSession{}
mockAdminRoleSession.On("ClientConfig", mock.Anything).Return(client.Config{
Config: &aws.Config{},
})
mockToken := &commonmock.TokenService{}
mockToken.On("NewSession", mock.Anything, test.GetAccountResult.AdminRoleArn).
Return(mockAdminRoleSession, nil)
mockToken.On("AssumeRole", mock.Anything).Return(nil, nil)
mockSession := &awsMocks.AwsSession{}

mockRoleManager := &roleMock.PolicyManager{}
mockRoleManager.On("SetIAMClient", mock.Anything).Return()
policyArn, _ := arn.Parse(fmt.Sprintf("arn:aws:iam::%s:policy/%s", test.GetAccountResult.ID, test.PrincipalPolicyName))
mockRoleManager.On("MergePolicy", &rolemanager.MergePolicyInput{
PolicyArn: policyArn,
PolicyName: test.PrincipalPolicyName,
PolicyDocument: test.StoragerPolicy,
}).Return(nil)

// Call transitionFinanceLock
err := processRecord(processRecordInput{
AccountID: test.GetAccountResult.ID,
DbSvc: &mockDB,
StoragerSvc: mockS3,
TokenSvc: mockToken,
AwsSession: mockSession,
RoleManager: mockRoleManager,
PrincipalRoleName: test.PrincipalRoleName,
PrincipalPolicyName: test.PrincipalPolicyName,
PrincipalIAMDenyTags: test.PrincipalIAMDenyTags,
})

// Assert expectations
if test.ExpectedError != nil {
require.Equal(t, test.ExpectedError.Error(), err.Error())
} else {
require.Nil(t, err)
}
}
}
Loading

0 comments on commit 9d78187

Please sign in to comment.