feat: initial draft for multi-region support

steadybit · Nov 27, 2024 · 271c619 · 271c619
1 parent 4f18c9f
commit 271c619
Show file tree

Hide file tree

Showing 36 changed files with 228 additions and 260 deletions.
diff --git a/extaz/availability_zone_attack_blackhole.go b/extaz/availability_zone_attack_blackhole.go
@@ -58,7 +58,7 @@ type azBlackholeImdsApi interface {
 func NewAzBlackholeAction() action_kit_sdk.Action[BlackholeState] {
 	return &azBlackholeAction{
 		clientProvider:             defaultClientProvider,
-		extensionRootAccountNumber: utils.Accounts.GetRootAccountNumber(),
+		extensionRootAccountNumber: utils.GetRootAccountNumber(),
 	}
 }
 
@@ -469,12 +469,12 @@ func getAllNACLsCreatedBySteadybit(clientEc2 azBlackholeEC2Api, ctx context.Cont
 }
 
 func defaultClientProvider(account string, region string) (azBlackholeEC2Api, azBlackholeImdsApi, error) {
-	awsAccount, err := utils.Accounts.GetAccount(account, region)
+	awsAccess, err := utils.GetAwsAccess(account, region)
 	if err != nil {
 		return nil, nil, err
 	}
-	clientEc2 := ec2.NewFromConfig(awsAccount.AwsConfig)
-	clientImds := imds.NewFromConfig(awsAccount.AwsConfig)
+	clientEc2 := ec2.NewFromConfig(awsAccess.AwsConfig)
+	clientImds := imds.NewFromConfig(awsAccess.AwsConfig)
 	if err != nil {
 		return nil, nil, err
 	}

diff --git a/extaz/availablity_zone_discovery.go b/extaz/availablity_zone_discovery.go
@@ -66,7 +66,7 @@ func (a *azDiscovery) DescribeTarget() discovery_kit_api.TargetDescription {
 }
 
 func (a *azDiscovery) DiscoverTargets(ctx context.Context) ([]discovery_kit_api.Target, error) {
-	return utils.ForEveryAccount(utils.Accounts, getTargetsForAccount, ctx, "availability zone")
+	return utils.ForEveryConfiguredAwsAccess(getTargetsForAccount, ctx, "availability zone")
 }
 
 func getTargetsForAccount(account *utils.AwsAccess, _ context.Context) ([]discovery_kit_api.Target, error) {

diff --git a/extaz/availablity_zone_discovery_test.go b/extaz/availablity_zone_discovery_test.go
@@ -16,8 +16,8 @@ type zoneMock struct {
 	mock.Mock
 }
 
-func (m *zoneMock) GetZones(awsAccountNumber string) []types.AvailabilityZone {
-	args := m.Called(awsAccountNumber)
+func (m *zoneMock) GetZones(awsAccountNumber string, region string) []types.AvailabilityZone {
+	args := m.Called(awsAccountNumber, region)
 	return args.Get(0).([]types.AvailabilityZone)
 }
 
@@ -31,10 +31,10 @@ func TestGetAllAvailabilityZones(t *testing.T) {
 			ZoneId:     discovery_kit_api.Ptr("euc1-az3"),
 		},
 	}
-	mockedApi.On("GetZones", mock.Anything).Return(mockedReturnValue)
+	mockedApi.On("GetZones", mock.Anything, mock.Anything).Return(mockedReturnValue)
 
 	// When
-	targets := getAllAvailabilityZones(mockedApi, "42")
+	targets := getAllAvailabilityZones(mockedApi, "42", "eu-central-1")
 
 	// Then
 	assert.Equal(t, 1, len(targets))

diff --git a/extec2/instance_attack_state.go b/extec2/instance_attack_state.go
@@ -152,9 +152,9 @@ func (e *ec2InstanceStateAction) Start(ctx context.Context, state *InstanceState
 }
 
 func defaultClientProvider(account string, region string) (ec2InstanceStateChangeApi, error) {
-	awsAccount, err := utils.Accounts.GetAccount(account, region)
+	awsAccess, err := utils.GetAwsAccess(account, region)
 	if err != nil {
 		return nil, err
 	}
-	return ec2.NewFromConfig(awsAccount.AwsConfig), nil
+	return ec2.NewFromConfig(awsAccess.AwsConfig), nil
 }
diff --git a/extec2/instance_discovery.go b/extec2/instance_discovery.go
@@ -249,7 +249,7 @@ func (e *ec2Discovery) DescribeAttributes() []discovery_kit_api.AttributeDescrip
 }
 
 func (e *ec2Discovery) DiscoverTargets(ctx context.Context) ([]discovery_kit_api.Target, error) {
-	return utils.ForEveryAccount(utils.Accounts, getTargetsForAccount, ctx, "ec2-instance")
+	return utils.ForEveryConfiguredAwsAccess(getTargetsForAccount, ctx, "ec2-instance")
 }
 
 func getTargetsForAccount(account *utils.AwsAccess, ctx context.Context) ([]discovery_kit_api.Target, error) {

diff --git a/extec2/instance_discovery_test.go b/extec2/instance_discovery_test.go
@@ -32,8 +32,8 @@ type zoneMock struct {
 	mock.Mock
 }
 
-func (m *zoneMock) GetZone(awsAccountNumber string, awsZone string) *types.AvailabilityZone {
-	args := m.Called(awsAccountNumber, awsZone)
+func (m *zoneMock) GetZone(awsAccountNumber string, awsZone string, region string) *types.AvailabilityZone {
+	args := m.Called(awsAccountNumber, awsZone, region)
 	return args.Get(0).(*types.AvailabilityZone)
 }
 
@@ -76,7 +76,7 @@ func TestGetAllEc2Instances(t *testing.T) {
 		RegionName: discovery_kit_api.Ptr("us-east-1"),
 		ZoneId:     discovery_kit_api.Ptr("us-east-1b-id"),
 	}
-	mockedZoneUtil.On("GetZone", mock.Anything, mock.Anything).Return(&mockedZone)
+	mockedZoneUtil.On("GetZone", mock.Anything, mock.Anything, mock.Anything).Return(&mockedZone)
 
 	// When
 	targets, err := GetAllEc2Instances(context.Background(), mockedApi, mockedZoneUtil, "42", "us-east-1")
@@ -126,7 +126,7 @@ func TestGetAllEc2InstancesWithFilteredAttributes(t *testing.T) {
 		RegionName: discovery_kit_api.Ptr("us-east-1"),
 		ZoneId:     discovery_kit_api.Ptr("us-east-1b-id"),
 	}
-	mockedZoneUtil.On("GetZone", mock.Anything, mock.Anything).Return(&mockedZone)
+	mockedZoneUtil.On("GetZone", mock.Anything, mock.Anything, mock.Anything).Return(&mockedZone)
 
 	// When
 	targets, err := GetAllEc2Instances(context.Background(), mockedApi, mockedZoneUtil, "42", "us-east-1")
@@ -179,7 +179,7 @@ func TestNameNotSet(t *testing.T) {
 		RegionName: discovery_kit_api.Ptr("us-east-1"),
 		ZoneId:     discovery_kit_api.Ptr("us-east-1b-id"),
 	}
-	mockedZoneUtil.On("GetZone", mock.Anything, mock.Anything).Return(&mockedZone)
+	mockedZoneUtil.On("GetZone", mock.Anything, mock.Anything, mock.Anything).Return(&mockedZone)
 
 	// When
 	targets, err := GetAllEc2Instances(context.Background(), mockedApi, mockedZoneUtil, "42", "us-east-1")
@@ -204,7 +204,7 @@ func TestGetAllEc2InstancesError(t *testing.T) {
 		RegionName: discovery_kit_api.Ptr("us-east-1"),
 		ZoneId:     discovery_kit_api.Ptr("us-east-1b-id"),
 	}
-	mockedZoneUtil.On("GetZone", mock.Anything, mock.Anything).Return(&mockedZone)
+	mockedZoneUtil.On("GetZone", mock.Anything, mock.Anything, mock.Anything).Return(&mockedZone)
 
 	// When
 	_, err := GetAllEc2Instances(context.Background(), mockedApi, mockedZoneUtil, "42", "us-east-1")

diff --git a/extecs/service_attack_scale.go b/extecs/service_attack_scale.go
@@ -145,9 +145,9 @@ func (e *ecsServiceScaleAction) Stop(ctx context.Context, state *ServiceScaleSta
 }
 
 func defaultClientProviderService(account string, region string) (ecsServiceScaleApi, error) {
-	awsAccount, err := utils.Accounts.GetAccount(account, region)
+	awsAccess, err := utils.GetAwsAccess(account, region)
 	if err != nil {
 		return nil, err
 	}
-	return ecs.NewFromConfig(awsAccount.AwsConfig), nil
+	return ecs.NewFromConfig(awsAccess.AwsConfig), nil
 }
diff --git a/extecs/service_description_poller.go b/extecs/service_description_poller.go
@@ -214,9 +214,9 @@ func (p EcsServiceDescriptionPoller) pollAll(ctx context.Context) {
 }
 
 func defaultDescribeServiceProvider(account string, region string) (ecsDescribeServicesApi, error) {
-	awsAccount, err := utils.Accounts.GetAccount(account, region)
+	awsAccess, err := utils.GetAwsAccess(account, region)
 	if err != nil {
 		return nil, err
 	}
-	return ecs.NewFromConfig(awsAccount.AwsConfig), nil
+	return ecs.NewFromConfig(awsAccess.AwsConfig), nil
 }
diff --git a/extecs/service_discovery.go b/extecs/service_discovery.go
@@ -104,7 +104,7 @@ func (e *ecsServiceDiscovery) DescribeAttributes() []discovery_kit_api.Attribute
 }
 
 func (e *ecsServiceDiscovery) DiscoverTargets(ctx context.Context) ([]discovery_kit_api.Target, error) {
-	return utils.ForEveryAccount(utils.Accounts, getEcsServicesForAccount, ctx, "ecs-service")
+	return utils.ForEveryConfiguredAwsAccess(getEcsServicesForAccount, ctx, "ecs-service")
 }
 
 func getEcsServicesForAccount(account *utils.AwsAccess, ctx context.Context) ([]discovery_kit_api.Target, error) {

diff --git a/extecs/task_attack_ssm.go b/extecs/task_attack_ssm.go
@@ -320,9 +320,9 @@ func (e *ecsTaskSsmAction) findManagedInstance(ctx context.Context, client ecsTa
 }
 
 func defaultTaskSsmClientProvider(account string, region string) (ecsTaskSsmApi, error) {
-	awsAccount, err := utils.Accounts.GetAccount(account, region)
+	awsAccess, err := utils.GetAwsAccess(account, region)
 	if err != nil {
 		return nil, err
 	}
-	return ssm.NewFromConfig(awsAccount.AwsConfig), nil
+	return ssm.NewFromConfig(awsAccess.AwsConfig), nil
 }
diff --git a/extecs/task_attack_stop.go b/extecs/task_attack_stop.go
@@ -112,9 +112,9 @@ func (e *ecsTaskStopAction) Start(ctx context.Context, state *TaskStopState) (*a
 }
 
 func defaultTaskStopClientProvider(account string, region string) (ecsTaskStopApi, error) {
-	awsAccount, err := utils.Accounts.GetAccount(account, region)
+	awsAccess, err := utils.GetAwsAccess(account, region)
 	if err != nil {
 		return nil, err
 	}
-	return ecs.NewFromConfig(awsAccount.AwsConfig), nil
+	return ecs.NewFromConfig(awsAccess.AwsConfig), nil
 }
diff --git a/extecs/task_discovery.go b/extecs/task_discovery.go
@@ -107,7 +107,7 @@ func (e *ecsTaskDiscovery) DescribeAttributes() []discovery_kit_api.AttributeDes
 }
 
 func (e *ecsTaskDiscovery) DiscoverTargets(ctx context.Context) ([]discovery_kit_api.Target, error) {
-	return utils.ForEveryAccount(utils.Accounts, getTargetsForAccount, ctx, "ecs-task")
+	return utils.ForEveryConfiguredAwsAccess(getTargetsForAccount, ctx, "ecs-task")
 }
 
 func getTargetsForAccount(account *utils.AwsAccess, ctx context.Context) ([]discovery_kit_api.Target, error) {

diff --git a/extecs/task_discovery_test.go b/extecs/task_discovery_test.go
@@ -47,8 +47,8 @@ type zoneMock struct {
 	mock.Mock
 }
 
-func (m *zoneMock) GetZone(awsAccountNumber string, awsZone string) *ec2types.AvailabilityZone {
-	args := m.Called(awsAccountNumber, awsZone)
+func (m *zoneMock) GetZone(awsAccountNumber string, awsZone string, region string) *ec2types.AvailabilityZone {
+	args := m.Called(awsAccountNumber, awsZone, region)
 	return args.Get(0).(*ec2types.AvailabilityZone)
 }
 
@@ -104,7 +104,7 @@ func TestGetAllEcsTasks(t *testing.T) {
 		RegionName: discovery_kit_api.Ptr("us-east-1"),
 		ZoneId:     discovery_kit_api.Ptr("us-east-1b-id"),
 	}
-	mockedZoneUtil.On("GetZone", mock.Anything, mock.Anything).Return(&mockedZone)
+	mockedZoneUtil.On("GetZone", mock.Anything, mock.Anything, mock.Anything).Return(&mockedZone)
 
 	// When
 	targets, err := GetAllEcsTasks(context.Background(), mockedApi, mockedZoneUtil, "42", "us-east-1")

diff --git a/extelasticache/common_elasticache.go b/extelasticache/common_elasticache.go
@@ -30,9 +30,9 @@ type ElasticacheApi interface {
 }
 
 func defaultElasticacheClientProvider(account string, region string) (ElasticacheApi, error) {
-	awsAccount, err := utils.Accounts.GetAccount(account, region)
+	awsAccess, err := utils.GetAwsAccess(account, region)
 	if err != nil {
 		return nil, err
 	}
-	return elasticache.NewFromConfig(awsAccount.AwsConfig), nil
+	return elasticache.NewFromConfig(awsAccess.AwsConfig), nil
 }
diff --git a/extelasticache/node_group_discovery.go b/extelasticache/node_group_discovery.go
@@ -126,7 +126,7 @@ func (r *elasticacheReplicationGroupDiscovery) DescribeAttributes() []discovery_
 }
 
 func (r *elasticacheReplicationGroupDiscovery) DiscoverTargets(ctx context.Context) ([]discovery_kit_api.Target, error) {
-	return utils.ForEveryAccount(utils.Accounts, getClusterTargetsForAccount, ctx, "replication-group")
+	return utils.ForEveryConfiguredAwsAccess(getClusterTargetsForAccount, ctx, "replication-group")
 }
 
 func getClusterTargetsForAccount(account *utils.AwsAccess, ctx context.Context) ([]discovery_kit_api.Target, error) {

diff --git a/extelb/alb_attack_static_response.go b/extelb/alb_attack_static_response.go
@@ -596,9 +596,9 @@ func restoreOldPriorities(ctx context.Context, client *albStaticResponseApi, sta
 }
 
 func defaultClientProviderService(account string, region string) (albStaticResponseApi, error) {
-	awsAccount, err := utils.Accounts.GetAccount(account, region)
+	awsAccess, err := utils.GetAwsAccess(account, region)
 	if err != nil {
 		return nil, err
 	}
-	return elasticloadbalancingv2.NewFromConfig(awsAccount.AwsConfig), nil
+	return elasticloadbalancingv2.NewFromConfig(awsAccess.AwsConfig), nil
 }
diff --git a/extelb/alb_discovery.go b/extelb/alb_discovery.go
@@ -99,7 +99,7 @@ func (e *albDiscovery) DescribeAttributes() []discovery_kit_api.AttributeDescrip
 }
 
 func (e *albDiscovery) DiscoverTargets(ctx context.Context) ([]discovery_kit_api.Target, error) {
-	return utils.ForEveryAccount(utils.Accounts, getTargetsForAccount, ctx, "ecs-task")
+	return utils.ForEveryConfiguredAwsAccess(getTargetsForAccount, ctx, "ecs-task")
 }
 
 func getTargetsForAccount(account *utils.AwsAccess, ctx context.Context) ([]discovery_kit_api.Target, error) {

diff --git a/extelb/alb_discovery_test.go b/extelb/alb_discovery_test.go
@@ -48,8 +48,8 @@ type zoneMock struct {
 	mock.Mock
 }
 
-func (m *zoneMock) GetZone(awsAccountNumber string, awsZone string) *ec2types.AvailabilityZone {
-	args := m.Called(awsAccountNumber, awsZone)
+func (m *zoneMock) GetZone(awsAccountNumber string, awsZone string, region string) *ec2types.AvailabilityZone {
+	args := m.Called(awsAccountNumber, awsZone, region)
 	return args.Get(0).(*ec2types.AvailabilityZone)
 }
 
@@ -146,10 +146,10 @@ func TestGetAllAlbTargets(t *testing.T) {
 	}
 	mockedZoneUtil.On("GetZone", mock.Anything, mock.MatchedBy(func(params string) bool {
 		return params == "us-east-1a"
-	})).Return(&mockedZone1a)
+	}), mock.Anything).Return(&mockedZone1a)
 	mockedZoneUtil.On("GetZone", mock.Anything, mock.MatchedBy(func(params string) bool {
 		return params == "us-east-1b"
-	})).Return(&mockedZone1b)
+	}), mock.Anything).Return(&mockedZone1b)
 
 	// When
 	targets, err := GetAlbs(context.Background(), mockedApi, mockedZoneUtil, "42", "us-east-1")

diff --git a/extfis/start_experiment.go b/extfis/start_experiment.go
@@ -100,11 +100,11 @@ func (f FisExperimentAction) Prepare(_ context.Context, state *FisExperimentStat
 
 func (f FisExperimentAction) Start(ctx context.Context, state *FisExperimentState) (*action_kit_api.StartResult, error) {
 	return startExperiment(ctx, state, func(account string, region string) (FisStartExperimentClient, error) {
-		awsAccount, err := utils.Accounts.GetAccount(account, region)
+		awsAccess, err := utils.GetAwsAccess(account, region)
 		if err != nil {
 			return nil, err
 		}
-		return fis.NewFromConfig(awsAccount.AwsConfig), nil
+		return fis.NewFromConfig(awsAccess.AwsConfig), nil
 	})
 }
 
@@ -139,11 +139,11 @@ func startExperiment(ctx context.Context, state *FisExperimentState, clientProvi
 
 func (f FisExperimentAction) Status(ctx context.Context, state *FisExperimentState) (*action_kit_api.StatusResult, error) {
 	return statusExperiment(ctx, state, func(account string, region string) (FisStatusExperimentClient, error) {
-		awsAccount, err := utils.Accounts.GetAccount(account, region)
+		awsAccess, err := utils.GetAwsAccess(account, region)
 		if err != nil {
 			return nil, err
 		}
-		return fis.NewFromConfig(awsAccount.AwsConfig), nil
+		return fis.NewFromConfig(awsAccess.AwsConfig), nil
 	})
 }
 
@@ -218,11 +218,11 @@ type FisStopExperimentClient interface {
 
 func (f FisExperimentAction) Stop(ctx context.Context, state *FisExperimentState) (*action_kit_api.StopResult, error) {
 	return stopExperiment(ctx, state, func(account string, region string) (FisStopExperimentClient, error) {
-		awsAccount, err := utils.Accounts.GetAccount(account, region)
+		awsAccess, err := utils.GetAwsAccess(account, region)
 		if err != nil {
 			return nil, err
 		}
-		return fis.NewFromConfig(awsAccount.AwsConfig), nil
+		return fis.NewFromConfig(awsAccess.AwsConfig), nil
 	})
 }
 

diff --git a/extfis/template_discovery.go b/extfis/template_discovery.go
@@ -97,7 +97,7 @@ func (f *fisTemplateDiscovery) DescribeAttributes() []discovery_kit_api.Attribut
 	}
 }
 func (f *fisTemplateDiscovery) DiscoverTargets(ctx context.Context) ([]discovery_kit_api.Target, error) {
-	return utils.ForEveryAccount(utils.Accounts, getTargetsForAccount, ctx, "fis-template")
+	return utils.ForEveryConfiguredAwsAccess(getTargetsForAccount, ctx, "fis-template")
 }
 func getTargetsForAccount(account *utils.AwsAccess, ctx context.Context) ([]discovery_kit_api.Target, error) {
 	client := fis.NewFromConfig(account.AwsConfig)

diff --git a/extlambda/attack.go b/extlambda/attack.go
@@ -138,10 +138,10 @@ func (a *lambdaAction) Stop(ctx context.Context, state *LambdaActionState) (*act
 }
 
 func defaultClientProvider(account string, region string) (ssmApi, error) {
-	awsAccount, err := utils.Accounts.GetAccount(account, region)
+	awsAccess, err := utils.GetAwsAccess(account, region)
 	if err != nil {
 		return nil, err
 	}
-	client := ssm.NewFromConfig(awsAccount.AwsConfig)
+	client := ssm.NewFromConfig(awsAccess.AwsConfig)
 	return client, nil
 }
diff --git a/extlambda/discovery.go b/extlambda/discovery.go
@@ -173,7 +173,7 @@ func (*lambdaDiscovery) DescribeAttributes() []discovery_kit_api.AttributeDescri
 }
 
 func (l *lambdaDiscovery) DiscoverTargets(ctx context.Context) ([]discovery_kit_api.Target, error) {
-	return utils.ForEveryAccount(utils.Accounts, getTargetsForAccount, ctx, "lambda")
+	return utils.ForEveryConfiguredAwsAccess(getTargetsForAccount, ctx, "lambda")
 }
 
 func getTargetsForAccount(account *utils.AwsAccess, ctx context.Context) ([]discovery_kit_api.Target, error) {

diff --git a/extmsk/cluster_discovery.go b/extmsk/cluster_discovery.go
@@ -151,7 +151,7 @@ func (r *mskClusterDiscovery) DescribeAttributes() []discovery_kit_api.Attribute
 }
 
 func (r *mskClusterDiscovery) DiscoverTargets(ctx context.Context) ([]discovery_kit_api.Target, error) {
-	return utils.ForEveryAccount(utils.Accounts, getClusterTargetsForAccount, ctx, "msk-cluster")
+	return utils.ForEveryConfiguredAwsAccess(getClusterTargetsForAccount, ctx, "msk-cluster")
 }
 
 func getClusterTargetsForAccount(account *utils.AwsAccess, ctx context.Context) ([]discovery_kit_api.Target, error) {

diff --git a/extmsk/common.go b/extmsk/common.go
@@ -30,9 +30,9 @@ type MskApi interface {
 }
 
 func defaultMskClientProvider(account string, region string) (MskApi, error) {
-	awsAccount, err := utils.Accounts.GetAccount(account, region)
+	awsAccess, err := utils.GetAwsAccess(account, region)
 	if err != nil {
 		return nil, err
 	}
-	return kafka.NewFromConfig(awsAccount.AwsConfig), nil
+	return kafka.NewFromConfig(awsAccess.AwsConfig), nil
 }
diff --git a/extrds/cluster_discovery.go b/extrds/cluster_discovery.go
@@ -121,7 +121,7 @@ func (r *rdsClusterDiscovery) DescribeAttributes() []discovery_kit_api.Attribute
 }
 
 func (r *rdsClusterDiscovery) DiscoverTargets(ctx context.Context) ([]discovery_kit_api.Target, error) {
-	return utils.ForEveryAccount(utils.Accounts, getClusterTargetsForAccount, ctx, "rds-cluster")
+	return utils.ForEveryConfiguredAwsAccess(getClusterTargetsForAccount, ctx, "rds-cluster")
 }
 
 func getClusterTargetsForAccount(account *utils.AwsAccess, ctx context.Context) ([]discovery_kit_api.Target, error) {

diff --git a/extrds/common_cluster.go b/extrds/common_cluster.go
@@ -33,9 +33,9 @@ func convertClusterAttackState(request action_kit_api.PrepareActionRequestBody,
 	return nil
 }
 func defaultClusterClientProvider(account string, region string) (rdsDBClusterApi, error) {
-	awsAccount, err := utils.Accounts.GetAccount(account, region)
+	awsAccess, err := utils.GetAwsAccess(account, region)
 	if err != nil {
 		return nil, err
 	}
-	return rds.NewFromConfig(awsAccount.AwsConfig), nil
+	return rds.NewFromConfig(awsAccess.AwsConfig), nil
 }