modules/Azure/Infrastructure/Public/Invoke-AzureApi.ps1

function Invoke-AzureApi {
    <#
    .SYNOPSIS
        Invokes Azure REST API (ARM, Graph, or KeyVault) with standardized error handling.

    .DESCRIPTION
        Single point of entry for all Azure API calls. Handles authentication
        internally - callers should not deal with tokens or auth logic.

        Supports three calling patterns:
        - ByUri: Pass a full URL via -Uri
        - ByPath: Pass a relative path via -Path with a mandatory -Api selector
        - Batch: Pass Graph sub-requests via -Requests to call Microsoft Graph $batch
    #>

    [CmdletBinding(DefaultParameterSetName = 'ByUri')]
    [OutputType([PSObject])]
    param(
        [Parameter(Mandatory, ParameterSetName = 'ByUri')]
        [ValidateNotNullOrEmpty()]
        [string]$Uri,

        [Parameter(Mandatory, ParameterSetName = 'ByPath')]
        [ValidateNotNullOrEmpty()]
        [string]$Path,

        [Parameter(Mandatory, ParameterSetName = 'Batch')]
        [AllowEmptyCollection()]
        [hashtable[]]$Requests,

        [Parameter(ParameterSetName = 'ByUri')]
        [Parameter(Mandatory, ParameterSetName = 'ByPath')]
        [Parameter(Mandatory, ParameterSetName = 'Batch')]
        [ValidateSet('ARM', 'Graph', 'GraphBeta', 'KeyVault')]
        [string]$Api,

        [Parameter(ParameterSetName = 'ByPath')]
        [string[]]$SubscriptionId,

        [Parameter(Mandatory)]
        [ValidateNotNullOrEmpty()]
        [string]$ResourceName,

        [Parameter(ParameterSetName = 'ByUri')]
        [Parameter(ParameterSetName = 'ByPath')]
        [ValidateSet('GET', 'POST', 'PUT', 'PATCH', 'DELETE')]
        [string]$Method = 'GET',

        [Parameter(ParameterSetName = 'ByUri')]
        [Parameter(ParameterSetName = 'ByPath')]
        [object]$Body,

        [Parameter(ParameterSetName = 'ByUri')]
        [Parameter(ParameterSetName = 'ByPath')]
        [switch]$Raw
    )

    $ProgressPreference = 'SilentlyContinue'
    $shouldThrow = $ErrorActionPreference -eq 'Stop'
    $ErrorActionPreference = 'Stop'

    function ConvertTo-HeaderMap {
        param(
            [Parameter()]
            [object]$HeaderSource
        )

        $headerMap = @{}
        if (-not $HeaderSource) {
            return $headerMap
        }

        if ($HeaderSource -is [System.Net.Http.Headers.HttpHeaders]) {
            foreach ($entry in $HeaderSource) {
                $headerMap[$entry.Key.ToLowerInvariant()] = (@($entry.Value) | ForEach-Object { [string]$_ }) -join ','
            }
            return $headerMap
        }

        if ($HeaderSource -is [System.Collections.IDictionary]) {
            foreach ($key in @($HeaderSource.Keys)) {
                if ($null -eq $key) {
                    continue
                }

                $value = $HeaderSource[$key]
                if ($null -eq $value) {
                    continue
                }

                $headerMap[[string]$key.ToLowerInvariant()] = if ($value -is [System.Array]) {
                    @($value) -join ','
                }
                else {
                    [string]$value
                }
            }
            return $headerMap
        }

        foreach ($property in $HeaderSource.PSObject.Properties) {
            if ($null -eq $property.Value) {
                continue
            }

            $headerMap[$property.Name.ToLowerInvariant()] = if ($property.Value -is [System.Array]) {
                @($property.Value) -join ','
            }
            else {
                [string]$property.Value
            }
        }

        $headerMap
    }

    function Get-HeaderValue {
        param(
            [Parameter()]
            [hashtable]$Headers,
            [Parameter(Mandatory)]
            [string]$Name
        )

        if (-not $Headers) {
            return $null
        }

        $lookupName = $Name.ToLowerInvariant()
        if ($Headers.ContainsKey($lookupName)) {
            return $Headers[$lookupName]
        }

        $null
    }

    function Invoke-SafeRestMethod {
        param(
            [Parameter(Mandatory)]
            [string]$RequestUri,
            [Parameter(Mandatory)]
            [hashtable]$RequestHeaders,
            [Parameter(Mandatory)]
            [string]$RequestMethod,
            [Parameter()]
            [string]$JsonBody
        )

        try {
            $responseHeaders = $null
            $restParams = @{
                Uri                     = $RequestUri
                Method                  = $RequestMethod
                Headers                 = $RequestHeaders
                ErrorAction             = 'Stop'
                ResponseHeadersVariable = 'responseHeaders'
            }

            if ($JsonBody) {
                $restParams.Body = $JsonBody
                $restParams.ContentType = 'application/json'
            }

            $restResponse = Invoke-RestMethod @restParams

            [pscustomobject]@{
                StatusCode = 200
                Content = $restResponse | ConvertTo-Json -Depth 20 -Compress
                Headers = ConvertTo-HeaderMap -HeaderSource $responseHeaders
            }
        }
        catch [Microsoft.PowerShell.Commands.HttpResponseException] {
            $errorBody = $null
            try {
                $errorBody = $_.ErrorDetails.Message
            }
            catch {}

            [pscustomobject]@{
                StatusCode = [int]$_.Exception.Response.StatusCode
                Content = $errorBody
                Headers = ConvertTo-HeaderMap -HeaderSource $_.Exception.Response.Headers
            }
        }
        catch {
            [pscustomobject]@{
                StatusCode = 0
                Content = $_.Exception.Message
                Headers = @{}
            }
        }
    }

    function Get-RetryDelaySeconds {
        param(
            [Parameter(Mandatory)]
            [object]$Response,
            [Parameter(Mandatory)]
            [int]$RetryAttempt
        )

        $retryAfterHeader = Get-HeaderValue -Headers $Response.Headers -Name 'Retry-After'
        if ($retryAfterHeader) {
            $parsedRetryAfter = $null
            if ([System.Net.Http.Headers.RetryConditionHeaderValue]::TryParse($retryAfterHeader, [ref]$parsedRetryAfter)) {
                if ($parsedRetryAfter.Delta -and $parsedRetryAfter.Delta.TotalSeconds -gt 0) {
                    return [math]::Ceiling($parsedRetryAfter.Delta.TotalSeconds)
                }

                if ($parsedRetryAfter.Date) {
                    $headerDelay = [math]::Ceiling(($parsedRetryAfter.Date - [DateTimeOffset]::UtcNow).TotalSeconds)
                    if ($headerDelay -gt 0) {
                        return $headerDelay
                    }
                }
            }
        }

        try {
            if ($Response.Content) {
                $errorContent = $Response.Content | ConvertFrom-Json -ErrorAction Stop
                if ($errorContent.retryAfter) {
                    $bodyDelay = [double]$errorContent.retryAfter
                    if ($bodyDelay -gt 0) {
                        return $bodyDelay
                    }
                }
            }
        }
        catch {}

        $baseDelay = [math]::Pow(2, $RetryAttempt - 1)
        $jitter = Get-Random -Minimum 0.8 -Maximum 1.2
        [math]::Round($baseDelay * $jitter, 2)
    }

    function Invoke-AzureRequestWithRetry {
        param(
            [Parameter(Mandatory)]
            [string]$RequestUri,
            [Parameter(Mandatory)]
            [hashtable]$RequestHeaders,
            [Parameter(Mandatory)]
            [string]$RequestMethod,
            [Parameter()]
            [string]$JsonBody,
            [Parameter(Mandatory)]
            [string]$FriendlyName
        )

        $maxRetries = 5
        $retryAttempts = 0

        while ($true) {
            $response = Invoke-SafeRestMethod -RequestUri $RequestUri -RequestHeaders $RequestHeaders -RequestMethod $RequestMethod -JsonBody $JsonBody
            if ($response.StatusCode -ne 429) {
                return $response
            }

            if ($retryAttempts -ge $maxRetries) {
                throw "Failed to load $FriendlyName - retries exhausted after $maxRetries retries."
            }

            $retryAttempts++
            $retryDelay = Get-RetryDelaySeconds -Response $response -RetryAttempt $retryAttempts
            Write-Verbose "[$FriendlyName] Rate limited (429). Retry $retryAttempts of $maxRetries after ${retryDelay}s..."
            InvokeCIEMAzureSleep -Seconds $retryDelay
        }
    }

    function Parse-ResponseContent {
        param(
            [Parameter()]
            [string]$Content
        )

        if (-not $Content) {
            return $null
        }

        try {
            $Content | ConvertFrom-Json -ErrorAction Stop
        }
        catch {
            $Content
        }
    }

    function Get-ParsedErrorMessage {
        param(
            [Parameter()]
            [object]$ParsedContent
        )

        if (-not $ParsedContent) {
            return $null
        }

        if ($ParsedContent -is [string]) {
            return $ParsedContent
        }

        if ($ParsedContent.PSObject.Properties.Name -contains 'error' -and $ParsedContent.error) {
            if ($ParsedContent.error.PSObject.Properties.Name -contains 'message') {
                return $ParsedContent.error.message
            }
        }

        $ParsedContent | ConvertTo-Json -Depth 20 -Compress
    }

    function Invoke-AzureBatchRequests {
        param(
            [Parameter(Mandatory)]
            [AllowEmptyCollection()]
            [hashtable[]]$BatchRequests,
            [Parameter(Mandatory)]
            [string]$BatchApi,
            [Parameter(Mandatory)]
            [hashtable]$BatchHeaders,
            [Parameter(Mandatory)]
            [string]$BatchResourceName
        )

        if (@($BatchRequests).Count -eq 0) {
            throw 'Invoke-AzureApi batch mode requires at least one batch request.'
        }

        if ($BatchApi -notin @('Graph', 'GraphBeta')) {
            throw "Invoke-AzureApi batch mode supports Graph APIs only. Received '$BatchApi'."
        }

        $batchEndpoint = (Get-CIEMAzureProviderApi -Name $BatchApi).BaseUrl.TrimEnd('/')
        $batchUri = "$batchEndpoint/`$batch"
        $results = @{}

        for ($offset = 0; $offset -lt $BatchRequests.Count; $offset += $script:CIEMGraphBatchSize) {
            $remaining = $BatchRequests.Count - $offset
            $chunkSize = [Math]::Min($script:CIEMGraphBatchSize, $remaining)
            $pendingRequests = @($BatchRequests[$offset..($offset + $chunkSize - 1)])
            $retryAttempts = @{}

            while ($pendingRequests.Count -gt 0) {
                $payloadRequests = foreach ($request in $pendingRequests) {
                    if (-not $request.Id) {
                        throw 'Invoke-AzureApi batch requests must include Id.'
                    }
                    if (-not $request.Method) {
                        throw "Invoke-AzureApi batch request '$($request.Id)' is missing Method."
                    }
                    if (-not $request.Path) {
                        throw "Invoke-AzureApi batch request '$($request.Id)' is missing Path."
                    }

                    @{
                        id = [string]$request.Id
                        method = [string]$request.Method
                        url = $request.Path.TrimStart('/')
                    }
                }

                $payload = @{
                    requests = $payloadRequests
                }

                $payloadJson = $payload | ConvertTo-Json -Depth 20 -Compress
                $batchResponse = Invoke-AzureRequestWithRetry -RequestUri $batchUri -RequestHeaders $BatchHeaders -RequestMethod 'POST' -JsonBody $payloadJson -FriendlyName "$BatchResourceName batch"

                if ($batchResponse.StatusCode -ne 200) {
                    $parsedError = Parse-ResponseContent -Content $batchResponse.Content
                    $errorDetail = Get-ParsedErrorMessage -ParsedContent $parsedError
                    throw "Failed to load $BatchResourceName batch - Status: $($batchResponse.StatusCode) - $errorDetail"
                }

                $parsedBatchResponse = Parse-ResponseContent -Content $batchResponse.Content
                if (-not $parsedBatchResponse -or -not ($parsedBatchResponse.PSObject.Properties.Name -contains 'responses')) {
                    throw "Failed to load $BatchResourceName batch - malformed batch response."
                }

                $subResponsesById = @{}
                foreach ($subResponse in @($parsedBatchResponse.responses)) {
                    $subResponsesById[[string]$subResponse.id] = $subResponse
                }

                $nextPending = [System.Collections.Generic.List[hashtable]]::new()
                $retryDelays = [System.Collections.Generic.List[double]]::new()

                foreach ($request in $pendingRequests) {
                    $requestId = [string]$request.Id
                    if (-not $subResponsesById.ContainsKey($requestId)) {
                        throw "Failed to load $BatchResourceName batch - missing response for request '$requestId'."
                    }

                    $subResponse = $subResponsesById[$requestId]
                    $statusCode = [int]$subResponse.status

                    if ($statusCode -eq 429) {
                        $currentRetryCount = if ($retryAttempts.ContainsKey($requestId)) { $retryAttempts[$requestId] } else { 0 }
                        if ($currentRetryCount -ge 5) {
                            throw "Failed to load $BatchResourceName - batch sub-request '$requestId' retries exhausted."
                        }

                        $retryAttempts[$requestId] = $currentRetryCount + 1
                        $retryResponse = [pscustomobject]@{
                            StatusCode = 429
                            Content = if ($subResponse.body) { $subResponse.body | ConvertTo-Json -Depth 20 -Compress } else { $null }
                            Headers = ConvertTo-HeaderMap -HeaderSource $subResponse.headers
                        }
                        $retryDelays.Add((Get-RetryDelaySeconds -Response $retryResponse -RetryAttempt $retryAttempts[$requestId]))
                        $nextPending.Add($request)
                        continue
                    }

                    $subBody = $subResponse.body
                    if ($statusCode -lt 200 -or $statusCode -ge 300) {
                        $results[$requestId] = [pscustomobject]@{
                            Id = $requestId
                            Success = $false
                            StatusCode = $statusCode
                            Items = @()
                            Content = $subBody
                            Error = Get-ParsedErrorMessage -ParsedContent $subBody
                        }
                        continue
                    }

                    $items = @()
                    if ($subBody) {
                        if ($subBody.PSObject.Properties.Name -contains 'value') {
                            $items = @($subBody.value)
                        }
                        elseif ($subBody -is [System.Array]) {
                            $items = @($subBody)
                        }
                        else {
                            $items = @($subBody)
                        }

                        if ($subBody.PSObject.Properties.Name -contains '@odata.nextLink' -and $subBody.'@odata.nextLink') {
                            $items += @(Invoke-AzureApi -Uri $subBody.'@odata.nextLink' -Api $BatchApi -ResourceName "$BatchResourceName/$requestId" -ErrorAction Stop)
                        }
                    }

                    $results[$requestId] = [pscustomobject]@{
                        Id = $requestId
                        Success = $true
                        StatusCode = $statusCode
                        Items = @($items)
                        Content = $subBody
                        Error = $null
                    }
                }

                if ($nextPending.Count -gt 0) {
                    $sleepSeconds = ($retryDelays | Measure-Object -Maximum).Maximum
                    InvokeCIEMAzureSleep -Seconds $sleepSeconds
                    $pendingRequests = @($nextPending.ToArray())
                    continue
                }

                break
            }
        }

        $results
    }

    if ($PSCmdlet.ParameterSetName -eq 'ByPath') {
        $apiRecord = Get-CIEMAzureProviderApi -Name $Api
        if (-not $apiRecord) {
            $msg = "No API endpoint record found for '$Api' in azure_provider_apis table."
            if ($shouldThrow) { throw $msg }
            Write-Warning $msg
            return
        }

        $baseUrl = $apiRecord.BaseUrl.TrimEnd('/')
        if ($SubscriptionId) {
            $results = @{}
            foreach ($subId in $SubscriptionId) {
                $fullUri = "$baseUrl/subscriptions/$subId/$($Path.TrimStart('/'))"
                $results[$subId] = Invoke-AzureApi -Uri $fullUri -Api $Api -ResourceName "$ResourceName ($subId)" -Method $Method -Body $Body -Raw:$Raw
            }
            return $results
        }

        $Uri = "$baseUrl/$($Path.TrimStart('/'))"
    }

    if (-not $Api) {
        $Api = if ($Uri -match 'graph\.microsoft\.com/beta') {
            'GraphBeta'
        }
        elseif ($Uri -match 'graph\.microsoft\.com') {
            'Graph'
        }
        elseif ($Uri -match '\.vault\.azure\.net') {
            'KeyVault'
        }
        else {
            'ARM'
        }
    }

    if (-not $script:AzureAuthContext -or -not $script:AzureAuthContext.IsConnected) {
        $msg = 'Not connected to Azure. Run Connect-CIEM first.'
        if ($shouldThrow) { throw $msg }
        Write-Warning $msg
        return
    }

    $jsonBody = $null
    if ($Body) {
        $jsonBody = $Body | ConvertTo-Json -Depth 20 -Compress
    }

    $token = switch ($Api) {
        'Graph' { $script:AzureAuthContext.GraphToken }
        'GraphBeta' { $script:AzureAuthContext.GraphToken }
        'ARM' { $script:AzureAuthContext.ARMToken }
        'KeyVault' { $script:AzureAuthContext.KeyVaultToken }
    }

    if (-not $token) {
        $msg = "$Api API call requested but no $Api token available. Run Connect-CIEM first."
        if ($shouldThrow) { throw $msg }
        Write-Warning $msg
        return
    }

    $headers = @{ Authorization = "Bearer $token" }

    if ($PSCmdlet.ParameterSetName -eq 'Batch') {
        return Invoke-AzureBatchRequests -BatchRequests $Requests -BatchApi $Api -BatchHeaders $headers -BatchResourceName $ResourceName
    }

    Write-Verbose "Loading $ResourceName..."

    $currentUri = $Uri
    $currentMethod = $Method
    $currentJsonBody = $jsonBody
    $currentBodyObject = $Body
    $pageCount = 0
    $lastNextLink = $null
    $hasWrittenPaginationProgress = $false

    try {
        while ($currentUri) {
            $currentResponse = Invoke-AzureRequestWithRetry -RequestUri $currentUri -RequestHeaders $headers -RequestMethod $currentMethod -JsonBody $currentJsonBody -FriendlyName $ResourceName

            if ($Raw) {
                return $currentResponse
            }

            switch ($currentResponse.StatusCode) {
                200 {
                    $content = Parse-ResponseContent -Content $currentResponse.Content
                    if (-not $content) {
                        break
                    }

                    $pageCount++
                    $isValueCollection = $content.PSObject.Properties.Name -contains 'value'
                    $isDataCollection = $content.PSObject.Properties.Name -contains 'data'

                    if ($isValueCollection -or $isDataCollection) {
                        $items = if ($isValueCollection) { @($content.value) } else { @($content.data) }
                        $nextLink = if ($content.PSObject.Properties.Name -contains '@odata.nextLink') {
                            $content.'@odata.nextLink'
                        }
                        elseif ($content.PSObject.Properties.Name -contains 'nextLink') {
                            $content.nextLink
                        }
                        else {
                            $null
                        }

                        if ($items.Count -eq 0 -and $nextLink) {
                            $currentUri = $null
                            break
                        }

                        if ($nextLink) {
                            if ($lastNextLink -and $nextLink -eq $lastNextLink) {
                                throw "Failed to load $ResourceName - pagination cycle detected at '$nextLink'."
                            }

                            $lastNextLink = $nextLink
                            $hasWrittenPaginationProgress = $true
                            Write-Progress -Activity "Loading $ResourceName" -Status "Page $pageCount" -CurrentOperation $nextLink
                        }
                        elseif ($hasWrittenPaginationProgress) {
                            Write-Progress -Activity "Loading $ResourceName" -Status "Page $pageCount"
                        }

                        $items

                        if ($nextLink) {
                            $currentUri = $nextLink
                            $currentMethod = 'GET'
                            $currentJsonBody = $null
                            continue
                        }

                        if ($content.PSObject.Properties.Name -contains '$skipToken' -and $content.'$skipToken' -and $currentMethod -eq 'POST') {
                            if ($currentBodyObject) {
                                $nextBody = $currentBodyObject | ConvertTo-Json -Depth 20 -Compress | ConvertFrom-Json -AsHashtable
                            }
                            else {
                                $nextBody = @{}
                            }

                            $nextBody['$skipToken'] = $content.'$skipToken'
                            $currentJsonBody = $nextBody | ConvertTo-Json -Depth 20 -Compress
                            $currentBodyObject = $nextBody
                            $hasWrittenPaginationProgress = $true
                            Write-Progress -Activity "Loading $ResourceName" -Status "Page $pageCount"
                            continue
                        }

                        $currentUri = $null
                        break
                    }

                    return $content
                }
                401 {
                    $msg = "Unauthorized loading $ResourceName - invalid or expired token"
                    if ($shouldThrow) { throw $msg }
                    Write-Warning $msg
                    $currentUri = $null
                    break
                }
                403 {
                    $msg = "Access denied loading $ResourceName - missing permissions"
                    if ($shouldThrow) { throw $msg }
                    Write-Warning $msg
                    $currentUri = $null
                    break
                }
                404 {
                    $msg = "Resource not found: $ResourceName"
                    if ($shouldThrow) { throw $msg }
                    Write-Verbose $msg
                    $currentUri = $null
                    break
                }
                0 {
                    $detail = if ($currentResponse.Content) { $currentResponse.Content } else { 'Unknown error' }
                    $msg = "Failed to load $ResourceName - $detail"
                    if ($shouldThrow) { throw $msg }
                    Write-Warning $msg
                    $currentUri = $null
                    break
                }
                default {
                    $parsedError = Parse-ResponseContent -Content $currentResponse.Content
                    $detail = Get-ParsedErrorMessage -ParsedContent $parsedError
                    $msg = "Failed to load $ResourceName - Status: $($currentResponse.StatusCode)"
                    if ($detail) {
                        $msg += " - $detail"
                    }

                    if ($shouldThrow) { throw $msg }
                    Write-Warning $msg
                    $currentUri = $null
                    break
                }
            }
        }
    }
    finally {
        if ($hasWrittenPaginationProgress) {
            Write-Progress -Activity "Loading $ResourceName" -Completed
        }
    }
}