Add support for reading deployment IP addresses

pkg/polaris/aws/aws.go

@@ -106,21 +106,21 @@ func (a API) toCloudAccountID(ctx context.Context, id IdentityFunc) (uuid.UUID,
 	}
 	identity, err := id(ctx)
 	if err != nil {
-		return uuid.Nil, fmt.Errorf("failed to lookup identity: %v", err)
+		return uuid.Nil, fmt.Errorf("failed to lookup identity: %s", err)
 	}
 
 	if identity.internal {
 		id, err := uuid.Parse(identity.id)
 		if err != nil {
-			return uuid.Nil, fmt.Errorf("failed to parse identity: %v", err)
+			return uuid.Nil, fmt.Errorf("failed to parse identity: %s", err)
 		}
 
 		return id, nil
 	}
 
 	accountsWithFeatures, err := aws.Wrap(a.client).CloudAccountsWithFeatures(ctx, core.FeatureCloudNativeProtection, identity.id)
 	if err != nil {
-		return uuid.Nil, fmt.Errorf("failed to get account: %v", err)
+		return uuid.Nil, fmt.Errorf("failed to get account: %s", err)
 	}
 
 	// Find the exact match.
@@ -143,7 +143,7 @@ func (a API) toNativeID(ctx context.Context, id IdentityFunc) (string, error) {
 	}
 	identity, err := id(ctx)
 	if err != nil {
-		return "", fmt.Errorf("failed to lookup identity: %v", err)
+		return "", fmt.Errorf("failed to lookup identity: %s", err)
 	}
 
 	if !identity.internal {
@@ -152,12 +152,12 @@ func (a API) toNativeID(ctx context.Context, id IdentityFunc) (string, error) {
 
 	uid, err := uuid.Parse(identity.id)
 	if err != nil {
-		return "", fmt.Errorf("failed to parse identity: %v", err)
+		return "", fmt.Errorf("failed to parse identity: %s", err)
 	}
 
-	accountWithFeatures, err := aws.Wrap(a.client).CloudAccountWithFeatures(ctx, uid, core.FeatureCloudNativeProtection)
+	accountWithFeatures, err := aws.Wrap(a.client).CloudAccountWithFeatures(ctx, uid, core.FeatureAll)
 	if err != nil {
-		return "", fmt.Errorf("failed to get account: %v", err)
+		return "", fmt.Errorf("failed to get account: %s", err)
 	}
 
 	return accountWithFeatures.Account.NativeID, nil
@@ -195,36 +195,55 @@ func (a API) Account(ctx context.Context, id IdentityFunc, feature core.Feature)
 	}
 	identity, err := id(ctx)
 	if err != nil {
-		return CloudAccount{}, fmt.Errorf("failed to lookup identity: %v", err)
+		return CloudAccount{}, fmt.Errorf("failed to lookup identity: %s", err)
 	}
 
 	if identity.internal {
 		cloudAccountID, err := uuid.Parse(identity.id)
 		if err != nil {
-			return CloudAccount{}, fmt.Errorf("failed to parse identity: %v", err)
+			return CloudAccount{}, fmt.Errorf("failed to parse identity: %s", err)
 		}
 
+		// We need to list all accounts and filter on the cloud account id since
+		// the API that looks up cloud accounts returns archived accounts too.
+		// Note, as of now, this endpoint doesn't return the permission groups
+		// for the accounts.
 		accountsWithFeatures, err := aws.Wrap(a.client).CloudAccountsWithFeatures(ctx, feature, "")
 		if err != nil {
-			return CloudAccount{}, fmt.Errorf("failed to get account: %v", err)
+			return CloudAccount{}, fmt.Errorf("failed to get account: %s", err)
 		}
 
 		// Find the exact match.
 		for _, accountWithFeatures := range accountsWithFeatures {
 			if accountWithFeatures.Account.ID == cloudAccountID {
-				return toCloudAccount(accountWithFeatures), nil
+				// Explicitly get the permission groups for the account.
+				account, err := aws.Wrap(a.client).CloudAccountWithFeatures(ctx, accountWithFeatures.Account.ID, feature)
+				if err != nil {
+					return CloudAccount{}, fmt.Errorf("failed to get account with permission groups: %s", err)
+				}
+
+				return toCloudAccount(account), nil
 			}
 		}
 	} else {
+		// We need to list accounts and filter on the native id since there is
+		// no API to look up an account by native id. Note, as of now, this
+		// endpoint doesn't return the permission groups for the accounts.
 		accountsWithFeatures, err := aws.Wrap(a.client).CloudAccountsWithFeatures(ctx, feature, identity.id)
 		if err != nil {
-			return CloudAccount{}, fmt.Errorf("failed to get account: %v", err)
+			return CloudAccount{}, fmt.Errorf("failed to get account: %s", err)
 		}
 
 		// Find the exact match.
 		for _, accountWithFeatures := range accountsWithFeatures {
 			if accountWithFeatures.Account.NativeID == identity.id {
-				return toCloudAccount(accountWithFeatures), nil
+				// Explicitly get the permission groups for the account.
+				account, err := aws.Wrap(a.client).CloudAccountWithFeatures(ctx, accountWithFeatures.Account.ID, feature)
+				if err != nil {
+					return CloudAccount{}, fmt.Errorf("failed to get account with permission groups: %s", err)
+				}
+
+				return toCloudAccount(account), nil
 			}
 		}
 	}
@@ -237,14 +256,22 @@ func (a API) Account(ctx context.Context, id IdentityFunc, feature core.Feature)
 func (a API) Accounts(ctx context.Context, feature core.Feature, filter string) ([]CloudAccount, error) {
 	a.log.Print(log.Trace)
 
+	// Note, as of now, this endpoint doesn't return the permission groups for
+	// the accounts.
 	accountsWithFeatures, err := aws.Wrap(a.client).CloudAccountsWithFeatures(ctx, feature, filter)
 	if err != nil {
-		return nil, fmt.Errorf("failed to get account: %v", err)
+		return nil, fmt.Errorf("failed to get accounts: %s", err)
 	}
 
 	accounts := make([]CloudAccount, 0, len(accountsWithFeatures))
 	for _, accountWithFeatures := range accountsWithFeatures {
-		accounts = append(accounts, toCloudAccount(accountWithFeatures))
+		// Explicitly get the permission groups for each account.
+		account, err := aws.Wrap(a.client).CloudAccountWithFeatures(ctx, accountWithFeatures.Account.ID, feature)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get accounts with permission groups: %s", err)
+		}
+
+		accounts = append(accounts, toCloudAccount(account))
 	}
 
 	return accounts, nil

pkg/polaris/aws/aws_test.go

@@ -234,6 +234,9 @@ func TestAwsAccountAddAndRemoveWithPermissionGroups(t *testing.T) {
 	if account.Features[0].Status != core.StatusConnected {
 		t.Fatalf("invalid feature status: %v", account.Features[0].Status)
 	}
+	if groups := account.Features[0].PermissionGroups; !reflect.DeepEqual(groups, []core.PermissionGroup{core.PermissionGroupBasic}) {
+		t.Fatalf("invalid permission groups: %v", groups)
+	}
 	if !account.Features[1].Equal(core.FeatureExocompute) {
 		t.Fatalf("invalid feature name: %v", account.Features[1].Name)
 	}
@@ -243,7 +246,9 @@ func TestAwsAccountAddAndRemoveWithPermissionGroups(t *testing.T) {
 	if account.Features[1].Status != core.StatusConnected {
 		t.Fatalf("invalid feature status: %v", account.Features[0].Status)
 	}
-
+	if groups := account.Features[1].PermissionGroups; !reflect.DeepEqual(groups, []core.PermissionGroup{core.PermissionGroupBasic, core.PermissionGroupRSCManagedCluster}) {
+		t.Fatalf("invalid permission groups: %v", groups)
+	}
 	// Remove AWS account from RSC.
 	err = awsClient.RemoveAccount(ctx, Profile(testAccount.Profile), features, false)
 	if err != nil {

pkg/polaris/graphql/core/core.go

@@ -368,8 +368,30 @@ func (a API) DeploymentVersion(ctx context.Context) (string, error) {
 	return payload.Data.DeploymentVersion, nil
 }
 
-// AllEnabledFeaturesForAccount returns all features enable for the RSC account.
-func (a API) AllEnabledFeaturesForAccount(ctx context.Context) ([]Feature, error) {
+// DeploymentIPAddresses returns the deployment IP addresses.
+func (a API) DeploymentIPAddresses(ctx context.Context) ([]string, error) {
+	a.log.Print(log.Trace)
+
+	buf, err := a.GQL.Request(ctx, allDeploymentIpAddressesQuery, struct{}{})
+	if err != nil {
+		return nil, fmt.Errorf("failed to request allDeploymentIpAddresses: %w", err)
+	}
+	a.log.Printf(log.Debug, "allDeploymentIpAddresses(): %s", string(buf))
+
+	var payload struct {
+		Data struct {
+			DeploymentIPAddresses []string `json:"allDeploymentIpAddresses"`
+		} `json:"data"`
+	}
+	if err := json.Unmarshal(buf, &payload); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal allDeploymentIpAddresses: %v", err)
+	}
+
+	return payload.Data.DeploymentIPAddresses, nil
+}
+
+// EnabledFeaturesForAccount returns all features enable for the RSC account.
+func (a API) EnabledFeaturesForAccount(ctx context.Context) ([]Feature, error) {
 	a.log.Print(log.Trace)
 
 	buf, err := a.GQL.Request(ctx, allEnabledFeaturesForAccountQuery, struct{}{})

pkg/polaris/graphql/core/queries/all_deployment_ip_addresses.graphql

@@ -0,0 +1,3 @@
+query RubrikPolarisSDKRequest {
+  allDeploymentIpAddresses
+}