Src/Private/Connect-EntraIDSession.ps1

function Connect-EntraIDSession {
    <#
    .SYNOPSIS
    Establishes authenticated connections to Microsoft Graph for Entra ID reporting.
    .DESCRIPTION
        Connects to Microsoft Graph API with the scopes required to read Entra ID
        identity and security configuration.
 
        Note: Graph sub-modules are pre-loaded by Invoke-AsBuiltReport.Microsoft.EntraID
        before the PScribo document engine starts, which avoids assembly-already-loaded
        conflicts. This function handles authentication only.
 
        Required Roles (any one of):
          - Global Administrator
          - Global Reader
          - Security Administrator
          - Security Reader
          - Authentication Policy Administrator
 
    .NOTES
        Version: 0.1.20
        Author: Pai Wei Sing
 
    .EXAMPLE
        Connect-EntraIDSession -UserPrincipalName 'admin@contoso.onmicrosoft.com'
    #>

    [CmdletBinding()]
    param (
        [Parameter(Mandatory)]
        [string]$UserPrincipalName
    )

    #region Validate UPN
    if (-not (Test-UserPrincipalName -UserPrincipalName $UserPrincipalName)) {
        $errorMsg = "Invalid User Principal Name format: '$UserPrincipalName'. Expected format: user@domain.com"
        Write-TranscriptLog $errorMsg 'ERROR' 'AUTH'
        throw $errorMsg
    }
    #endregion

    Write-TranscriptLog "Starting connection to Microsoft Graph for Entra ID report: $UserPrincipalName" 'INFO' 'AUTH'

    #region Microsoft Graph -- reuse or connect
    $ExistingGraph = $null
    try { $ExistingGraph = Get-MgContext -ErrorAction SilentlyContinue } catch { }

    if ($ExistingGraph -and $ExistingGraph.TenantId) {
        Write-TranscriptLog "Reusing existing Microsoft Graph session (TenantId: $($ExistingGraph.TenantId))" 'SUCCESS' 'AUTH'
    } else {
        Write-Host " - Connecting to Microsoft Graph..."
        Write-TranscriptLog "Connecting to Microsoft Graph with required Entra ID scopes" 'INFO' 'AUTH'

        $GraphScopes = @(
            'Organization.Read.All'
            'Directory.Read.All'
            'User.Read.All'
            'Group.Read.All'
            'Application.Read.All'
            'Policy.Read.All'
            'RoleManagement.Read.Directory'
            'AuditLog.Read.All'
            'UserAuthenticationMethod.Read.All'
            'Device.Read.All'
            'Reports.Read.All'
            'IdentityRiskyUser.Read.All'
            'IdentityRiskEvent.Read.All'
            'Agreement.Read.All'
            'AccessReview.Read.All'
            'EntitlementManagement.Read.All'
        )

        Invoke-WithRetry -ScriptBlock {
            Connect-MgGraph -Scopes $GraphScopes -ErrorAction Stop
        } -OperationName 'Connect to Microsoft Graph (Entra ID)'

        $MgCtx = Get-MgContext -ErrorAction SilentlyContinue
        if ($MgCtx -and $MgCtx.TenantId) {
            Write-TranscriptLog "Microsoft Graph connection verified (TenantId: $($MgCtx.TenantId))" 'SUCCESS' 'AUTH'
        } else {
            throw "Connect-MgGraph succeeded but Get-MgContext returned no context."
        }
    }
    #endregion

    Write-Host " - Microsoft Graph connected successfully." -ForegroundColor Green
    Write-TranscriptLog "Microsoft Graph connection established for Entra ID report: $UserPrincipalName" 'SUCCESS' 'AUTH'
}


function Disconnect-EntraIDSession {
    <#
    .SYNOPSIS
    Cleanly disconnects Microsoft Graph session used for Entra ID reporting.
    #>

    [CmdletBinding()]
    param()

    Write-TranscriptLog "Disconnecting Microsoft Graph session" 'INFO' 'AUTH'

    try {
        $GraphCtx = Get-MgContext -ErrorAction SilentlyContinue
        if ($GraphCtx) {
            Write-Host " - Disconnecting Microsoft Graph..."
            $null = Disconnect-MgGraph -ErrorAction SilentlyContinue
            Write-TranscriptLog "Microsoft Graph session disconnected" 'SUCCESS' 'AUTH'
        }
    } catch {
        Write-TranscriptLog "Microsoft Graph disconnect warning: $($_.Exception.Message)" 'WARNING' 'AUTH'
    }

    Write-Host " - Session disconnected." -ForegroundColor Green
    Write-TranscriptLog "Entra ID session disconnected" 'SUCCESS' 'AUTH'
}