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. Optional for ManagedIdentity (taken
        from IMDS metadata) and UserDelegated. When omitted for UserDelegated
        you are prompted for a tenant ID (GUID) or verified domain
        (e.g. contoso.onmicrosoft.com) to sign in to — the device-code flow
        needs a concrete tenant, so it cannot be discovered after sign-in.

    .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. Optional — when
        omitted, the module lists the workspaces you can access and prompts
        you to select one or create a new workspace.

    .PARAMETER EventhouseId
        Fabric Eventhouse item id. Endpoints are resolved via Fabric REST.
        Optional — when omitted, the module lists the Eventhouses in the
        selected workspace and prompts you to select one or create a new
        Eventhouse.

    .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
        # Fully interactive: sign in, pick (or create) tenant / workspace / Eventhouse.
        Connect-FDAObservability -AuthMethod UserDelegated

    .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,

        [string] $WorkspaceId,

        [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' {
            # 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' }
            $cid = $ClientId
            $provider = {
                param($Scope)
                # The raw device-code endpoint needs a concrete tenant: the
                # tenant-less 'organizations'/'common' authorities are rejected
                # with AADSTS50059. Connect-FDAObservability resolves a tenant
                # (parameter or interactive prompt) before any token is fetched.
                $tenant = $script:FDAState.TenantId
                if (-not $tenant) {
                    throw 'UserDelegated sign-in requires a tenant. Pass -TenantId, or supply a tenant ID/domain when prompted.'
                }
                $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 tenant / workspace / Eventhouse. Anything not supplied as a
    # parameter is selected interactively (this is where the device-code
    # sign-in happens for UserDelegated).
    # ---------------------------------------------------------------------

    # Tenant: only UserDelegated needs an interactive prompt. SP uses the
    # supplied -TenantId; ManagedIdentity takes it from IMDS. The device-code
    # flow can't bootstrap without a concrete tenant, so ask for one up front.
    if (-not $script:FDAState.TenantId -and $AuthMethod -eq 'UserDelegated') {
        $script:FDAState.TenantId = Resolve-FDATenant
        Write-Host "Signing in to tenant: $($script:FDAState.TenantId)" -ForegroundColor Green
    }

    # Workspace.
    if (-not $WorkspaceId) {
        $WorkspaceId = Resolve-FDAWorkspace
    }
    $script:FDAState.WorkspaceId = $WorkspaceId

    # Eventhouse.
    if (-not $EventhouseId) {
        $EventhouseId = Resolve-FDAEventhouse -WorkspaceId $WorkspaceId
    }
    $script:FDAState.EventhouseId = $EventhouseId

    # 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()]
    [OutputType([string])]
    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"
}