migrate role assignment cleanup from bash to Go

Dockerfile

@@ -1,5 +1,3 @@
-FROM mcr.microsoft.com/cbl-mariner/base/core:2.0
-RUN tdnf install -y azure-cli jq && tdnf clean all
-COPY rg-cleanup.sh /usr/local/bin
-COPY bin/rg-cleanup ./bin
-ENTRYPOINT [ "rg-cleanup.sh" ]
+FROM alpine:3.21
+COPY bin/rg-cleanup /usr/local/bin
+ENTRYPOINT [ "rg-cleanup" ]

Makefile

@@ -1,6 +1,6 @@
 IMAGE_REGISTRY ?= k8sprowcomm.azurecr.io
 IMAGE_NAME := rg-cleanup
-IMAGE_VERSION ?= v0.4.6
+IMAGE_VERSION ?= v0.4.7
 
 .PHONY: all
 all: build

go.mod

@@ -3,12 +3,12 @@ module github.com/chewong/rg-cleanup
 go 1.13
 
 require (
-	github.com/Azure/azure-sdk-for-go v36.2.0+incompatible
-	github.com/Azure/azure-sdk-for-go/sdk/azcore v1.6.0
-	github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.3.0
+	github.com/Azure/azure-sdk-for-go/sdk/azcore v1.15.0
+	github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.7.0
+	github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/authorization/armauthorization/v2 v2.2.0
 	github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources v1.1.1
 	github.com/Azure/go-autorest/autorest v0.9.2 // indirect
 	github.com/Azure/go-autorest/autorest/adal v0.8.0 // indirect
 	github.com/Azure/go-autorest/autorest/to v0.3.0
-	github.com/Azure/go-autorest/autorest/validation v0.2.0 // indirect
+	github.com/microsoftgraph/msgraph-sdk-go v1.51.0
 )

main.go

@@ -13,7 +13,10 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/arm"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/cloud"
 	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
+	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/authorization/armauthorization/v2"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources"
+	msgraphsdk "github.com/microsoftgraph/msgraph-sdk-go"
+	"github.com/microsoftgraph/msgraph-sdk-go/serviceprincipals"
 )
 
 const (
@@ -40,15 +43,16 @@ var rfc3339Layouts = []string{
 }
 
 type options struct {
-	clientID       string
-	clientSecret   string
-	tenantID       string
-	subscriptionID string
-	dryRun         bool
-	ttl            time.Duration
-	identity       bool
-	regex          string
-	cli            bool
+	clientID        string
+	clientSecret    string
+	tenantID        string
+	subscriptionID  string
+	dryRun          bool
+	ttl             time.Duration
+	identity        bool
+	regex           string
+	cli             bool
+	roleAssignments bool
 }
 
 func (o *options) validate() error {
@@ -84,11 +88,14 @@ func defineOptions() *options {
 	flag.BoolVar(&o.cli, "az-cli", false, "Set to true if we should use az cli for AUTH")
 	flag.DurationVar(&o.ttl, "ttl", defaultTTL, "The duration we allow resource groups to live before we consider them to be stale.")
 	flag.StringVar(&o.regex, "regex", defaultRegex, "Only delete resource groups matching regex")
+	flag.BoolVar(&o.roleAssignments, "role-assignments", false, "Set to true if we should delete role assignments assigned to principals which no longer exist")
 	flag.Parse()
 	return &o
 }
 
 func main() {
+	ctx := context.Background()
+
 	log.Println("Initializing rg-cleanup")
 	log.Printf("args: %v\n", os.Args)
 
@@ -102,19 +109,51 @@ func main() {
 		log.Println("Dry-run enabled - printing logs but not actually deleting resource groups")
 	}
 
-	r, err := getResourceGroupClient(*o)
+	options := arm.ClientOptions{
+		ClientOptions: azcore.ClientOptions{
+			Cloud: cloud.AzurePublic,
+		},
+	}
+
+	cred, err := getAzureCredential(*o)
 	if err != nil {
 		log.Printf("Error when obtaining resource group client: %v", err)
 		panic(err)
 	}
 
-	if err := run(context.Background(), r, o.ttl, o.dryRun, o.regex); err != nil {
-		log.Printf("Error when running rg-cleanup: %v", err)
+	resourceGroupClient, err := armresources.NewResourceGroupsClient(o.subscriptionID, cred, &options)
+	if err != nil {
+		log.Printf("Error when obtaining resource group client: %v", err)
 		panic(err)
 	}
+
+	if err := runResourceGroupCleanup(ctx, resourceGroupClient, o.ttl, o.dryRun, o.regex); err != nil {
+		log.Printf("Error when cleaning up resource groups: %v", err)
+		panic(err)
+	}
+
+	if o.roleAssignments {
+		roleAssignmentClient, err := armauthorization.NewRoleAssignmentsClient(o.subscriptionID, cred, &options)
+		if err != nil {
+			log.Printf("Error when obtaining role assignment client: %v", err)
+			panic(err)
+		}
+
+		graph, err := msgraphsdk.NewGraphServiceClientWithCredentials(cred, nil)
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		if err := runRoleAssignmentCleanup(ctx, o.subscriptionID, roleAssignmentClient, graph, o.dryRun); err != nil {
+			log.Printf("Error when cleaning up role assignments: %v", err)
+			panic(err)
+		}
+	} else {
+		log.Println("Skipping role assignment cleanup")
+	}
 }
 
-func run(ctx context.Context, r *armresources.ResourceGroupsClient, ttl time.Duration, dryRun bool, regex string) error {
+func runResourceGroupCleanup(ctx context.Context, r *armresources.ResourceGroupsClient, ttl time.Duration, dryRun bool, regex string) error {
 	log.Println("Scanning for stale resource groups")
 
 	pager := r.NewListPager(nil)
@@ -199,12 +238,7 @@ func regexMatchesResourceGroupName(regex string, rgName string) (bool, error) {
 	return false, nil
 }
 
-func getResourceGroupClient(o options) (*armresources.ResourceGroupsClient, error) {
-	options := arm.ClientOptions{
-		ClientOptions: azcore.ClientOptions{
-			Cloud: cloud.AzurePublic,
-		},
-	}
+func getAzureCredential(o options) (*azidentity.ChainedTokenCredential, error) {
 	possibleTokens := []azcore.TokenCredential{}
 	if o.identity {
 		micOptions := azidentity.ManagedIdentityCredentialOptions{
@@ -230,13 +264,79 @@ func getResourceGroupClient(o options) (*armresources.ResourceGroupsClient, erro
 	} else {
 		log.Println("unknown login option. login may not succeed")
 	}
-	chain, err := azidentity.NewChainedTokenCredential(possibleTokens, nil)
-	if err != nil {
-		return nil, err
+	return azidentity.NewChainedTokenCredential(possibleTokens, nil)
+}
+
+func runRoleAssignmentCleanup(ctx context.Context, subscriptionID string, roleAssignments *armauthorization.RoleAssignmentsClient, graph *msgraphsdk.GraphServiceClient, dryRun bool) error {
+	log.Println("Scanning for stale role assignments")
+
+	// Role assignments that might be able to be deleted, by principalID to which they're assigned.
+	principalToAssignmentIDs := map[string][]string{}
+	filter := "atScope()" // ignore assignments scoped more narrowly than the subscription
+	pager := roleAssignments.NewListForSubscriptionPager(&armauthorization.RoleAssignmentsClientListForSubscriptionOptions{
+		Filter: &filter,
+	})
+	for pager.More() {
+		assignments, err := pager.NextPage(ctx)
+		if err != nil {
+			return err
+		}
+		for _, assignment := range assignments.Value {
+			if assignment.Properties.PrincipalType == nil || *assignment.Properties.PrincipalType != armauthorization.PrincipalTypeServicePrincipal {
+				continue
+			}
+			// The atScope() filter doesn't ignore assignments scoped more broadly than the subscription
+			if assignment.Properties.Scope == nil || *assignment.Properties.Scope != "/subscriptions/"+subscriptionID {
+				continue
+			}
+			if assignment.Properties.PrincipalID != nil && assignment.ID != nil {
+				pid := *assignment.Properties.PrincipalID
+				principalToAssignmentIDs[pid] = append(principalToAssignmentIDs[pid], *assignment.ID)
+			}
+		}
+	}
+	if len(principalToAssignmentIDs) == 0 {
+		log.Println("No role assignments found")
+		return nil
+	}
+
+	assignedPrincipalIDs := make([]string, 0, len(principalToAssignmentIDs))
+	for k := range principalToAssignmentIDs {
+		assignedPrincipalIDs = append(assignedPrincipalIDs, k)
 	}
-	resourceGroupClient, err := armresources.NewResourceGroupsClient(o.subscriptionID, chain, &options)
+	idReq := serviceprincipals.NewGetByIdsPostRequestBody()
+	idReq.SetIds(assignedPrincipalIDs)
+	idRes, err := graph.ServicePrincipals().GetByIds().PostAsGetByIdsPostResponse(ctx, idReq, &serviceprincipals.GetByIdsRequestBuilderPostRequestConfiguration{})
 	if err != nil {
-		return nil, err
+		return fmt.Errorf("error querying graph: %w", err)
 	}
-	return resourceGroupClient, nil
+
+	// When a role assignment refers to a principal ID that exists, it should not be deleted.
+	for _, id := range idRes.GetValue() {
+		if existingID := id.GetId(); existingID != nil {
+			delete(principalToAssignmentIDs, *existingID)
+		}
+	}
+
+	if len(principalToAssignmentIDs) == 0 {
+		log.Printf("No unattached role assignments found")
+		return nil
+	}
+
+	// The remaining assigned principals no longer exist. Role assignments associated with them should be deleted.
+	for _, assignments := range principalToAssignmentIDs {
+		for _, assignment := range assignments {
+			if dryRun {
+				log.Printf("Dry-run: skip deletion of eligible role assignment %s", assignment)
+				continue
+			}
+			_, err := roleAssignments.DeleteByID(ctx, assignment, nil)
+			if err != nil {
+				return fmt.Errorf("failed to delete role assignment %s: %w", assignment, err)
+			}
+			log.Printf("Deleted role assignment %s", assignment)
+		}
+	}
+
+	return nil
 }

templates/rg-cleaner-logic-app-uami.bicep

@@ -14,7 +14,7 @@ param location string = resourceGroup().location
 param rg_name string = resourceGroup().name
 
 var default_container_cmd = [
-  'rg-cleanup.sh'
+  'rg-cleanup'
   '--identity'
 ]
 var dryrun_cmd = [