Private/Resolve-RBACContext.ps1

function Get-RBACContextAccountId {
    <#
    .SYNOPSIS
        StrictMode-safe extraction of an Az context's account id (for logging).
    #>

    [CmdletBinding()]
    param([Parameter()][object]$Context)
    if ($Context -and ($Context.PSObject.Properties.Name -contains 'Account') -and $Context.Account) {
        return $Context.Account.Id
    }
    return '<unknown>'
}

function Get-RBACContextTenantId {
    <#
    .SYNOPSIS
        StrictMode-safe extraction of an Az context's tenant id.
    #>

    [CmdletBinding()]
    param([Parameter()][object]$Context)
    if ($Context -and ($Context.PSObject.Properties.Name -contains 'Tenant') -and $Context.Tenant) {
        return $Context.Tenant.Id
    }
    return $null
}

function Resolve-RBACContext {
    <#
    .SYNOPSIS
        Builds the probe execution context (the identity probes run *as*).
    .DESCRIPTION
        Internal. PSAutoRBAC separates two identities:
 
          * the *probe identity* - who runs the access checks / live probe, and
          * the *target caller* (-CallerId on the public cmdlets) - whose access
            is being evaluated and, optionally, granted.
 
        By default the probe identity is the ambient session (Get-AzContext /
        Get-MgContext). Supply -RunAsCredential, -RunAsServicePrincipal, or
        -RunAsManagedIdentity to run the probe as a different identity without
        disturbing the caller's interactive session: a child Az context is
        created (scoped to this module) and torn down by Disconnect().
 
        The returned context object exposes:
          IsRunAs [bool]
          TenantId [string]
          AzContext [object] the Az context to use (ambient or run-as)
          GetAccessToken [scriptblock] param([string]$ResourceUrl) -> token string
          Disconnect [scriptblock] tears down a run-as context (no-op for ambient)
    .OUTPUTS
        PSCustomObject (PSAutoRBAC.Context)
    #>

    [CmdletBinding(DefaultParameterSetName = 'Ambient')]
    [OutputType([psobject])]
    param(
        [Parameter()]
        [string]$TenantId,

        [Parameter(ParameterSetName = 'Credential')]
        [pscredential]$RunAsCredential,

        [Parameter(ParameterSetName = 'ServicePrincipal')]
        [pscredential]$RunAsServicePrincipal,

        [Parameter(ParameterSetName = 'ServicePrincipal')]
        [string]$RunAsTenantId,

        [Parameter(ParameterSetName = 'ManagedIdentity')]
        [switch]$RunAsManagedIdentity,

        [Parameter(ParameterSetName = 'ManagedIdentity')]
        [string]$RunAsManagedIdentityClientId
    )

    $isRunAs    = $PSCmdlet.ParameterSetName -ne 'Ambient'
    $azContext  = $null
    $disconnect = { }

    Write-PSFMessage -Level Verbose -Message "Resolving probe context (mode: $($PSCmdlet.ParameterSetName))." -Tag 'PSAutoRBAC', 'Context'

    $hasAz = [bool](Get-Command -Name 'Connect-AzAccount' -ErrorAction SilentlyContinue)

    if ($isRunAs) {
        if (-not $hasAz) {
            Write-PSFMessage -Level Error -Message 'Run-as probing requires Az.Accounts, which is not available.' -Tag 'PSAutoRBAC', 'Context'
            throw 'Run-as probing requires the Az.Accounts module (Connect-AzAccount). Install it or use the ambient session.'
        }
        Write-PSFMessage -Level Significant -Message "Establishing a run-as probe context via $($PSCmdlet.ParameterSetName)." -Tag 'PSAutoRBAC', 'Context'

        $connect = @{ Scope = 'Process'; ErrorAction = 'Stop' }
        switch ($PSCmdlet.ParameterSetName) {
            'Credential' {
                $connect['Credential'] = $RunAsCredential
            }
            'ServicePrincipal' {
                $connect['ServicePrincipal'] = $true
                $connect['Credential']       = $RunAsServicePrincipal
                $tid = if ($RunAsTenantId) { $RunAsTenantId } else { $TenantId }
                if (-not $tid) { throw 'Service-principal run-as requires -RunAsTenantId (or -TenantId).' }
                $connect['Tenant'] = $tid
            }
            'ManagedIdentity' {
                $connect['Identity'] = $true
                if ($RunAsManagedIdentityClientId) {
                    $connect['AccountId'] = $RunAsManagedIdentityClientId
                }
            }
        }
        if ($TenantId -and -not $connect.ContainsKey('Tenant')) { $connect['Tenant'] = $TenantId }

        $login     = Connect-AzAccount @connect
        $azContext = $login.Context
        Write-PSFMessage -Level Verbose -Message "Run-as context established (account '$(Get-RBACContextAccountId $azContext)', tenant '$(Get-RBACContextTenantId $azContext)')." -Tag 'PSAutoRBAC', 'Context'
        $disconnect = {
            try {
                Write-PSFMessage -Level Verbose -Message "Tearing down run-as context '$($azContext.Name)'." -Tag 'PSAutoRBAC', 'Context'
                Disconnect-AzAccount -ContextName $azContext.Name -Scope Process -ErrorAction SilentlyContinue | Out-Null
            }
            catch { Write-PSFMessage -Level Warning -Message "Run-as context teardown skipped: $($_.Exception.Message)" -Tag 'PSAutoRBAC', 'Context' }
        }.GetNewClosure()
    }
    elseif ($hasAz) {
        $azContext = Get-AzContext -ErrorAction SilentlyContinue
        Write-PSFMessage -Level Debug -Message "Using ambient Az context: $(if ($azContext) { Get-RBACContextAccountId $azContext } else { '<none>' })." -Tag 'PSAutoRBAC', 'Context'
    }
    else {
        Write-PSFMessage -Level Debug -Message 'Az.Accounts not present; context is metadata-only (offline).' -Tag 'PSAutoRBAC', 'Context'
    }

    $effectiveTenant = $TenantId
    if (-not $effectiveTenant -and $azContext) { $effectiveTenant = Get-RBACContextTenantId $azContext }

    $getToken = {
        param([Parameter(Mandatory)][string]$ResourceUrl)
        if (-not (Get-Command -Name 'Get-AzAccessToken' -ErrorAction SilentlyContinue)) {
            Write-PSFMessage -Level Error -Message "Get-AzAccessToken unavailable; cannot acquire a token for '$ResourceUrl'." -Tag 'PSAutoRBAC', 'Context'
            throw "Get-AzAccessToken is unavailable. Import Az.Accounts (and Connect-AzAccount) to acquire a token for '$ResourceUrl'."
        }
        Write-PSFMessage -Level Debug -Message "Acquiring access token for resource '$ResourceUrl'." -Tag 'PSAutoRBAC', 'Context'
        $params = @{ ResourceUrl = $ResourceUrl; ErrorAction = 'Stop' }
        if ($azContext) { $params['DefaultProfile'] = $azContext }
        $tok = Get-AzAccessToken @params
        # Az 12+ may return the token as a SecureString; normalize to plain text.
        if ($tok.Token -is [System.Security.SecureString]) {
            return [System.Net.NetworkCredential]::new('', $tok.Token).Password
        }
        return $tok.Token
    }.GetNewClosure()

    [pscustomobject]@{
        PSTypeName     = 'PSAutoRBAC.Context'
        IsRunAs        = $isRunAs
        TenantId       = $effectiveTenant
        AzContext      = $azContext
        GetAccessToken = $getToken
        Disconnect     = $disconnect
    }
}