Public/Connect-FDAObservability.ps1

function Connect-FDAObservability {
    <#
    .SYNOPSIS
        Authenticate the session against Fabric / Eventhouse / FDA endpoints.
    .DESCRIPTION
        Installs a token-provider closure into module-scope state. All
        subsequent cmdlets resolve tokens through this provider, refreshing
        transparently on expiry.

        Three auth methods supported. The caller picks which one their
        environment uses; the module behaves the same after connect.

    .PARAMETER AuthMethod
        ServicePrincipal | ManagedIdentity | UserDelegated

    .PARAMETER TenantId
        Required for ServicePrincipal and UserDelegated. Optional for
        ManagedIdentity (taken from IMDS metadata).

    .PARAMETER ClientId
        Required for ServicePrincipal. For UserDelegated, defaults to the
        well-known Power BI public client id.

    .PARAMETER ClientSecret
        Required for ServicePrincipal when not using certificate auth.

    .PARAMETER Certificate
        Optional X509 certificate for ServicePrincipal cert-based auth.

    .PARAMETER ManagedIdentityClientId
        Optional. For user-assigned managed identity, the client id to use.

    .PARAMETER WorkspaceId
        Fabric workspace id where the FDAObs database lives.

    .PARAMETER EventhouseId
        Fabric Eventhouse item id. Endpoints are resolved via Fabric REST.

    .PARAMETER DatabaseName
        KQL database name. Defaults to 'FDAObs'.

    .EXAMPLE
        Connect-FDAObservability -AuthMethod ServicePrincipal `
            -TenantId 'a...' -ClientId 'b...' -ClientSecret $sec `
            -WorkspaceId 'w...' -EventhouseId 'e...'

    .EXAMPLE
        Connect-FDAObservability -AuthMethod ManagedIdentity `
            -WorkspaceId 'w...' -EventhouseId 'e...'

    .EXAMPLE
        Connect-FDAObservability -AuthMethod UserDelegated `
            -TenantId 'a...' -WorkspaceId 'w...' -EventhouseId 'e...'
    #>

    [CmdletBinding(DefaultParameterSetName = 'ServicePrincipal')]
    param(
        [Parameter(Mandatory)]
        [ValidateSet('ServicePrincipal', 'ManagedIdentity', 'UserDelegated')]
        [string] $AuthMethod,

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

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

        [Parameter(ParameterSetName = 'ServicePrincipal')]
        [securestring] $ClientSecret,

        [Parameter(ParameterSetName = 'ServicePrincipal')]
        [System.Security.Cryptography.X509Certificates.X509Certificate2] $Certificate,

        [Parameter(ParameterSetName = 'ManagedIdentity')]
        [string] $ManagedIdentityClientId,

        [Parameter(Mandatory)]
        [string] $WorkspaceId,

        [Parameter(Mandatory)]
        [string] $EventhouseId,

        [string] $DatabaseName = 'FDAObs'
    )

    $script:FDAState.AuthMethod = $AuthMethod
    $script:FDAState.TenantId = $TenantId
    $script:FDAState.WorkspaceId = $WorkspaceId
    $script:FDAState.EventhouseId = $EventhouseId
    $script:FDAState.DatabaseName = $DatabaseName

    # Install token-provider closures keyed by scope. '*' is the fallback.
    switch ($AuthMethod) {
        'ServicePrincipal' {
            if (-not $TenantId -or -not $ClientId) {
                throw 'ServicePrincipal requires -TenantId and -ClientId.'
            }
            if (-not $ClientSecret -and -not $Certificate) {
                throw 'ServicePrincipal requires either -ClientSecret or -Certificate.'
            }
            $tenant = $TenantId
            $cid = $ClientId
            $sec = $ClientSecret
            $cert = $Certificate
            $provider = {
                param($Scope)
                $tokenUrl = "https://login.microsoftonline.com/$tenant/oauth2/v2.0/token"
                $form = @{
                    client_id  = $cid
                    grant_type = 'client_credentials'
                    scope      = $Scope
                }
                if ($cert) {
                    # Client assertion (cert) flow.
                    $jwt = New-FDAClientAssertion -ClientId $cid -TenantId $tenant -Certificate $cert
                    $form['client_assertion_type'] = 'urn:ietf:params:oauth:client-assertion-type:jwt-bearer'
                    $form['client_assertion'] = $jwt
                } else {
                    $plain = [System.Net.NetworkCredential]::new('', $sec).Password
                    $form['client_secret'] = $plain
                }
                $resp = Invoke-RestMethod -Method Post -Uri $tokenUrl -Body $form -ErrorAction Stop
                [pscustomobject]@{
                    Token     = $resp.access_token
                    ExpiresOn = (Get-Date).AddSeconds([int]$resp.expires_in)
                }
            }.GetNewClosure()
            $script:FDAState.TokenProviders['*'] = $provider
        }
        'ManagedIdentity' {
            $miCid = $ManagedIdentityClientId
            $provider = {
                param($Scope)
                # IMDS endpoint. resource = scope-without-/.default suffix.
                $resource = $Scope -replace '/\.default$', ''
                $imds = 'http://169.254.169.254/metadata/identity/oauth2/token?api-version=2018-02-01&resource={0}' -f [uri]::EscapeDataString($resource)
                if ($miCid) {
                    $imds += '&client_id=' + [uri]::EscapeDataString($miCid)
                }
                $headers = @{ Metadata = 'true' }
                # Azure Arc / App Service variants would use env-supplied endpoints. Detect:
                if ($env:IDENTITY_ENDPOINT -and $env:IDENTITY_HEADER) {
                    $imds = '{0}?resource={1}&api-version=2019-08-01' -f $env:IDENTITY_ENDPOINT, [uri]::EscapeDataString($resource)
                    if ($miCid) { $imds += '&client_id=' + [uri]::EscapeDataString($miCid) }
                    $headers = @{ 'X-IDENTITY-HEADER' = $env:IDENTITY_HEADER }
                }
                $resp = Invoke-RestMethod -Method Get -Uri $imds -Headers $headers -ErrorAction Stop
                $expires = if ($resp.expires_on) {
                    [DateTimeOffset]::FromUnixTimeSeconds([long]$resp.expires_on).LocalDateTime
                } else {
                    (Get-Date).AddSeconds([int]$resp.expires_in)
                }
                [pscustomobject]@{
                    Token     = $resp.access_token
                    ExpiresOn = $expires
                }
            }.GetNewClosure()
            $script:FDAState.TokenProviders['*'] = $provider
        }
        'UserDelegated' {
            if (-not $TenantId) { throw 'UserDelegated requires -TenantId.' }
            # Default to the Power BI / Azure PowerShell well-known public client
            # if no ClientId was supplied. Device code flow.
            if (-not $ClientId) { $ClientId = '1950a258-227b-4e31-a9cf-717495945fc8' }
            $tenant = $TenantId
            $cid = $ClientId
            $provider = {
                param($Scope)
                $deviceUrl = "https://login.microsoftonline.com/$tenant/oauth2/v2.0/devicecode"
                $form = @{ client_id = $cid; scope = $Scope }
                $dc = Invoke-RestMethod -Method Post -Uri $deviceUrl -Body $form -ErrorAction Stop
                Write-Host ''
                Write-Host '====================================================================='
                Write-Host 'Open a browser to:' -ForegroundColor Yellow
                Write-Host " $($dc.verification_uri)" -ForegroundColor Cyan
                Write-Host 'And enter code:' -ForegroundColor Yellow
                Write-Host " $($dc.user_code)" -ForegroundColor Cyan
                Write-Host '====================================================================='
                Write-Host ''
                $tokenUrl = "https://login.microsoftonline.com/$tenant/oauth2/v2.0/token"
                $deadline = (Get-Date).AddSeconds([int]$dc.expires_in)
                while ((Get-Date) -lt $deadline) {
                    Start-Sleep -Seconds ([int]$dc.interval)
                    $tokenForm = @{
                        client_id   = $cid
                        grant_type  = 'urn:ietf:params:oauth:grant-type:device_code'
                        device_code = $dc.device_code
                    }
                    try {
                        $resp = Invoke-RestMethod -Method Post -Uri $tokenUrl -Body $tokenForm -ErrorAction Stop
                        return [pscustomobject]@{
                            Token     = $resp.access_token
                            ExpiresOn = (Get-Date).AddSeconds([int]$resp.expires_in)
                        }
                    } catch {
                        $err = $_.ErrorDetails.Message | ConvertFrom-Json -ErrorAction SilentlyContinue
                        if ($err -and $err.error -eq 'authorization_pending') { continue }
                        throw
                    }
                }
                throw 'Device code flow timed out before user completed sign-in.'
            }.GetNewClosure()
            $script:FDAState.TokenProviders['*'] = $provider
        }
    }

    $script:FDAState.Connected = $true

    # Resolve Eventhouse endpoints.
    $endpoints = Get-FDAEventhouseEndpoint -WorkspaceId $WorkspaceId -EventhouseId $EventhouseId
    $script:FDAState.EventhouseClusterUri = $endpoints.QueryServiceUri
    $script:FDAState.EventhouseIngestUri = $endpoints.IngestionServiceUri

    # Load config & log levels from the database (best-effort; empty on new install).
    try { Get-FDAObservabilityConfig | Out-Null } catch { Write-Verbose "Config load skipped: $($_.Exception.Message)" }
    try { Get-FDALogLevel | Out-Null }            catch { Write-Verbose "Log levels load skipped: $($_.Exception.Message)" }

    # Drain any spooled events from the previous session.
    try { Restore-FDASpool } catch { Write-Verbose "Spool restore skipped: $($_.Exception.Message)" }

    # Start the background flush timer.
    Start-FDAFlushTimer

    [pscustomobject]@{
        Connected           = $true
        AuthMethod          = $AuthMethod
        WorkspaceId         = $WorkspaceId
        EventhouseId        = $EventhouseId
        EventhouseDisplay   = $endpoints.DisplayName
        DatabaseName        = $DatabaseName
        ClusterUri          = $endpoints.QueryServiceUri
        IngestionUri        = $endpoints.IngestionServiceUri
        SessionId           = $script:FDAState.SessionId
    }
}

function New-FDAClientAssertion {
    <#
    .SYNOPSIS
        Build a signed JWT client assertion for cert-based SP auth.
    #>

    [CmdletBinding()]
    param(
        [Parameter(Mandatory)] [string] $ClientId,
        [Parameter(Mandatory)] [string] $TenantId,
        [Parameter(Mandatory)] [System.Security.Cryptography.X509Certificates.X509Certificate2] $Certificate
    )
    # Build header.
    $thumbprintBytes = [System.Convert]::FromHexString($Certificate.Thumbprint)
    $x5t = [System.Convert]::ToBase64String($thumbprintBytes).Replace('+','-').Replace('/','_').TrimEnd('=')
    $header = @{ alg = 'RS256'; typ = 'JWT'; x5t = $x5t } | ConvertTo-Json -Compress
    $now = [DateTimeOffset]::UtcNow
    $payload = @{
        aud = "https://login.microsoftonline.com/$TenantId/oauth2/v2.0/token"
        iss = $ClientId
        sub = $ClientId
        jti = [guid]::NewGuid().ToString()
        nbf = $now.ToUnixTimeSeconds()
        exp = $now.AddMinutes(10).ToUnixTimeSeconds()
    } | ConvertTo-Json -Compress
    function _b64url([string]$s) {
        [System.Convert]::ToBase64String([Text.Encoding]::UTF8.GetBytes($s)).Replace('+','-').Replace('/','_').TrimEnd('=')
    }
    function _b64url_bytes([byte[]]$b) {
        [System.Convert]::ToBase64String($b).Replace('+','-').Replace('/','_').TrimEnd('=')
    }
    $headerB64 = _b64url $header
    $payloadB64 = _b64url $payload
    $signingInput = "$headerB64.$payloadB64"
    $rsa = $Certificate.GetRSAPrivateKey()
    if (-not $rsa) { throw 'Certificate does not expose an RSA private key.' }
    $signature = $rsa.SignData([Text.Encoding]::UTF8.GetBytes($signingInput), [Security.Cryptography.HashAlgorithmName]::SHA256, [Security.Cryptography.RSASignaturePadding]::Pkcs1)
    $sigB64 = _b64url_bytes $signature
    "$signingInput.$sigB64"
}