add condition for determining different graph endpoints per environment

cisagov · Jan 18, 2025 · 2b5330f · 2b5330f
1 parent 2b5b6c6
commit 2b5330f
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 5 deletions.
diff --git a/PowerShell/ScubaGear/Modules/Providers/ExportAADProvider.psm1 b/PowerShell/ScubaGear/Modules/Providers/ExportAADProvider.psm1
@@ -184,7 +184,7 @@ function Export-AADProvider {
     Import-Module $PSScriptRoot/ProviderHelpers/AADRiskyPermissionsHelper.psm1
 
     $RiskyApps = $Tracker.TryCommand("Get-ApplicationsWithRiskyPermissions")
-    $RiskySPs = $Tracker.TryCommand("Get-ServicePrincipalsWithRiskyPermissions")
+    $RiskySPs = $Tracker.TryCommand("Get-ServicePrincipalsWithRiskyPermissions", @{"M365Environment"=$M365Environment})
 
     $RiskyApps = if ($null -eq $RiskyApps -or $RiskyApps.Count -eq 0) { $null } else { $RiskyApps }
     $RiskySPs = if ($null -eq $RiskySPs -or $RiskySPs.Count -eq 0) { $null } else { $RiskySPs }
@@ -197,7 +197,6 @@ function Export-AADProvider {
         $AggregateRiskyApps = "{}"
         $ThirdPartySPs = "{}"
     }
-
     ##### End block
 
     $SuccessfulCommands = ConvertTo-Json @($Tracker.GetSuccessfulCommands())

diff --git a/PowerShell/ScubaGear/Modules/Providers/ProviderHelpers/AADRiskyPermissionsHelper.psm1 b/PowerShell/ScubaGear/Modules/Providers/ProviderHelpers/AADRiskyPermissionsHelper.psm1
@@ -1,5 +1,3 @@
-using module "..\..\ScubaConfig\ScubaConfig.psm1"
-
 function Get-RiskyPermissionsJson {
     process {
         try {
@@ -230,6 +228,11 @@ function Get-ServicePrincipalsWithRiskyPermissions {
     .Functionality
     #Internal
     ##>
+    param (
+        [ValidateNotNullOrEmpty()]
+        [string]
+        $M365Environment
+    )
     process {
         try {
             $RiskyPermissionsJson = Get-RiskyPermissionsJson
@@ -247,6 +250,17 @@ function Get-ServicePrincipalsWithRiskyPermissions {
                 $Chunks.Add($ServicePrincipalIds[$i..([math]::Min($i + $ChunkSize - 1, $ServicePrincipalIds.Count - 1))])
             }
 
+            $endpoint = '/beta/$batch'
+            if ($M365Environment -eq "gcchigh") {
+                $endpoint = "https://graph.microsoft.us" + $endpoint
+            }
+            elseif ($M365Environment -eq "dod") {
+                $endpoint = "https://dod-graph.microsoft.us" + $endpoint
+            }
+            else {
+                $endpoint = "https://graph.microsoft.com" + $endpoint
+            }
+
             # Process each chunk
             foreach ($Chunk in $Chunks) {
                 $BatchBody = @{
@@ -262,7 +276,7 @@ function Get-ServicePrincipalsWithRiskyPermissions {
                 }
 
                 # Send the batch request
-                $Response = Invoke-MgGraphRequest -Method POST -Uri 'https://graph.microsoft.com/beta/$batch' -Body (
+                $Response = Invoke-MgGraphRequest -Method POST -Uri $endpoint -Body (
                     $BatchBody | ConvertTo-Json -Depth 5
                 )