Private/Auth/Connect-GraphSession.ps1

# Copyright (c) 2026 Sandy Zeng. All rights reserved.
# Source-available. All rights reserved. See LICENSE file.

<#
    Connect-GraphSession.ps1 — Pure-PowerShell interactive sign-in to Microsoft Graph
    using OAuth 2.0 Authorization Code Flow with PKCE on a loopback HttpListener.
    No MSAL, no Microsoft.Graph SDK, no embedded WebView, no WAM broker — and no
    on-disk token cache. All tokens live in memory only.
 
    Author: Sandy Zeng
    Project: IntuneDiff
 
    Credits:
        The browser/PKCE flow is adapted from "Connect-MgGraphViaBrowser.ps1"
        by Ugur Koc (https://github.com/ugurkocde/Intune) — used with credit.
        Original: https://raw.githubusercontent.com/ugurkocde/Intune/refs/heads/main/Connect-MgGraphViaBrowser/Connect-MgGraphViaBrowser.ps1
 
    Version History:
    1.0.0 Initial release.
    1.0.1 Switched from WAM broker to system browser for reliability.
    1.0.2 Token cache is in-memory only; no tokens written to disk.
    2.0.0 Removed all dependence on Microsoft.Graph.Authentication / MSAL.
            Replaced with hand-rolled PKCE auth-code flow (credit: Ugur Koc).
#>


# ---------------------------------------------------------------------------
# Module-scoped state
# ---------------------------------------------------------------------------
$script:GraphClientId   = '14d82eec-204b-4c2f-b7e8-296a70dab67e'  # Microsoft Graph Command Line Tools
$script:GraphResource   = 'https://graph.microsoft.com'
$script:AuthorityHost   = 'https://login.microsoftonline.com'

if (-not $script:GraphAccounts) { $script:GraphAccounts = @{} }  # HomeAccountId -> account record
if (-not $script:GraphCurrent)  { $script:GraphCurrent  = $null } # HomeAccountId of currently active account

# Remove any legacy on-disk caches left by earlier versions
foreach ($legacy in @(
    (Join-Path $env:LOCALAPPDATA 'IntuneDiff\msalcache.bin'),
    (Join-Path ([Environment]::GetFolderPath('LocalApplicationData')) 'Connect-MgGraphViaBrowser\tokens.json')
)) {
    if ($legacy -and (Test-Path -LiteralPath $legacy)) {
        Remove-Item -LiteralPath $legacy -Force -ErrorAction SilentlyContinue
    }
}

# ---------------------------------------------------------------------------
# Public: required / missing scopes
# ---------------------------------------------------------------------------
function Get-IntuneDiffRequiredScopes {
    [CmdletBinding()]
    [OutputType([string[]])]
    param()
    return @(
        'DeviceManagementConfiguration.Read.All'
    )
}

function Get-IntuneDiffMissingScopes {
    <#
    .SYNOPSIS
        Returns any required Graph scopes absent from the current signed-in context.
        ReadWrite variants satisfy a Read requirement.
    #>

    [CmdletBinding()]
    [OutputType([string[]])]
    param([string[]]$Required = (Get-IntuneDiffRequiredScopes))
    if (-not (Test-GraphConnection)) { return $Required }

    $have = @()
    if ($script:SignedInUser -and $script:SignedInUser.Scopes) {
        # Strip the "https://graph.microsoft.com/" prefix if present so comparison is by short name
        $have = @($script:SignedInUser.Scopes | ForEach-Object {
            if ($_ -like "$script:GraphResource/*") { $_.Substring($script:GraphResource.Length + 1) } else { $_ }
        })
    }

    $missing = [System.Collections.Generic.List[string]]::new()
    foreach ($scope in $Required) {
        if ($have -contains $scope) { continue }
        $parts = $scope.Split('.')
        if ($parts.Length -ge 2 -and $parts[1] -eq 'Read') {
            $rwParts = $parts.Clone()
            $rwParts[1] = 'ReadWrite'
            if ($have -contains ($rwParts -join '.')) { continue }
        }
        $missing.Add($scope) | Out-Null
    }
    return @($missing)
}

# ---------------------------------------------------------------------------
# Public: tenant display info (uses our own Invoke-IntuneDiffRequest)
# ---------------------------------------------------------------------------
function Get-TenantDisplayInfo {
    [CmdletBinding()]
    param()
    if (-not (Test-GraphConnection)) { return $null }
    try {
        $resp = Invoke-IntuneDiffRequest -Method GET -Uri 'https://graph.microsoft.com/v1.0/organization?$select=id,displayName,verifiedDomains'
        $org = $resp.value | Select-Object -First 1
        if (-not $org) { return $null }
        $defaultDomain = ($org.verifiedDomains | Where-Object { $_.isDefault } | Select-Object -First 1).name
        [pscustomobject]@{
            TenantId      = $org.id
            DisplayName   = $org.displayName
            DefaultDomain = $defaultDomain
        }
    } catch { $null }
}

# ---------------------------------------------------------------------------
# PKCE / OAuth helpers (adapted from Ugur Koc's Connect-MgGraphViaBrowser.ps1)
# ---------------------------------------------------------------------------
function New-IDCodeVerifier {
    $bytes = [byte[]]::new(32)
    [System.Security.Cryptography.RandomNumberGenerator]::Fill($bytes)
    return [Convert]::ToBase64String($bytes).TrimEnd('=').Replace('+', '-').Replace('/', '_')
}

function New-IDCodeChallenge {
    param([Parameter(Mandatory)][string]$Verifier)
    $hash = [System.Security.Cryptography.SHA256]::HashData([System.Text.Encoding]::UTF8.GetBytes($Verifier))
    return [Convert]::ToBase64String($hash).TrimEnd('=').Replace('+', '-').Replace('/', '_')
}

function Get-IDFreePort {
    $tcp = [System.Net.Sockets.TcpListener]::new([System.Net.IPAddress]::Loopback, 0)
    $tcp.Start()
    try { return $tcp.LocalEndpoint.Port } finally { $tcp.Stop() }
}

function Resolve-IDScopes {
    <# Normalize bare scope names to fully-qualified Graph scopes. Always include
       openid + profile + offline_access so the token response carries an id_token
       (with preferred_username / oid / tid) and a refresh token. #>

    param([string[]]$Scopes)
    $oidc = @('openid', 'profile', 'offline_access', 'email')
    $resolved = @(foreach ($s in $Scopes) {
        if ($s -like 'https://*' -or $s -in $oidc) { $s } else { "$script:GraphResource/$s" }
    })
    foreach ($required in 'openid','profile','offline_access') {
        if ($required -notin $resolved) { $resolved += $required }
    }
    return ,$resolved
}

function Invoke-IDTokenEndpoint {
    param(
        [Parameter(Mandatory)][string]$TenantSegment,
        [Parameter(Mandatory)][hashtable]$Body
    )
    $url = "$script:AuthorityHost/$TenantSegment/oauth2/v2.0/token"
    return Invoke-RestMethod -Method POST -Uri $url -Body $Body `
        -ContentType 'application/x-www-form-urlencoded' -ErrorAction Stop
}

function Get-IDTokenViaRefresh {
    param(
        [Parameter(Mandatory)][string]$RefreshToken,
        [Parameter(Mandatory)][string[]]$Scopes,
        [Parameter(Mandatory)][string]$TenantSegment
    )
    $body = @{
        client_id     = $script:GraphClientId
        grant_type    = 'refresh_token'
        refresh_token = $RefreshToken
        scope         = ($Scopes -join ' ')
    }
    return Invoke-IDTokenEndpoint -TenantSegment $TenantSegment -Body $body
}

function Get-IDTokenViaBrowser {
    param(
        [Parameter(Mandatory)][string]$TenantSegment,
        [Parameter(Mandatory)][string[]]$Scopes,
        [string] $LoginHint,
        [switch] $ForceConsent,
        [switch] $SelectAccount
    )

    $port        = Get-IDFreePort
    $redirectUri = "http://localhost:$port/"

    $verifier  = New-IDCodeVerifier
    $challenge = New-IDCodeChallenge -Verifier $verifier
    $state     = [Guid]::NewGuid().ToString('N')

    $promptValue = if ($ForceConsent) { 'consent' }
                   elseif ($SelectAccount -and -not $LoginHint) { 'select_account' }
                   else { $null }

    $authParams = [ordered]@{
        client_id             = $script:GraphClientId
        response_type         = 'code'
        redirect_uri          = $redirectUri
        response_mode         = 'query'
        scope                 = ($Scopes -join ' ')
        state                 = $state
        code_challenge        = $challenge
        code_challenge_method = 'S256'
    }
    if ($promptValue) { $authParams['prompt'] = $promptValue }
    if ($LoginHint)   { $authParams['login_hint'] = $LoginHint }

    $query = ($authParams.GetEnumerator() | ForEach-Object {
        "$($_.Key)=$([Uri]::EscapeDataString([string]$_.Value))"
    }) -join '&'

    $authUrl = "$script:AuthorityHost/$TenantSegment/oauth2/v2.0/authorize?$query"

    $listener = [System.Net.HttpListener]::new()
    $listener.Prefixes.Add($redirectUri)
    $listener.Start()

    try {
        Write-IDLog "Opening system browser for sign-in (listening on $redirectUri)"
        Start-Process $authUrl | Out-Null

        $ctxTask = $listener.GetContextAsync()

        # Pump the WPF dispatcher while waiting so the host UI stays responsive
        $dispatcher = [System.Windows.Threading.Dispatcher]::CurrentDispatcher
        $deadline   = [DateTime]::UtcNow.AddMinutes(5)
        while (-not $ctxTask.IsCompleted) {
            if ([DateTime]::UtcNow -gt $deadline) {
                throw 'Authentication timed out after 5 minutes.'
            }
            if ($dispatcher) {
                $dispatcher.Invoke([System.Action]{ }, [System.Windows.Threading.DispatcherPriority]::Background)
            }
            [System.Threading.Thread]::Sleep(50)
        }
        if ($ctxTask.IsFaulted) { throw $ctxTask.Exception.InnerException }
        $ctx = $ctxTask.Result

        $code          = $ctx.Request.QueryString['code']
        $err           = $ctx.Request.QueryString['error']
        $errDesc       = $ctx.Request.QueryString['error_description']
        $returnedState = $ctx.Request.QueryString['state']

        $successHtml = @'
<!DOCTYPE html><html><head><meta charset="utf-8">
<title>IntuneDiff &mdash; Signed In</title>
<style>
  *{margin:0;padding:0;box-sizing:border-box;}
  body{background:#111827;color:#F3F4F6;font-family:-apple-system,Segoe UI,sans-serif;
       display:flex;align-items:center;justify-content:center;height:100vh;}
  .card{background:#1F2937;border:1px solid #374151;border-radius:16px;
        padding:48px 56px;text-align:center;max-width:480px;}
  .icon{font-size:52px;margin-bottom:20px;}
  h1{font-size:22px;font-weight:700;margin-bottom:10px;color:#F9FAFB;}
  p{font-size:14px;color:#9CA3AF;line-height:1.6;}
</style></head>
<body><div class="card">
  <div class="icon">&#x2705;</div>
  <h1>Authentication complete</h1>
  <p>You are signed in. You can close this tab and return to IntuneDiff.</p>
</div></body></html>
'@

        $errorHtml = "<html><body style='font-family:Segoe UI,sans-serif;text-align:center;padding-top:80px;color:#c0392b;'><h2>Authentication Failed</h2><p>$([System.Net.WebUtility]::HtmlEncode($errDesc))</p></body></html>"

        $html  = if ($code) { $successHtml } else { $errorHtml }
        $bytes = [System.Text.Encoding]::UTF8.GetBytes($html)
        $ctx.Response.ContentType     = 'text/html'
        $ctx.Response.ContentLength64 = $bytes.Length
        $ctx.Response.OutputStream.Write($bytes, 0, $bytes.Length)
        $ctx.Response.Close()

        if ($returnedState -ne $state) {
            throw 'OAuth state mismatch - possible CSRF. Aborting.'
        }
        if (-not $code) {
            throw "Authorization failed: $err - $errDesc"
        }

        $body = @{
            client_id     = $script:GraphClientId
            grant_type    = 'authorization_code'
            code          = $code
            redirect_uri  = $redirectUri
            code_verifier = $verifier
            scope         = ($Scopes -join ' ')
        }
        return Invoke-IDTokenEndpoint -TenantSegment $TenantSegment -Body $body
    }
    finally {
        $listener.Stop()
        $listener.Close()
    }
}

# ---------------------------------------------------------------------------
# Token record helpers
# ---------------------------------------------------------------------------
function ConvertFrom-IDIdToken {
    <# Decode the unsigned JWT payload of an id_token for upn / tid / oid claims. #>
    param([Parameter(Mandatory)][string]$IdToken)
    $parts = $IdToken.Split('.')
    if ($parts.Length -lt 2) { return $null }
    $payload = $parts[1].Replace('-', '+').Replace('_', '/')
    switch ($payload.Length % 4) {
        2 { $payload += '==' }
        3 { $payload += '='  }
    }
    try {
        $json = [System.Text.Encoding]::UTF8.GetString([Convert]::FromBase64String($payload))
        return ($json | ConvertFrom-Json)
    } catch { return $null }
}

function New-IDAccountRecord {
    <# Build / update an account record from a token-endpoint response. #>
    param(
        [Parameter(Mandatory)][pscustomobject]$Tokens,
        [string[]]$RequestedScopes,
        [string]$TenantSegment
    )

    $claims = $null
    if ($Tokens.id_token) { $claims = ConvertFrom-IDIdToken -IdToken $Tokens.id_token }

    $username = if ($claims -and $claims.preferred_username) { $claims.preferred_username }
                elseif ($claims -and $claims.upn)            { $claims.upn }
                elseif ($claims -and $claims.email)          { $claims.email }
                else                                          { $null }

    $tenantId = if ($claims -and $claims.tid) { $claims.tid }
                elseif ($TenantSegment -and $TenantSegment -ne 'common' -and $TenantSegment -ne 'organizations') { $TenantSegment }
                else { $null }

    $oid = if ($claims -and $claims.oid) { $claims.oid } else { $null }

    # Fallback: id_token absent or missing claims — ask Graph /me directly with the new access token
    if ((-not $username -or -not $tenantId -or -not $oid) -and $Tokens.access_token) {
        try {
            $meResp = Invoke-RestMethod -Method GET -Uri 'https://graph.microsoft.com/v1.0/me' `
                -Headers @{ Authorization = "Bearer $($Tokens.access_token)"; Accept = 'application/json' } `
                -TimeoutSec 15 -ErrorAction Stop
            if (-not $username) { $username = $meResp.userPrincipalName }
            if (-not $oid)      { $oid      = $meResp.id }
            # /me does not return tenantId — leave as-is
        } catch {
            Write-IDLog "Could not resolve identity via /me: $($_.Exception.Message)"
        }
    }

    if (-not $username) { $username = '<unknown>' }

    $homeId = if ($oid -and $tenantId) { "$oid.$tenantId" }
              elseif ($oid)            { $oid }
              elseif ($username -and $username -ne '<unknown>') { $username.ToLowerInvariant() }
              else                     { [Guid]::NewGuid().ToString('N') }

    $expiresOn = if ($Tokens.expires_in) {
        [DateTime]::UtcNow.AddSeconds([int]$Tokens.expires_in - 60)
    } else {
        [DateTime]::UtcNow.AddMinutes(50)
    }

    $scopesGranted = if ($Tokens.scope) { @($Tokens.scope.Split(' ')) } else { @($RequestedScopes) }

    return [pscustomobject]@{
        HomeAccountId = $homeId
        Username      = $username
        TenantId      = $tenantId
        Environment   = 'login.microsoftonline.com'
        AccessToken   = $Tokens.access_token
        RefreshToken  = $Tokens.refresh_token
        ExpiresOn     = $expiresOn
        Scopes        = $scopesGranted
    }
}

function Set-IDCurrentAccount {
    param([Parameter(Mandatory)][pscustomobject]$Record)
    $script:GraphAccounts[$Record.HomeAccountId] = $Record
    $script:GraphCurrent = $Record.HomeAccountId
    $script:SignedInUser = [pscustomobject]@{
        Account       = $Record.Username
        TenantId      = $Record.TenantId
        Scopes        = $Record.Scopes
        HomeAccountId = $Record.HomeAccountId
    }
}

# ---------------------------------------------------------------------------
# Public: cached accounts (mimics MSAL IAccount surface)
# ---------------------------------------------------------------------------
function Get-IDCachedAccounts {
    [CmdletBinding()]
    [OutputType([object[]])]
    param()
    return @($script:GraphAccounts.Values | Sort-Object Username)
}

function Remove-IDCachedAccount {
    [CmdletBinding()]
    param([Parameter(Mandatory)][object]$Account)
    if ($Account -and $Account.HomeAccountId -and $script:GraphAccounts.ContainsKey($Account.HomeAccountId)) {
        $script:GraphAccounts.Remove($Account.HomeAccountId)
        if ($script:GraphCurrent -eq $Account.HomeAccountId) {
            $script:GraphCurrent = $null
        }
    }
}

# ---------------------------------------------------------------------------
# Public: token accessor with silent refresh (used by Invoke-IntuneDiffRequest)
# ---------------------------------------------------------------------------
function Get-IDAccessToken {
    <#
    .SYNOPSIS
        Returns the current access token, refreshing it silently if it has expired.
    #>

    [CmdletBinding()]
    [OutputType([string])]
    param([switch]$ForceRefresh)

    if (-not $script:GraphCurrent -or -not $script:GraphAccounts.ContainsKey($script:GraphCurrent)) {
        throw 'Not signed in. Call Connect-GraphSession first.'
    }
    $record = $script:GraphAccounts[$script:GraphCurrent]

    $needsRefresh = $ForceRefresh -or ($record.ExpiresOn -le [DateTime]::UtcNow)
    if (-not $needsRefresh) { return $record.AccessToken }

    if (-not $record.RefreshToken) {
        throw 'Access token expired and no refresh token is available. Please sign in again.'
    }

    Write-IDLog "Refreshing access token for $($record.Username)..."
    $tenantSegment = if ($record.TenantId) { $record.TenantId } else { 'organizations' }
    $tokens = Get-IDTokenViaRefresh `
        -RefreshToken  $record.RefreshToken `
        -Scopes        $record.Scopes `
        -TenantSegment $tenantSegment

    $updated = New-IDAccountRecord -Tokens $tokens -RequestedScopes $record.Scopes -TenantSegment $tenantSegment
    # Preserve identity if the id_token wasn't reissued
    if ($updated.Username -eq '<unknown>') { $updated.Username = $record.Username }
    if (-not $updated.TenantId)            { $updated.TenantId = $record.TenantId }
    $updated.HomeAccountId = $record.HomeAccountId
    Set-IDCurrentAccount -Record $updated
    return $updated.AccessToken
}

# ---------------------------------------------------------------------------
# Public: Connect-GraphSession
# ---------------------------------------------------------------------------
function Connect-GraphSession {
    <#
    .SYNOPSIS
        Signs in to Microsoft Graph via the system browser using OAuth 2.0
        Authorization Code Flow with PKCE on a loopback HttpListener.
 
    .PARAMETER Account
        An account object from Get-IDCachedAccounts. If supplied and a refresh
        token is available, attempts silent refresh first.
 
    .PARAMETER TenantId
        Optional tenant ID or verified domain. Defaults to 'organizations'.
 
    .PARAMETER Scopes
        Scopes to request. Bare names are auto-prefixed with the Graph resource.
        If omitted, requests https://graph.microsoft.com/.default so already-
        consented permissions come back without re-prompting for consent.
 
    .PARAMETER ForceConsent
        Force the consent prompt to appear (prompt=consent).
    #>

    [CmdletBinding()]
    [Diagnostics.CodeAnalysis.SuppressMessageAttribute('PSAvoidUsingPlainTextForPassword', '',
        Justification = 'PowerShell scope variable, not a credential parameter.')]
    param(
        [object]  $Account,
        [string]  $TenantId,
        [string[]]$Scopes,
        [switch]  $ForceConsent
    )

    if (-not $Scopes -or $Scopes.Count -eq 0) {
        $Scopes = @('.default')
    }
    $resolvedScopes = Resolve-IDScopes -Scopes $Scopes

    $tenantSegment = if ($TenantId) { $TenantId }
                     elseif ($Account -and $Account.TenantId) { $Account.TenantId }
                     else { 'organizations' }

    Write-IDLog "Connect-GraphSession: Account=$($Account.Username), Tenant=$tenantSegment, Scopes=$($resolvedScopes -join ', ')"

    $tokens = $null

    # 1) Silent refresh path
    if ($Account -and $Account.RefreshToken -and -not $ForceConsent) {
        try {
            Write-IDLog "Attempting silent refresh for $($Account.Username)..."
            $tokens = Get-IDTokenViaRefresh `
                -RefreshToken  $Account.RefreshToken `
                -Scopes        $resolvedScopes `
                -TenantSegment $tenantSegment
            Write-IDLog 'Silent refresh succeeded.'
        } catch {
            Write-IDLog "Silent refresh failed, falling back to interactive: $($_.Exception.Message)"
            $tokens = $null
        }
    }

    # 2) Interactive browser path
    if (-not $tokens) {
        $loginHint = if ($Account -and $Account.Username -and $Account.Username -ne '<unknown>') { $Account.Username } else { $null }
        $tokens = Get-IDTokenViaBrowser `
            -TenantSegment $tenantSegment `
            -Scopes        $resolvedScopes `
            -LoginHint     $loginHint `
            -ForceConsent: $ForceConsent `
            -SelectAccount: (-not $loginHint)
    }

    if (-not $tokens -or -not $tokens.access_token) {
        throw 'Sign-in did not complete. No token was acquired.'
    }

    $record = New-IDAccountRecord -Tokens $tokens -RequestedScopes $resolvedScopes -TenantSegment $tenantSegment
    # If we came in via silent-refresh of an existing account, preserve its identity
    if ($Account -and $Account.HomeAccountId -and ($record.Username -eq '<unknown>' -or -not $record.TenantId)) {
        if ($record.Username -eq '<unknown>') { $record.Username = $Account.Username }
        if (-not $record.TenantId)             { $record.TenantId = $Account.TenantId }
        $record.HomeAccountId = $Account.HomeAccountId
    }

    Set-IDCurrentAccount -Record $record
    Write-IDLog "Token acquired for $($record.Username) (Tenant: $($record.TenantId))"
    return $script:SignedInUser
}