Public/Connect-IdentityOps.ps1

function Connect-IdentityOps {
    <#
    .SYNOPSIS
        Connects to Microsoft Graph with the required scopes for IdentityOps commands.
    .DESCRIPTION
        Supports all major authentication flows: Interactive, Device Code, Client Secret,
        Client Certificate (thumbprint or file), Managed Identity, and Access Token.
    .EXAMPLE
        Connect-IdentityOps
        # Interactive browser-based sign-in (default)
    .EXAMPLE
        Connect-IdentityOps -DeviceCode
        # Device code flow for headless / SSH sessions
    .EXAMPLE
        Connect-IdentityOps -TenantId $tid -ClientId $cid -ClientSecret $sec
        # App-only with client secret (SecureString)
    .EXAMPLE
        Connect-IdentityOps -TenantId $tid -ClientId $cid -CertificateThumbprint "AB12..."
        # App-only with certificate thumbprint from local cert store
    .EXAMPLE
        Connect-IdentityOps -ManagedIdentity
        # System-assigned managed identity (Azure hosted)
    .EXAMPLE
        Connect-IdentityOps -AccessToken $secureToken
        # Pre-acquired access token (SecureString)
    #>

    [CmdletBinding(DefaultParameterSetName = 'Interactive')]
    param(
        # ── Interactive (default) ──────────────────────────────────────────────
        [Parameter(ParameterSetName = 'Interactive')]
        [switch]$Interactive,

        # ── Device Code ────────────────────────────────────────────────────────
        [Parameter(ParameterSetName = 'DeviceCode', Mandatory)]
        [switch]$DeviceCode,

        # ── Tenant & Client for app-only flows ─────────────────────────────────
        [Parameter(ParameterSetName = 'ClientSecret', Mandatory)]
        [Parameter(ParameterSetName = 'CertThumbprint', Mandatory)]
        [Parameter(ParameterSetName = 'CertFile', Mandatory)]
        [ValidatePattern('^[0-9a-fA-F]{8}-([0-9a-fA-F]{4}-){3}[0-9a-fA-F]{12}$',
            ErrorMessage = 'TenantId must be a valid GUID.')]
        [string]$TenantId,

        [Parameter(ParameterSetName = 'ClientSecret', Mandatory)]
        [Parameter(ParameterSetName = 'CertThumbprint', Mandatory)]
        [Parameter(ParameterSetName = 'CertFile', Mandatory)]
        [Parameter(ParameterSetName = 'ManagedIdentityUser', Mandatory)]
        [ValidatePattern('^[0-9a-fA-F]{8}-([0-9a-fA-F]{4}-){3}[0-9a-fA-F]{12}$',
            ErrorMessage = 'ClientId must be a valid GUID.')]
        [string]$ClientId,

        # ── Client Secret ──────────────────────────────────────────────────────
        [Parameter(ParameterSetName = 'ClientSecret', Mandatory)]
        [System.Security.SecureString]$ClientSecret,

        # ── Certificate Thumbprint ─────────────────────────────────────────────
        [Parameter(ParameterSetName = 'CertThumbprint', Mandatory)]
        [ValidatePattern('^[0-9A-Fa-f]{40}$',
            ErrorMessage = 'CertificateThumbprint must be a 40-character hex string.')]
        [string]$CertificateThumbprint,

        # ── Certificate File ──────────────────────────────────────────────────
        [Parameter(ParameterSetName = 'CertFile', Mandatory)]
        [ValidateScript({
            if (-not (Test-Path $_ -PathType Leaf)) {
                throw "Certificate file not found: $_"
            }
            $true
        })]
        [string]$CertificatePath,

        [Parameter(ParameterSetName = 'CertFile')]
        [System.Security.SecureString]$CertificatePassword,

        # ── Managed Identity ──────────────────────────────────────────────────
        [Parameter(ParameterSetName = 'ManagedIdentity', Mandatory)]
        [Parameter(ParameterSetName = 'ManagedIdentityUser', Mandatory)]
        [switch]$ManagedIdentity,

        # ── Access Token ──────────────────────────────────────────────────────
        [Parameter(ParameterSetName = 'AccessToken', Mandatory)]
        [System.Security.SecureString]$AccessToken,

        # ── Common ────────────────────────────────────────────────────────────
        [string[]]$Scopes,

        [switch]$NoBanner
    )

    # ── Default scopes for delegated flows ─────────────────────────────────
    $defaultScopes = @(
        'Application.Read.All',
        'AuditLog.Read.All',
        'Directory.Read.All',
        'Group.Read.All',
        'Policy.Read.All',
        'RoleManagement.Read.Directory',
        'User.Read.All',
        'UserAuthenticationMethod.Read.All',
        'CrossTenantInformation.ReadBasic.All'
    )

    if (-not $Scopes) {
        $Scopes = $defaultScopes
    }

    # ── Disconnect any existing session ────────────────────────────────────
    try { Disconnect-MgGraph -ErrorAction SilentlyContinue } catch { }

    # ── Connect based on chosen flow ───────────────────────────────────────
    $authFlowName = $PSCmdlet.ParameterSetName
    try {
        switch ($PSCmdlet.ParameterSetName) {

            'Interactive' {
                Write-IOLog 'Connecting via interactive browser sign-in...' -Level Verbose
                Connect-MgGraph -Scopes $Scopes -ErrorAction Stop -NoWelcome
                $authFlowName = 'Interactive (Browser)'
            }

            'DeviceCode' {
                Write-IOLog 'Connecting via device code flow...' -Level Verbose
                Connect-MgGraph -Scopes $Scopes -UseDeviceCode -ErrorAction Stop -NoWelcome
                $authFlowName = 'Device Code'
            }

            'ClientSecret' {
                Write-IOLog 'Connecting via client secret...' -Level Verbose
                $credential = [PSCredential]::new($ClientId, $ClientSecret)
                Connect-MgGraph -TenantId $TenantId -ClientSecretCredential $credential -ErrorAction Stop -NoWelcome
                $authFlowName = 'Client Secret (App-Only)'
            }

            'CertThumbprint' {
                Write-IOLog 'Connecting via certificate thumbprint...' -Level Verbose
                Connect-MgGraph -TenantId $TenantId -ClientId $ClientId -CertificateThumbprint $CertificateThumbprint -ErrorAction Stop -NoWelcome
                $authFlowName = 'Certificate Thumbprint (App-Only)'
            }

            'CertFile' {
                Write-IOLog 'Connecting via certificate file...' -Level Verbose
                $resolvedCertPath = (Resolve-Path $CertificatePath).Path
                $bstrPtr = [System.IntPtr]::Zero
                try {
                    $plaintextPwd = $null
                    if ($CertificatePassword) {
                        $bstrPtr = [System.Runtime.InteropServices.Marshal]::SecureStringToBSTR($CertificatePassword)
                        $plaintextPwd = [System.Runtime.InteropServices.Marshal]::PtrToStringBSTR($bstrPtr)
                    }
                    $cert = [System.Security.Cryptography.X509Certificates.X509Certificate2]::new(
                        $resolvedCertPath,
                        $plaintextPwd,
                        [System.Security.Cryptography.X509Certificates.X509KeyStorageFlags]::EphemeralKeySet
                    )
                }
                finally {
                    # Zero and free BSTR to prevent secret lingering in memory
                    if ($bstrPtr -ne [System.IntPtr]::Zero) {
                        [System.Runtime.InteropServices.Marshal]::ZeroFreeBSTR($bstrPtr)
                    }
                    $plaintextPwd = $null
                }
                try {
                    Connect-MgGraph -TenantId $TenantId -ClientId $ClientId -Certificate $cert -ErrorAction Stop -NoWelcome
                }
                finally {
                    # Dispose certificate to release private key material from memory
                    if ($cert) { $cert.Dispose() }
                }
                $authFlowName = 'Certificate File (App-Only)'
            }

            'ManagedIdentity' {
                Write-IOLog 'Connecting via system-assigned managed identity...' -Level Verbose
                Connect-MgGraph -Identity -ErrorAction Stop -NoWelcome
                $authFlowName = 'Managed Identity (System)'
            }

            'ManagedIdentityUser' {
                Write-IOLog 'Connecting via user-assigned managed identity...' -Level Verbose
                Connect-MgGraph -Identity -ClientId $ClientId -ErrorAction Stop -NoWelcome
                $authFlowName = 'Managed Identity (User-Assigned)'
            }

            'AccessToken' {
                Write-IOLog 'Connecting via pre-acquired access token...' -Level Verbose
                Connect-MgGraph -AccessToken $AccessToken -ErrorAction Stop -NoWelcome
                $authFlowName = 'Access Token'
            }
        }
    }
    catch {
        $script:IOConnection.Connected = $false
        throw [System.Management.Automation.ErrorRecord]::new(
            [System.Security.Authentication.AuthenticationException]::new(
                "Authentication failed ($authFlowName): $($_.Exception.Message)"
            ),
            'IO_AuthenticationFailed',
            [System.Management.Automation.ErrorCategory]::AuthenticationError,
            $authFlowName
        )
    }

    # ── Populate connection state ──────────────────────────────────────────
    $ctx = Get-MgContext -ErrorAction SilentlyContinue

    $tenantName = $null
    try {
        $org = Invoke-MgGraphRequest -Method GET -Uri 'v1.0/organization?$select=displayName,id' -ErrorAction Stop -OutputType PSObject
        if ($org.value -and $org.value.Count -gt 0) {
            $tenantName = $org.value[0].displayName
        }
    }
    catch {
        Write-IOLog 'Could not retrieve tenant display name.' -Level Verbose
    }

    $script:IOConnection = @{
        Connected         = $true
        TenantId          = $ctx.TenantId
        TenantName        = $tenantName
        UserPrincipalName = $ctx.Account
        AuthFlow          = $authFlowName
        ConnectedAt       = Get-Date
        Scopes            = @($ctx.Scopes)
    }

    # ── Welcome banner ─────────────────────────────────────────────────────
    if (-not $NoBanner) {
        Show-IOWelcome -UserPrincipalName $ctx.Account `
                       -TenantId $ctx.TenantId `
                       -TenantName $tenantName `
                       -AuthFlow $authFlowName
    }
}