Public/NC.Connections.ps1

#Requires -Version 5.0
using namespace System.Management.Automation

# Nebula.Core: Connections helpers ==================================================================================================================

function Connect-EOL {
    <#
    .SYNOPSIS
        Establishes a connection to Exchange Online using EXO V3 module.
    .DESCRIPTION
        Imports the ExchangeOnlineManagement module (if needed) and invokes Connect-ExchangeOnline
        with opinionated defaults. When no explicit UserPrincipalName is supplied, the helper
        Find-UserConnected is used to auto-detect the interactive user.
    .PARAMETER UserPrincipalName
        UPN/e-mail used for the authentication prompt. Defaults to the current user.
    .PARAMETER DelegatedOrganization
        Optional customer tenant to target when running in delegated admin scenarios.
    .PARAMETER DisableWAM
        Disable Web Account Manager (WAM) for the Exchange Online sign-in flow.
    .PARAMETER Device
        Use device-code authentication instead of the default interactive flow.
    .PARAMETER NoWamFallback
        Do not automatically retry with -DisableWAM when the default WAM-based flow fails.
    .PARAMETER PassThru
        Return the Connect-ExchangeOnline result (session info) to the caller.
    .EXAMPLE
        Connect-EOL -UserPrincipalName 'admin@contoso.com'
    .EXAMPLE
        'admin@contoso.com' | Connect-EOL -DelegatedOrganization 'customer.contoso.com' -PassThru
    #>

    [CmdletBinding()]
    param(
        [Parameter(ValueFromPipeline = $true, ValueFromPipelineByPropertyName = $true)]
        [Alias('UPN', 'User')]
        [string]$UserPrincipalName,
        [string]$DelegatedOrganization,
        [switch]$DisableWAM,
        [switch]$Device,
        [switch]$NoWamFallback,
        [switch]$PassThru
    )

    begin {
        $module = Get-Module -Name ExchangeOnlineManagement -ListAvailable | Select-Object -First 1
        if (-not $module) {
            throw "ExchangeOnlineManagement module not found. Install it with Install-Module ExchangeOnlineManagement."
        }

        if (-not (Get-Module -Name ExchangeOnlineManagement)) {
            Import-Module ExchangeOnlineManagement -ErrorAction Stop
        }

        $baseConnectParams = @{
            ShowBanner            = $false
            SkipLoadingCmdletHelp = $true
        }

        # Hide Connect-ExchangeOnline's cosmetic "WAM is enabled by default" notice unless the
        # caller wants verbose diagnostics (e.g. while investigating WAM/broker issues).
        if ($VerbosePreference -eq 'SilentlyContinue') {
            $baseConnectParams.WarningAction = 'SilentlyContinue'
        }
    }

    process {
        if (-not $PSBoundParameters.ContainsKey('UserPrincipalName') -or [string]::IsNullOrWhiteSpace($UserPrincipalName)) {
            $UserPrincipalName = Find-UserConnected
        }

        if ([string]::IsNullOrWhiteSpace($UserPrincipalName)) {
            throw "Unable to determine the UserPrincipalName. Provide -UserPrincipalName explicitly."
        }

        $connectParams = $baseConnectParams.Clone()
        $connectParams.UserPrincipalName = $UserPrincipalName

        if ($PSBoundParameters.ContainsKey('DelegatedOrganization') -and $DelegatedOrganization) {
            $connectParams.DelegatedOrganization = $DelegatedOrganization
        }

        # -DisableWAM/-Device only exist on ExchangeOnlineManagement 3.7.2+ (WAM-by-default releases).
        # Older, pinned installs don't expose them; passing them anyway would throw a parameter-binding error.
        $exoCmd = Get-Command -Name Connect-ExchangeOnline -Module ExchangeOnlineManagement -ErrorAction SilentlyContinue
        $supportsDisableWAM = $exoCmd -and $exoCmd.Parameters.ContainsKey('DisableWAM')
        $supportsDevice = $exoCmd -and $exoCmd.Parameters.ContainsKey('Device')

        if ($DisableWAM.IsPresent) {
            if ($supportsDisableWAM) {
                $connectParams.DisableWAM = $true
            }
            else {
                Write-Verbose "-DisableWAM requested but the installed ExchangeOnlineManagement version does not support it (pre-3.7.2, no WAM by default). Ignoring."
            }
        }

        if ($Device.IsPresent) {
            if ($supportsDevice) {
                $connectParams.Device = $true
            }
            else {
                Write-NCMessage "-Device requested but the installed ExchangeOnlineManagement version does not support it (pre-3.7.2). Falling back to the standard interactive sign-in." -Level WARNING
            }
        }

        $authMode = if ($Device.IsPresent -and $supportsDevice) {
            'device code'
        }
        elseif ($DisableWAM.IsPresent -and $supportsDisableWAM) {
            'interactive without WAM'
        }
        else {
            'interactive'
        }

        Write-NCMessage "Connecting to Exchange Online as $UserPrincipalName using $authMode ..." -Level INFO

        try {
            $session = Connect-ExchangeOnline @connectParams
        }
        catch {
            $shouldRetryWithoutWam = $supportsDisableWAM -and (-not $DisableWAM.IsPresent) -and (-not $Device.IsPresent) -and (-not $NoWamFallback.IsPresent)
            $exceptionText = $_ | Out-String

            if ($shouldRetryWithoutWam -and $exceptionText -match '(?i)(RuntimeBroker|Web Account Manager|\bWAM\b|Error Acquiring Token|NullReferenceException)') {
                Write-NCMessage "Exchange Online sign-in via WAM failed. Retrying with WAM disabled ..." -Level WARNING
                $connectParams.DisableWAM = $true
                $session = Connect-ExchangeOnline @connectParams
            }
            else {
                throw
            }
        }

        if ($PassThru.IsPresent) {
            return $session
        }
    }
}

function Connect-Nebula {
    <#
    .SYNOPSIS
        Entry point to establish both Exchange Online and Microsoft Graph sessions.
    .DESCRIPTION
        Uses the private Test-* helpers to ensure Exchange Online and Microsoft Graph are connected,
        connecting Microsoft Graph first so its authentication dependencies are loaded before
        Exchange Online and using a WAM-disabled EXO sign-in for the combined flow. It can also
        force reconnects, install modules, or skip Graph.
    .PARAMETER UserPrincipalName
        Optional explicit UPN for the Exchange Online connection.
    .PARAMETER GraphScopes
        Microsoft Graph delegated scopes to request (default User.Read.All).
    .PARAMETER GraphTenantId
        Optional tenant ID/domain for the Graph connection.
    .PARAMETER GraphDeviceCode
        Use device-code auth for Graph instead of launching a browser.
    .PARAMETER AutoInstall
        Install missing modules automatically without prompting.
    .PARAMETER ForceReconnect
        Force reconnect (skip health checks) for both services.
    .PARAMETER SkipGraph
        Only establish Exchange Online; skip Microsoft Graph.
    .PARAMETER GraphLoginHint
        UPN passed as -LoginHint to Connect-MgGraph, helping the WAM broker resolve the target
        account without repeatedly prompting. Defaults to UserPrincipalName when not specified.
    #>

    [CmdletBinding()]
    param(
        [string]$UserPrincipalName,
        [string[]]$GraphScopes = @('User.Read.All'),
        [string]$GraphTenantId,
        [switch]$GraphDeviceCode,
        [switch]$AutoInstall,
        [switch]$ForceReconnect,
        [switch]$SkipGraph,
        [string]$GraphLoginHint
    )

    Write-NCMessage "Welcome to Nebula! Connecting, please wait ..." -Level INFO

    try {
        $checkUpdates = $NCVars.CheckUpdatesOnConnect
        if ($checkUpdates -is [string]) {
            $checkUpdates = $checkUpdates -match '^(1|true|yes|y)$'
        }
        else {
            $checkUpdates = [bool]$checkUpdates
        }

        if ($checkUpdates) {
            Test-NebulaModuleUpdates | Out-Null
        }
    }
    catch {
        Write-NCMessage "Update check failed. $($_.Exception.Message)" -Level WARNING
    }

    if (-not $SkipGraph) {
        $activeGraphContext = $null
        try { $activeGraphContext = Get-MgContext -ErrorAction Stop } catch {}

        $graphLoginHint = if (-not [string]::IsNullOrWhiteSpace($GraphLoginHint)) {
            $GraphLoginHint
        }
        elseif (-not [string]::IsNullOrWhiteSpace($UserPrincipalName)) {
            $UserPrincipalName
        }
        elseif ($activeGraphContext -and -not [string]::IsNullOrWhiteSpace($activeGraphContext.Account)) {
            # Keep an active Graph session's account instead of switching to the workstation identity
            $activeGraphContext.Account
        }
        else {
            Find-UserConnected
        }

        # Keep the active session's tenant (e.g. a guest or delegated tenant) when the same account stays connected
        $graphTenant = $GraphTenantId
        if ([string]::IsNullOrWhiteSpace($graphTenant) -and $activeGraphContext -and $activeGraphContext.Account -eq $graphLoginHint) {
            $graphTenant = $activeGraphContext.TenantId
        }

        $graphConnected = Test-MgGraphConnection `
            -Scopes $GraphScopes `
            -TenantId $graphTenant `
            -UseDeviceCode:$GraphDeviceCode.IsPresent `
            -AutoInstall:$AutoInstall.IsPresent `
            -ForceReconnect:$ForceReconnect.IsPresent `
            -EnsureExchangeOnline:$false `
            -LoginHint $graphLoginHint

        if (-not $graphConnected) {
            throw "Failed to establish Microsoft Graph session."
        }
    }

    $exoConnected = Test-EOLConnection -UserPrincipalName $UserPrincipalName `
        -AutoInstall:$AutoInstall.IsPresent `
        -ForceReconnect:$ForceReconnect.IsPresent `
        -DisableWAM:$(-not $SkipGraph)

    if (-not $exoConnected) {
        throw "Failed to establish Exchange Online session."
    }

    if ($SkipGraph) {
        return [pscustomobject]@{
            ExchangeOnline = $true
            MicrosoftGraph = $false
        }
    }

    return [pscustomobject]@{
        ExchangeOnline = $true
        MicrosoftGraph = $true
    }
}

function Disconnect-Nebula {
    <#
    .SYNOPSIS
        Disconnects from Exchange Online and Microsoft Graph sessions.
    .DESCRIPTION
        Calls Disconnect-ExchangeOnline (if available) and Disconnect-MgGraph,
        suppressing common errors and allowing targeted disconnects.
    .PARAMETER ExchangeOnly
        Disconnect only Exchange Online.
    .PARAMETER GraphOnly
        Disconnect only Microsoft Graph.
    #>

    [CmdletBinding()]
    param(
        [switch]$ExchangeOnly,
        [switch]$GraphOnly
    )

    $disconnectExo = if ($GraphOnly) { $false } else { $true }
    $disconnectGraph = if ($ExchangeOnly) { $false } else { $true }

    if ($disconnectExo) {
        try {
            if (Get-Command -Name Disconnect-ExchangeOnline -ErrorAction SilentlyContinue) {
                Disconnect-ExchangeOnline -Confirm:$false -ErrorAction Stop
                Write-NCMessage "Exchange Online session disconnected." -Level INFO
            }
        }
        catch {
            Write-NCMessage "Failed to disconnect Exchange Online: $($_.Exception.Message)" -Level WARNING
        }
    }

    if ($disconnectGraph) {
        try {
            if (Get-Command -Name Disconnect-MgGraph -ErrorAction SilentlyContinue) {
                Disconnect-MgGraph -ErrorAction Stop | Out-Null
                Write-NCMessage "Microsoft Graph session disconnected." -Level INFO
            }
        }
        catch {
            Write-NCMessage "Failed to disconnect Microsoft Graph: $($_.Exception.Message)" -Level WARNING
        }
    }
}

function Get-NebulaConnections {
    <#
    .SYNOPSIS
        Shows active Nebula connection status for Exchange Online and Microsoft Graph.
    .DESCRIPTION
        Checks current Exchange Online and Microsoft Graph session contexts and performs lightweight
        health probes for both services (unless skipped). Those probes can also trigger provider-side
        token/session refresh when supported. Returns a compact status object for interactive and
        scripted usage.
    #>

    [CmdletBinding()]
    param(
        [switch]$SkipHealthCheck
    )

    $exoConnected = $false
    $exoHealthy = $false
    $exoHealthError = $null
    $exoUser = $null
    $exoOrganization = $null

    try {
        if (Get-Command -Name Get-ConnectionInformation -ErrorAction SilentlyContinue) {
            $exoSession = Get-ConnectionInformation -ErrorAction Stop | Select-Object -First 1
            if ($exoSession) {
                $exoConnected = $true
                $exoUser = $exoSession.UserPrincipalName
                $exoOrganization = $exoSession.Organization
            }
        }
    }
    catch {
        $exoConnected = $false
    }

    if ($SkipHealthCheck.IsPresent) {
        $exoHealthy = $exoConnected
    }
    elseif ($exoConnected) {
        try {
            if (Get-Command -Name Get-EXOMailbox -ErrorAction SilentlyContinue) {
                # Lightweight probe: if this fails, session is usually stale/expired.
                Get-EXOMailbox -ResultSize 1 -ErrorAction Stop | Out-Null
                $exoHealthy = $true
            }
            else {
                $exoHealthError = "Get-EXOMailbox command not available for validation."
            }
        }
        catch {
            $exoHealthy = $false
            $exoHealthError = $_.Exception.Message
        }
    }

    $graphConnected = $false
    $graphHealthy = $false
    $graphHealthError = $null
    $graphAccount = $null
    $graphTenantId = $null
    $graphScopes = @()

    try {
        if (Get-Command -Name Get-MgContext -ErrorAction SilentlyContinue) {
            $graphContext = Get-MgContext -ErrorAction Stop
            if ($graphContext -and $graphContext.Account) {
                $graphConnected = $true
                $graphAccount = $graphContext.Account
                $graphTenantId = $graphContext.TenantId
                $graphScopes = @($graphContext.Scopes)
            }
        }
    }
    catch {
        $graphConnected = $false
    }

    if ($SkipHealthCheck.IsPresent) {
        $graphHealthy = $graphConnected
    }
    elseif ($graphConnected) {
        try {
            if (Get-Command -Name Invoke-MgGraphRequest -ErrorAction SilentlyContinue) {
                # Lightweight probe against /me to validate current access token.
                Invoke-MgGraphRequest -Method GET -Uri 'v1.0/me?$select=id' -ErrorAction Stop | Out-Null
                $graphHealthy = $true
            }
            else {
                $graphHealthError = "Invoke-MgGraphRequest command not available for validation."
            }
        }
        catch {
            $graphHealthy = $false
            $graphHealthError = $_.Exception.Message
        }
    }

    [pscustomobject]@{
        ExchangeOnlineConnected = $exoConnected
        ExchangeOnlineHealthy   = $exoHealthy
        ExchangeOnlineError     = $exoHealthError
        ExchangeOnlineUser      = $exoUser
        ExchangeOnlineTenant    = $exoOrganization
        MicrosoftGraphConnected = $graphConnected
        MicrosoftGraphHealthy   = $graphHealthy
        MicrosoftGraphError     = $graphHealthError
        MicrosoftGraphAccount   = $graphAccount
        MicrosoftGraphTenantId  = $graphTenantId
        MicrosoftGraphScopes    = $graphScopes
    }
}

function Update-NebulaConnections {
    <#
    .SYNOPSIS
        Refreshes Nebula connections for Exchange Online and Microsoft Graph, repairing them if unhealthy.
    .DESCRIPTION
        Runs the same health probes as Get-NebulaConnections and, unless skipped, reconnects any
        service found disconnected or unhealthy (stale/broken session) using the same detected user
        and, for Microsoft Graph, the scopes already granted to the current session. Returns the
        resulting connection status.
    .PARAMETER SkipHealthCheck
        Skip probe/repair entirely and only report whether session contexts are currently present.
    .EXAMPLE
        Update-NebulaConnections
    .EXAMPLE
        Update-NebulaConnections -SkipHealthCheck
    #>

    [CmdletBinding()]
    param(
        [switch]$SkipHealthCheck
    )

    if ($SkipHealthCheck.IsPresent) {
        return Get-NebulaConnections -SkipHealthCheck
    }

    # Probe first: Test-MgGraphConnection trusts a cached context, so a stale Graph token needs -ForceReconnect
    $status = $null
    try {
        $status = Get-NebulaConnections
    }
    catch {}
    $forceGraphReconnect = [bool]($status -and $status.MicrosoftGraphConnected -and -not $status.MicrosoftGraphHealthy)

    $graphScopes = @()
    $graphAccount = $null
    $graphTenant = $null
    try {
        $graphContext = Get-MgContext -ErrorAction Stop
        $graphScopes = @($graphContext.Scopes | Where-Object { $_ })
        $graphAccount = $graphContext.Account
        # A reconnect without the tenant would sign the account in to its home tenant
        $graphTenant = $graphContext.TenantId
    }
    catch {}
    if (-not $graphScopes -or $graphScopes.Count -eq 0) {
        $graphScopes = @('User.Read.All')
    }
    # Repair the current Graph account instead of switching to the workstation identity
    if ([string]::IsNullOrWhiteSpace($graphAccount)) {
        $graphAccount = Find-UserConnected
    }

    # Graph before Exchange Online, as in Connect-Nebula, to avoid the cross-module authentication conflict
    $graphConnected = $false
    try {
        $graphConnected = [bool](Test-MgGraphConnection -Scopes $graphScopes -TenantId $graphTenant -EnsureExchangeOnline:$false -LoginHint $graphAccount -ForceReconnect:$forceGraphReconnect)
    }
    catch {
        Write-NCMessage "Microsoft Graph repair attempt failed. $($_.Exception.Message)" -Level WARNING
    }

    $exoParams = @{}
    try {
        $exoUser = (Get-ConnectionInformation -ErrorAction Stop | Select-Object -First 1).UserPrincipalName
        if (-not [string]::IsNullOrWhiteSpace($exoUser)) {
            $exoParams.UserPrincipalName = $exoUser
        }
    }
    catch {}

    try {
        # With a Graph session in the same process, EXO signs in without WAM (see Connect-Nebula)
        Test-EOLConnection @exoParams -DisableWAM:$graphConnected | Out-Null
    }
    catch {
        Write-NCMessage "Exchange Online repair attempt failed. $($_.Exception.Message)" -Level WARNING
    }

    Get-NebulaConnections
}

Set-Alias -Name Leave-Nebula -Value Disconnect-Nebula