Perf: batch Graph requests in guest and privileged-role tests

src/powershell/ZeroTrustAssessment.psd1

@@ -12,7 +12,7 @@
 RootModule = 'ZeroTrustAssessment.psm1'
 
 # Version number of this module.
-ModuleVersion = '2.3.0'
+ModuleVersion = '2.4.0'
 
 # Supported PSEditions
 CompatiblePSEditions = 'Core', 'Desktop'

src/powershell/tests/Test-Assessment.21818.ps1

@@ -100,55 +100,45 @@ ORDER BY rd.displayName;
         }
     )
 
-    $notificationRules = @()
-    # This flag is used to control the flow of a loop, allowing the script to break out of the outer loop when set to $true.
-    $exitLoop = $false
-    $passed = $true
-
-    # Query activation notification rules for administrators
-    # For each policy ID, retrieve notification rules for role activation events sent to administrators
+    # Build one request per (policy, rule) pair and fetch them all in batched calls
     Write-ZtProgress -Activity $activity -Status "Getting activation notification rules"
 
-    foreach ($policyAssignment in $resultsPolicyAssignments) {
-        $policyId = $policyAssignment.policyId
-        $roleDisplayName = $policyAssignment.roleDisplayName
-
+    $ruleRequests = foreach ($policyAssignment in $resultsPolicyAssignments) {
         foreach ($ruleId in $notifications.ruleId) {
-            try {
-                $uri = "policies/roleManagementPolicies/$policyId/rules/$ruleId"
-
-                $rule = Invoke-ZtGraphRequest -RelativeUri $uri -ApiVersion 'v1.0' -ErrorAction Stop
-
-                if ($rule) {
-                    $notificationRules += ($rule | Add-Member -MemberType NoteProperty -Name RoleDisplayName -Value $roleDisplayName -Force -PassThru)
-
-                    # TO-DO: When the performance of the API is improved, we can collect all rules and move the check outside the loop to determine if the test passes or fails.
-                    # Check if isDefaultRecipientsEnabled is false and notificationRecipients is an empty array
-                    if ($rule.isDefaultRecipientsEnabled -eq $false -and
-                        ($null -eq $rule.notificationRecipients -or $rule.notificationRecipients.Count -eq 0)) {
-                        $passed = $false
-                        $exitLoop = $true
-                        break  # Exit inner loop if condition is met
-                    }
-                }
-            }
-            catch {
-                Write-Error "Failed to retrieve rule $ruleId for policy $($policyId): $($_.Exception.Message)"
+            [PSCustomObject]@{
+                PolicyId        = $policyAssignment.policyId
+                RuleId          = $ruleId
+                RoleDisplayName = $policyAssignment.roleDisplayName
             }
         }
-        if ($exitLoop) {
-            break # Exit outer loop if condition is met
+    }
+
+    $notificationRules = [System.Collections.Generic.List[object]]::new()
+    if ($ruleRequests) {
+        $ruleResults = Invoke-ZtGraphBatchRequest -Path "policies/roleManagementPolicies/{0}/rules/{1}" -ArgumentList $ruleRequests -Properties PolicyId, RuleId -Matched -ErrorAction SilentlyContinue
+        foreach ($result in $ruleResults) {
+            if (-not $result.Success -or -not $result.Result) { continue }
+            $rule = $result.Result | Select-Object -First 1
+            $notificationRules.Add(($rule | Add-Member -MemberType NoteProperty -Name RoleDisplayName -Value $result.Argument.RoleDisplayName -Force -PassThru))
         }
     }
 
+    # A role is non-compliant when a notification rule has default recipients disabled and no
+    # additional recipients configured.
+    $failingRule = $notificationRules | Where-Object {
+        $_.isDefaultRecipientsEnabled -eq $false -and
+        ($null -eq $_.notificationRecipients -or $_.notificationRecipients.Count -eq 0)
+    } | Select-Object -First 1
+    $passed = -not $failingRule
+
     $testResultMarkdown = ""
 
     # Output the result of the check
     if ($passed) {
         $testResultMarkdown += "Role notifications are properly configured for privileged role.`n`n%TestResult%"
     }
     else {
-        $testResultMarkdown += "Role notifications are not properly configured.`n`nNote: To save time, this check stops when it finds the first role that does not have notifications. After fixing this role and all other roles, we recommend running the check again to verify.`n`n%TestResult%"
+        $testResultMarkdown += "Role notifications are not properly configured.`n`n%TestResult%"
     }
 
     # Build the detailed sections of the markdown

src/powershell/tests/Test-Assessment.21868.ps1

@@ -41,12 +41,6 @@ function Test-Assessment-21868 {
     $allApp = Invoke-DatabaseQuery -Database $Database -Sql $sqlApp
     $allSP = Invoke-DatabaseQuery -Database $Database -Sql $sqlSP
 
-    $queryParameters = '$select=id,displayName,userPrincipalName'
-
-    # Initialize lists for guest owners only
-    $guestAppOwners = [System.Collections.Generic.List[object]]::new()
-    $guestSpOwners = [System.Collections.Generic.List[object]]::new()
-
     # Get all guest users first (more efficient than repeated queries)
     $sqlGuests = @"
 SELECT id, userPrincipalName, displayName
@@ -61,34 +55,13 @@ WHERE userType = 'Guest'
         [void]$guestUserIds.Add($guest.id)
     }
 
-    # Filter owners to only include guests
-    foreach ($app in $allApp) {
-        $owners = Invoke-ZtGraphRequest -RelativeUri "applications/$($app.id)/owners/microsoft.graph.user?$queryParameters" -ApiVersion 'v1.0'
-        if ($owners) {
-            foreach ($owner in $owners) {
-                $owner | Add-Member -MemberType NoteProperty -Name 'appDisplayName' -Value $app.displayName -Force -PassThru |
-                    Add-Member -MemberType NoteProperty -Name 'appObjectId' -Value $app.id -Force -PassThru |
-                        Add-Member -MemberType NoteProperty -Name 'appId' -Value $app.appId -Force
-                if ($guestUserIds.Contains($owner.id)) {
-                    $guestAppOwners.Add($owner)
-                }
-            }
-        }
-    }
+    Write-ZtProgress -Activity $activity -Status "Getting application owners"
+    $guestAppOwners = @(Get-GuestResourceOwner -Resources $allApp -ResourceType 'applications' -GuestUserIds $guestUserIds `
+            -DisplayNameProperty 'appDisplayName' -ObjectIdProperty 'appObjectId' -AppIdProperty 'appId')
 
-    foreach ($sp in $allSP) {
-        $owners = Invoke-ZtGraphRequest -RelativeUri "servicePrincipals/$($sp.id)/owners/microsoft.graph.user?$queryParameters" -ApiVersion 'v1.0'
-        if ($owners) {
-            foreach ($owner in $owners) {
-                $owner | Add-Member -MemberType NoteProperty -Name 'spDisplayName' -Value $sp.displayName -Force -PassThru |
-                    Add-Member -MemberType NoteProperty -Name 'spObjectId' -Value $sp.id -Force -PassThru |
-                        Add-Member -MemberType NoteProperty -Name 'spAppId' -Value $sp.appId -Force
-                if ($guestUserIds.Contains($owner.id)) {
-                    $guestSpOwners.Add($owner)
-                }
-            }
-        }
-    }
+    Write-ZtProgress -Activity $activity -Status "Getting service principal owners"
+    $guestSpOwners = @(Get-GuestResourceOwner -Resources $allSP -ResourceType 'servicePrincipals' -GuestUserIds $guestUserIds `
+            -DisplayNameProperty 'spDisplayName' -ObjectIdProperty 'spObjectId' -AppIdProperty 'spAppId')
 
     $hasGuestAppOwners = $guestAppOwners.Count -gt 0
     $hasGuestSpOwners = $guestSpOwners.Count -gt 0
@@ -179,3 +152,43 @@ WHERE userType = 'Guest'
 
     Add-ZtTestResultDetail @params
 }
+
+function Get-GuestResourceOwner {
+    [CmdletBinding()]
+    param (
+        [object[]] $Resources,
+
+        [string] $ResourceType,
+
+        [System.Collections.Generic.HashSet[string]] $GuestUserIds,
+
+        [string] $DisplayNameProperty,
+
+        [string] $ObjectIdProperty,
+
+        [string] $AppIdProperty
+    )
+
+    $guestOwners = [System.Collections.Generic.List[object]]::new()
+    if (-not $Resources) { return $guestOwners }
+
+    $resourceById = @{}
+    foreach ($resource in $Resources) { $resourceById[$resource.id] = $resource }
+
+    $ownerPath = "$ResourceType/{0}/owners/microsoft.graph.user?`$select=id,displayName,userPrincipalName"
+    $ownerResults = Invoke-ZtGraphBatchRequest -Path $ownerPath -ArgumentList $Resources.id -Matched -ErrorAction SilentlyContinue
+
+    foreach ($result in $ownerResults) {
+        if (-not $result.Success) { continue }
+        $resource = $resourceById[$result.Argument]
+        foreach ($owner in $result.Result) {
+            if (-not $GuestUserIds.Contains($owner.id)) { continue }
+            $owner | Add-Member -NotePropertyName $DisplayNameProperty -NotePropertyValue $resource.displayName -Force
+            $owner | Add-Member -NotePropertyName $ObjectIdProperty -NotePropertyValue $resource.id -Force
+            $owner | Add-Member -NotePropertyName $AppIdProperty -NotePropertyValue $resource.appId -Force
+            $guestOwners.Add($owner)
+        }
+    }
+
+    return $guestOwners
+}

src/powershell/tests/Test-Assessment.21877.ps1

@@ -58,30 +58,22 @@ WHERE userType = 'Guest'
 
     Write-ZtProgress -Activity $activity -Status "Checking sponsors for $totalGuestCount guest users"
 
-    # Process guests and check sponsors efficiently
-    $guestsWithoutSponsors = [System.Collections.Generic.List[object]]::new()
-    $guestsWithSponsorsCount = 0
-
-    foreach ($guest in $guestUsers) {
-        try {
-            # Get the sponsors for the guest user
-            $guestUserWithSponsors = Invoke-ZtGraphRequest -RelativeUri "users/$($guest.id)?`$expand=sponsors" -ApiVersion 'v1.0'
-
-            # Check if guest has sponsors
-            if ($guestUserWithSponsors.sponsors -and $guestUserWithSponsors.sponsors.Count -gt 0) {
-                $guestsWithSponsorsCount++
-            }
-            else {
-                $guestsWithoutSponsors.Add($guestUserWithSponsors)
-            }
-        }
-        catch {
-            Write-PSFMessage "Failed to get sponsors for guest $($guest.userPrincipalName): $($_.Exception.Message)" -Level Verbose
-            # Treat as guest without sponsor if API call fails
-            $guestsWithoutSponsors.Add($guest)
+    # Collect the IDs of guests with a confirmed sponsor. Anything not confirmed (a failed or empty
+    # lookup) falls through to "without sponsor" by deriving that list from the original guest set.
+    $sponsoredIds = [System.Collections.Generic.HashSet[string]]::new()
+
+    $sponsorResults = Invoke-ZtGraphBatchRequest -Path "users/{0}?`$expand=sponsors" -ArgumentList $guestUsers.id -Matched -ErrorAction SilentlyContinue
+
+    foreach ($result in $sponsorResults) {
+        if (-not $result.Success) { continue }
+        $guestUserWithSponsors = $result.Result | Select-Object -First 1
+        if ($guestUserWithSponsors.sponsors.Count -gt 0) {
+            [void]$sponsoredIds.Add($guestUserWithSponsors.id)
         }
     }
 
+    $guestsWithSponsorsCount = $sponsoredIds.Count
+    $guestsWithoutSponsors = @($guestUsers | Where-Object { -not $sponsoredIds.Contains($_.id) })
     $guestsWithoutSponsorsCount = $guestsWithoutSponsors.Count
     $passed = $guestsWithoutSponsorsCount -eq 0