Get-Snapshot.ps1

# Syskit Discovery - Script Runner
# This script runs Syskit Discovery: M365 Security Assessment via the M365Snapshot module
# Usage examples:
# .\Get-Snapshot.ps1 -TenantId "contoso.onmicrosoft.com" -ClientId "existing-app-id"
# .\Get-Snapshot.ps1 -TenantId "contoso.onmicrosoft.com"
# .\Get-Snapshot.ps1 -TenantId "contoso.onmicrosoft.com" -RegisterClient
# .\Get-Snapshot.ps1 -TenantId "contoso.onmicrosoft.com" -Scopes Applications,Exchange

[CmdletBinding(PositionalBinding=$false)]
param(
    [Parameter(Mandatory=$true)]
    [string]$TenantId,

    [Parameter(Mandatory=$false)]
    [string]$ClientId,

    [Parameter(Mandatory=$false)]
    [switch]$RegisterClient,

    [Parameter(Mandatory=$false)]
    [switch]$GenerateHTMLReport,

    [Parameter(Mandatory=$false)]
    [switch]$NoGenerateHTMLReport,

    [Parameter(Mandatory=$false)]
    [string]$ReportPath,

    [Parameter(Mandatory=$false)]
    [ValidateRange(1, [int]::MaxValue)]
    [int]$MaxUsers = 10000,

    [Parameter(Mandatory=$false)]
    [ValidateRange(1, [int]::MaxValue)]
    [int]$MaxGroups = 10000,

    [Parameter(Mandatory=$false)]
    [ValidateRange(1, [int]::MaxValue)]
    [int]$MaxSites = 1000,

    [Parameter(Mandatory=$false)]
    [switch]$LoadAllUsers,

    [Parameter(Mandatory=$false)]
    [switch]$LoadAllGroups,

    [Parameter(Mandatory=$false)]
    [switch]$LoadAllSites,

    [Parameter(Mandatory=$false)]
    [ValidateSet('All', 'Entra', 'Exchange', 'Applications', 'SharePoint')]
    [string[]]$Scopes = @('All'),

    [Parameter(Mandatory=$false)]
    [ValidateRange(1, [int]::MaxValue)]
    [int]$MaxAppRegistrations = 1000,

    [Parameter(Mandatory=$false)]
    [switch]$LoadAllAppRegistrations,

    [Parameter(Mandatory=$false)]
    [ValidateRange(1, [int]::MaxValue)]
    [int]$MaxSharedMailboxes = 100,

    [Parameter(Mandatory=$false)]
    [ValidateRange(1, [int]::MaxValue)]
    [int]$MaxSecurityDistributionGroups = 100,

    [Parameter(Mandatory=$false)]
    [string[]]$ConfidentialSensitivityLabels = @(),

    [Parameter(Mandatory=$false)]
    [string]$ConfidentialSensitivityLabelsFile,

    [Parameter(Mandatory=$false)]
    [switch]$NoRelaunch
)

$script:AppName = "Syskit Discovery"
$script:SitesDelegatedScope = "Sites.FullControl.All"
$script:GroupDelegatedScope = "Group.Read.All"
$script:ReportsDelegatedScope = "Reports.Read.All"
$script:ApplicationsDelegatedScope = "Application.Read.All"
$script:AuditLogsDelegatedScope = "AuditLog.Read.All"
$script:MailboxSettingsDelegatedScope = "MailboxSettings.Read"
$script:AppRegistrationContractVersion = "1.0.0"
$script:AppStorageName = ($script:AppName -replace "[^a-zA-Z0-9]", "").Trim()
if ([string]::IsNullOrWhiteSpace($script:AppStorageName)) {
    $script:AppStorageName = "App"
}

if ($GenerateHTMLReport -and $NoGenerateHTMLReport) {
    Write-Host "ERROR: Use either -GenerateHTMLReport or -NoGenerateHTMLReport, not both." -ForegroundColor Red
    exit 1
}

if ($NoGenerateHTMLReport) {
    $GenerateHTMLReport = $false
}
elseif (-not $PSBoundParameters.ContainsKey('GenerateHTMLReport')) {
    $GenerateHTMLReport = $true
}

function Resolve-EnabledFeaturesFromScopes {
    param(
        [string[]]$SelectedScopes
    )

    $effectiveSelectedScopes = @()
    if ($null -ne $SelectedScopes) {
        $effectiveSelectedScopes = @($SelectedScopes | Where-Object { -not [string]::IsNullOrWhiteSpace([string]$_) } | ForEach-Object { ([string]$_).Trim() })
    }

    if ($effectiveSelectedScopes.Count -eq 0) {
        $effectiveSelectedScopes = @('All')
    }

    $allSelected = @($effectiveSelectedScopes | Where-Object { $_ -eq 'All' }).Count -gt 0

    return @{
        SelectedScopes = $effectiveSelectedScopes
        IncludeAppRegistrations = [bool]($allSelected -or ($effectiveSelectedScopes -contains 'Applications'))
        IncludeSharedMailboxes = [bool]($allSelected -or ($effectiveSelectedScopes -contains 'Exchange'))
        IncludeSecurityDistributionGroups = [bool]($allSelected -or ($effectiveSelectedScopes -contains 'Entra'))
    }
}

$enabledFeatures = Resolve-EnabledFeaturesFromScopes -SelectedScopes $Scopes
$IncludeAppRegistrations = [bool]$enabledFeatures.IncludeAppRegistrations
$IncludeSharedMailboxes = [bool]$enabledFeatures.IncludeSharedMailboxes
$IncludeSecurityDistributionGroups = [bool]$enabledFeatures.IncludeSecurityDistributionGroups

$global:SyskitDiscoveryTranscriptActive = $false

function Start-RunTranscript {
    param(
        [string]$StorageName,
        [switch]$Quiet
    )

    try {
        $logRoot = Join-Path $env:LOCALAPPDATA $StorageName
        $logDir = Join-Path $logRoot "Logs"
        if (-not (Test-Path $logDir)) {
            New-Item -Path $logDir -ItemType Directory -Force | Out-Null
        }

        $timestamp = (Get-Date).ToString("yyyyMMdd_HHmmss")
        $logPath = Join-Path $logDir ("Snapshot_{0}_{1}.log" -f $timestamp, $PID)
        Start-Transcript -Path $logPath -Force | Out-Null
        $global:SyskitDiscoveryTranscriptActive = $true

        if (-not $Quiet) {
            Write-Host "[INFO] Run log: $logPath" -ForegroundColor DarkGray
        }
    }
    catch {
        if (-not $Quiet) {
            Write-Host "[INFO] Unable to start run transcript logging: $($_.Exception.Message)" -ForegroundColor DarkGray
        }
    }
}

if (-not (Get-EventSubscriber -ErrorAction SilentlyContinue | Where-Object { $_.SourceIdentifier -eq 'SyskitDiscovery.TranscriptStop' })) {
    Register-EngineEvent -SourceIdentifier PowerShell.Exiting -Action {
        if ($global:SyskitDiscoveryTranscriptActive) {
            try { Stop-Transcript | Out-Null } catch {}
            $global:SyskitDiscoveryTranscriptActive = $false
        }
    } -SupportEvent -ErrorAction SilentlyContinue | Out-Null
}

if ($NoRelaunch) {
    Start-RunTranscript -StorageName $script:AppStorageName
}

if ([string]::IsNullOrWhiteSpace($ReportPath)) {
    $ReportPath = "$($script:AppStorageName)_Report.html"
}

if (-not [string]::IsNullOrWhiteSpace($ConfidentialSensitivityLabelsFile) -and (Test-Path $ConfidentialSensitivityLabelsFile)) {
    try {
        $labelsFromFile = Get-Content -Path $ConfidentialSensitivityLabelsFile -Raw | ConvertFrom-Json
        if ($labelsFromFile -is [System.Array]) {
            $ConfidentialSensitivityLabels = @($labelsFromFile | ForEach-Object { [string]$_ })
        }
        elseif ($null -ne $labelsFromFile) {
            $ConfidentialSensitivityLabels = @([string]$labelsFromFile)
        }
    }
    catch {
        Write-Host "WARNING: Failed to load ConfidentialSensitivityLabels from file '$ConfidentialSensitivityLabelsFile'." -ForegroundColor Yellow
    }
    finally {
        try {
            Remove-Item -Path $ConfidentialSensitivityLabelsFile -Force -ErrorAction SilentlyContinue
        }
        catch { }
    }
}

if (-not $NoRelaunch) {
    $pwshCmd = Get-Command -Name pwsh -ErrorAction SilentlyContinue
    if ($pwshCmd) {
        Write-Host "Starting clean PowerShell session for dependency isolation..." -ForegroundColor DarkGray

        $childArgs = @(
            "-NoProfile",
            "-ExecutionPolicy", "Bypass",
            "-File", $PSCommandPath,
            "-TenantId", $TenantId,
            "-NoRelaunch"
        )

        if (-not [string]::IsNullOrWhiteSpace($ClientId)) {
            $childArgs += @("-ClientId", $ClientId)
        }
        if ($RegisterClient) {
            $childArgs += "-RegisterClient"
        }
        if ($GenerateHTMLReport) {
            $childArgs += "-GenerateHTMLReport"
        }
        if ($NoGenerateHTMLReport) {
            $childArgs += "-NoGenerateHTMLReport"
        }
        if ($PSBoundParameters.ContainsKey("ReportPath") -and -not [string]::IsNullOrWhiteSpace($ReportPath)) {
            $childArgs += @("-ReportPath", $ReportPath)
        }
        if ($PSBoundParameters.ContainsKey("MaxUsers")) {
            $childArgs += @("-MaxUsers", $MaxUsers)
        }
        if ($PSBoundParameters.ContainsKey("MaxGroups")) {
            $childArgs += @("-MaxGroups", $MaxGroups)
        }
        if ($PSBoundParameters.ContainsKey("MaxSites")) {
            $childArgs += @("-MaxSites", $MaxSites)
        }
        if ($LoadAllUsers) {
            $childArgs += "-LoadAllUsers"
        }
        if ($LoadAllGroups) {
            $childArgs += "-LoadAllGroups"
        }
        if ($LoadAllSites) {
            $childArgs += "-LoadAllSites"
        }
        if ($PSBoundParameters.ContainsKey("Scopes")) {
            foreach ($scopeName in @($Scopes)) {
                $childArgs += @("-Scopes", $scopeName)
            }
        }
        if ($PSBoundParameters.ContainsKey("MaxAppRegistrations")) {
            $childArgs += @("-MaxAppRegistrations", $MaxAppRegistrations)
        }
        if ($LoadAllAppRegistrations) {
            $childArgs += "-LoadAllAppRegistrations"
        }
        if ($PSBoundParameters.ContainsKey("MaxSharedMailboxes")) {
            $childArgs += @("-MaxSharedMailboxes", $MaxSharedMailboxes)
        }
        if ($PSBoundParameters.ContainsKey("MaxSecurityDistributionGroups")) {
            $childArgs += @("-MaxSecurityDistributionGroups", $MaxSecurityDistributionGroups)
        }
        if ($ConfidentialSensitivityLabels.Count -gt 0) {
            $nonEmptyLabels = @($ConfidentialSensitivityLabels | Where-Object { -not [string]::IsNullOrWhiteSpace([string]$_) })
            if ($nonEmptyLabels.Count -gt 0) {
                $labelsFile = Join-Path ([System.IO.Path]::GetTempPath()) ("{0}_{1}.json" -f $script:AppStorageName, ([guid]::NewGuid().ToString("N")))
                $nonEmptyLabels | ConvertTo-Json -Depth 3 | Set-Content -Path $labelsFile -Encoding UTF8
                $childArgs += @("-ConfidentialSensitivityLabelsFile", $labelsFile)
            }
        }

        & $pwshCmd.Source @childArgs
        exit $LASTEXITCODE
    }

    Start-RunTranscript -StorageName $script:AppStorageName
}

$configDir = Join-Path $env:LOCALAPPDATA $script:AppStorageName
$configPath = Join-Path $configDir "config.json"
$legacyConfigDir = Join-Path $env:LOCALAPPDATA "M365Snapshot"
$legacyConfigPath = Join-Path $legacyConfigDir "config.json"

function Get-LocalConfig {
    $sourceConfigPath = $null
    if (Test-Path $configPath) {
        $sourceConfigPath = $configPath
    }
    elseif (Test-Path $legacyConfigPath) {
        $sourceConfigPath = $legacyConfigPath
    }

    if (-not $sourceConfigPath) {
        return @{
            appName = $script:AppName
            tenants = @{}
        }
    }

    try {
        $raw = Get-Content -Path $sourceConfigPath -Raw
        $parsed = ConvertFrom-Json -InputObject $raw -AsHashtable
        if (-not $parsed.ContainsKey("tenants")) {
            $parsed["tenants"] = @{}
        }
        if (-not $parsed.ContainsKey("appName")) {
            $parsed["appName"] = $script:AppName
        }
        return $parsed
    }
    catch {
        Write-Host "WARNING: Failed to parse config at $sourceConfigPath. A new config will be created." -ForegroundColor Yellow
        return @{
            appName = $script:AppName
            tenants = @{}
        }
    }
}

function Save-LocalConfig {
    param(
        [hashtable]$Config
    )

    if (-not (Test-Path $configDir)) {
        New-Item -Path $configDir -ItemType Directory -Force | Out-Null
    }

    $Config | ConvertTo-Json -Depth 10 | Set-Content -Path $configPath -Encoding UTF8
}

function Save-TenantClientId {
    param(
        [string]$Tenant,
        [string]$ResolvedClientId,
        [string]$RegisteredAppDisplayName,
        [string]$ContractVersion,
        [string]$ContractHash,
        [string[]]$DelegatedScopes = @()
    )

    $config = Get-LocalConfig
    $config["appName"] = $script:AppName
    $config["tenants"][$Tenant] = @{
        clientId = $ResolvedClientId
        appDisplayName = $RegisteredAppDisplayName
        appRegistrationContractVersion = $ContractVersion
        appRegistrationContractHash = $ContractHash
        delegatedScopes = @($DelegatedScopes | Sort-Object -Unique)
        updatedUtc = (Get-Date).ToUniversalTime().ToString("o")
    }
    Save-LocalConfig -Config $config
}

function Get-StoredTenantEntry {
    param([string]$Tenant)

    $config = Get-LocalConfig
    if ($config["tenants"].ContainsKey($Tenant)) {
        return $config["tenants"][$Tenant]
    }

    return $null
}

function Get-StoredClientId {
    param([string]$Tenant)

    $entry = Get-StoredTenantEntry -Tenant $Tenant
    if ($null -ne $entry -and $entry.ContainsKey("clientId") -and -not [string]::IsNullOrWhiteSpace($entry["clientId"])) {
        return [string]$entry["clientId"]
    }

    return $null
}

function Test-IsGuid {
    param([string]$Value)

    if ([string]::IsNullOrWhiteSpace($Value)) { return $false }
    return ($Value.Trim() -match '^[0-9a-fA-F]{8}\-[0-9a-fA-F]{4}\-[0-9a-fA-F]{4}\-[0-9a-fA-F]{4}\-[0-9a-fA-F]{12}$')
}

function Get-AppRegistrationContract {
    param(
        [switch]$EnableAppRegistrationRead,
        [switch]$EnableSharedMailboxRead
    )

    $requiredScopes = @(
        "User.Read.All",
        $script:GroupDelegatedScope,
        $script:ReportsDelegatedScope,
        $script:SitesDelegatedScope
    )

    if ($EnableAppRegistrationRead) {
        $requiredScopes += $script:ApplicationsDelegatedScope
        $requiredScopes += $script:AuditLogsDelegatedScope
    }

    if ($EnableSharedMailboxRead) {
        $requiredScopes += $script:MailboxSettingsDelegatedScope
    }

    $contractBody = [ordered]@{
        appName = $script:AppName
        graphResourceAppId = "00000003-0000-0000-c000-000000000000"
        delegatedScopes = @($requiredScopes | Sort-Object -Unique)
    }

    $contractJson = $contractBody | ConvertTo-Json -Compress -Depth 8
    $contractBytes = [System.Text.Encoding]::UTF8.GetBytes($contractJson)
    $sha256 = [System.Security.Cryptography.SHA256]::Create()
    try {
        $hashBytes = $sha256.ComputeHash($contractBytes)
        $contractHash = ([System.BitConverter]::ToString($hashBytes)).Replace('-', '').ToLowerInvariant()
    }
    finally {
        $sha256.Dispose()
    }

    return @{
        Version = $script:AppRegistrationContractVersion
        Hash = $contractHash
        DelegatedScopes = @($contractBody.delegatedScopes)
    }
}

function Test-AppContractSatisfiesRequiredScopes {
    param(
        [object]$StoredEntry,
        [string[]]$RequiredScopes
    )

    if ($null -eq $StoredEntry -or -not $StoredEntry.ContainsKey("delegatedScopes")) {
        return $false
    }

    $storedScopes = @($StoredEntry["delegatedScopes"] | ForEach-Object { [string]$_ } | Where-Object { -not [string]::IsNullOrWhiteSpace($_) } | Sort-Object -Unique)
    if ($storedScopes.Count -eq 0) {
        return $false
    }

    foreach ($requiredScope in @($RequiredScopes | Sort-Object -Unique)) {
        if ($storedScopes -notcontains $requiredScope) {
            return $false
        }
    }

    return $true
}

function Register-M365SnapshotClient {
    param(
        [string]$Tenant,
        [switch]$EnableAppRegistrationRead,
        [switch]$EnableSharedMailboxRead
    )

    $requiredModules = @(
        "Microsoft.Graph.Authentication",
        "Microsoft.Graph.Applications",
        "Microsoft.Graph.Identity.SignIns"
    )

    foreach ($requiredModule in $requiredModules) {
        if (-not (Get-Module -ListAvailable -Name $requiredModule)) {
            throw "Required module '$requiredModule' is not installed. Install it with: Install-Module $requiredModule -Scope CurrentUser"
        }
    }

    Import-Module Microsoft.Graph.Authentication -Force
    Import-Module Microsoft.Graph.Applications -Force
    Import-Module Microsoft.Graph.Identity.SignIns -Force

    $registrationDisplayName = "$($script:AppName) app"

    Write-Host "Connecting to Microsoft Graph to register app '$registrationDisplayName'..." -ForegroundColor Cyan
    try {
        Connect-MgGraph -TenantId $Tenant -Scopes @(
            "Application.ReadWrite.All",
            "DelegatedPermissionGrant.ReadWrite.All",
            "Directory.Read.All"
        ) | Out-Null
    }
    catch {
        $errorText = $_.Exception.Message
        if ($errorText -match "Could not load type") {
            throw "Microsoft Graph module dependency conflict detected. Close all PowerShell terminals and retry. If issue persists, reinstall Microsoft.Graph modules. Original error: $errorText"
        }
        throw
    }

    $graphAppId = "00000003-0000-0000-c000-000000000000"
    $graphSp = Get-MgServicePrincipal -Filter "appId eq '$graphAppId'"

    if (-not $graphSp) {
        throw "Could not resolve Microsoft Graph service principal in tenant '$Tenant'."
    }

    $graphUserReadAllScope = $graphSp.Oauth2PermissionScopes | Where-Object { $_.Value -eq "User.Read.All" } | Select-Object -First 1
    $graphSitesReadScope = $graphSp.Oauth2PermissionScopes | Where-Object { $_.Value -eq $script:SitesDelegatedScope } | Select-Object -First 1
    $graphGroupReadScope = $graphSp.Oauth2PermissionScopes | Where-Object { $_.Value -eq $script:GroupDelegatedScope } | Select-Object -First 1
    $graphReportsReadScope = $graphSp.Oauth2PermissionScopes | Where-Object { $_.Value -eq $script:ReportsDelegatedScope } | Select-Object -First 1
    $graphApplicationReadScope = if ($EnableAppRegistrationRead) { $graphSp.Oauth2PermissionScopes | Where-Object { $_.Value -eq $script:ApplicationsDelegatedScope } | Select-Object -First 1 } else { $null }
    $graphAuditLogReadScope = if ($EnableAppRegistrationRead) { $graphSp.Oauth2PermissionScopes | Where-Object { $_.Value -eq $script:AuditLogsDelegatedScope } | Select-Object -First 1 } else { $null }
    $graphMailboxSettingsReadScope = if ($EnableSharedMailboxRead) { $graphSp.Oauth2PermissionScopes | Where-Object { $_.Value -eq $script:MailboxSettingsDelegatedScope } | Select-Object -First 1 } else { $null }

    if (-not $graphUserReadAllScope -or -not $graphSitesReadScope -or -not $graphGroupReadScope -or -not $graphReportsReadScope -or ($EnableAppRegistrationRead -and (-not $graphApplicationReadScope -or -not $graphAuditLogReadScope)) -or ($EnableSharedMailboxRead -and -not $graphMailboxSettingsReadScope)) {
        $availableGraphScopes = ($graphSp.Oauth2PermissionScopes | Select-Object -ExpandProperty Value | Sort-Object -Unique) -join ", "
        if ($EnableAppRegistrationRead -and $EnableSharedMailboxRead) {
            throw "Could not resolve required delegated scopes (User.Read.All / $($script:GroupDelegatedScope) / $($script:ReportsDelegatedScope) / $($script:ApplicationsDelegatedScope) / $($script:AuditLogsDelegatedScope) / $($script:MailboxSettingsDelegatedScope) / $($script:SitesDelegatedScope)). Available Graph scopes for this tenant/app: $availableGraphScopes"
        }
        elseif ($EnableAppRegistrationRead) {
            throw "Could not resolve required delegated scopes (User.Read.All / $($script:GroupDelegatedScope) / $($script:ReportsDelegatedScope) / $($script:ApplicationsDelegatedScope) / $($script:AuditLogsDelegatedScope) / $($script:SitesDelegatedScope)). Available Graph scopes for this tenant/app: $availableGraphScopes"
        }
        elseif ($EnableSharedMailboxRead) {
            throw "Could not resolve required delegated scopes (User.Read.All / $($script:GroupDelegatedScope) / $($script:ReportsDelegatedScope) / $($script:MailboxSettingsDelegatedScope) / $($script:SitesDelegatedScope)). Available Graph scopes for this tenant/app: $availableGraphScopes"
        }
        else {
            throw "Could not resolve required delegated scopes (User.Read.All / $($script:GroupDelegatedScope) / $($script:ReportsDelegatedScope) / $($script:SitesDelegatedScope)). Available Graph scopes for this tenant/app: $availableGraphScopes"
        }
    }

    $resourceAccessEntries = @(
        @{
            id = $graphUserReadAllScope.Id
            type = "Scope"
        },
        @{
            id = $graphSitesReadScope.Id
            type = "Scope"
        },
        @{
            id = $graphGroupReadScope.Id
            type = "Scope"
        },
        @{
            id = $graphReportsReadScope.Id
            type = "Scope"
        }
    )

    if ($EnableAppRegistrationRead) {
        $resourceAccessEntries += @{
            id = $graphApplicationReadScope.Id
            type = "Scope"
        }
        $resourceAccessEntries += @{
            id = $graphAuditLogReadScope.Id
            type = "Scope"
        }
    }

    if ($EnableSharedMailboxRead) {
        $resourceAccessEntries += @{
            id = $graphMailboxSettingsReadScope.Id
            type = "Scope"
        }
    }

    $requiredResourceAccess = @(
        @{
            resourceAppId = $graphAppId
            resourceAccess = $resourceAccessEntries
        }
    )

    $application = New-MgApplication `
        -DisplayName $registrationDisplayName `
        -SignInAudience "AzureADMyOrg" `
        -PublicClient @{ redirectUris = @(
            "https://login.microsoftonline.com/common/oauth2/nativeclient",
            "http://localhost"
        ) } `
        -RequiredResourceAccess $requiredResourceAccess

    $servicePrincipal = New-MgServicePrincipal -AppId $application.AppId

    $grantScopes = @("User.Read.All", $graphGroupReadScope.Value, $graphReportsReadScope.Value, $graphSitesReadScope.Value)
    if ($EnableAppRegistrationRead) {
        $grantScopes += $graphApplicationReadScope.Value
        $grantScopes += $graphAuditLogReadScope.Value
    }

    if ($EnableSharedMailboxRead) {
        $grantScopes += $graphMailboxSettingsReadScope.Value
    }

    New-MgOauth2PermissionGrant -BodyParameter @{
        clientId = $servicePrincipal.Id
        consentType = "AllPrincipals"
        resourceId = $graphSp.Id
        scope = ($grantScopes -join " ")
    } | Out-Null

    Write-Host "App registration completed. ClientId: $($application.AppId)" -ForegroundColor Green

    return @{
        ClientId = $application.AppId
        DisplayName = $registrationDisplayName
    }
}

if (-not [string]::IsNullOrWhiteSpace($ClientId)) {
    if (-not (Test-IsGuid -Value $ClientId)) {
        Write-Host "ERROR: -ClientId must be a GUID (appId). Received: '$ClientId'" -ForegroundColor Red
        Write-Host "Tip: Pass multiple confidential labels via -ConfidentialSensitivityLabels, not in -ClientId." -ForegroundColor Yellow
        exit 1
    }

    $currentContract = Get-AppRegistrationContract -EnableAppRegistrationRead:$IncludeAppRegistrations -EnableSharedMailboxRead:$IncludeSharedMailboxes
    Save-TenantClientId -Tenant $TenantId -ResolvedClientId $ClientId -RegisteredAppDisplayName "ProvidedByUser" -ContractVersion $currentContract.Version -ContractHash $currentContract.Hash -DelegatedScopes $currentContract.DelegatedScopes
}
elseif ($RegisterClient) {
    try {
        $registered = Register-M365SnapshotClient -Tenant $TenantId -EnableAppRegistrationRead:$IncludeAppRegistrations -EnableSharedMailboxRead:$IncludeSharedMailboxes
        $ClientId = $registered.ClientId
        $currentContract = Get-AppRegistrationContract -EnableAppRegistrationRead:$IncludeAppRegistrations -EnableSharedMailboxRead:$IncludeSharedMailboxes
        Save-TenantClientId -Tenant $TenantId -ResolvedClientId $ClientId -RegisteredAppDisplayName $registered.DisplayName -ContractVersion $currentContract.Version -ContractHash $currentContract.Hash -DelegatedScopes $currentContract.DelegatedScopes
    }
    catch {
        Write-Host "ERROR: Failed to register client app. $_" -ForegroundColor Red
        exit 1
    }
}
else {
    $currentContract = Get-AppRegistrationContract -EnableAppRegistrationRead:$IncludeAppRegistrations -EnableSharedMailboxRead:$IncludeSharedMailboxes
    $storedClientId = Get-StoredClientId -Tenant $TenantId
    $storedEntry = Get-StoredTenantEntry -Tenant $TenantId

    if (-not [string]::IsNullOrWhiteSpace($storedClientId)) {
        $storedHasNeededScopes = Test-AppContractSatisfiesRequiredScopes -StoredEntry $storedEntry -RequiredScopes $currentContract.DelegatedScopes
        $storedContractVersion = if ($null -ne $storedEntry -and $storedEntry.ContainsKey("appRegistrationContractVersion")) { [string]$storedEntry["appRegistrationContractVersion"] } else { "" }
        $storedContractHash = if ($null -ne $storedEntry -and $storedEntry.ContainsKey("appRegistrationContractHash")) { [string]$storedEntry["appRegistrationContractHash"] } else { "" }
        $contractChanged = -not $storedHasNeededScopes
        if (-not $contractChanged -and -not [string]::IsNullOrWhiteSpace($storedContractVersion) -and -not [string]::IsNullOrWhiteSpace($storedContractHash)) {
            $contractChanged = ($storedContractVersion -ne $currentContract.Version) -or ($storedContractHash -ne $currentContract.Hash)
        }

        if ($contractChanged) {
            Write-Host "Stored app registration does not satisfy selected snapshot scopes for tenant '$TenantId'. Re-registering app..." -ForegroundColor Yellow
            try {
                $registered = Register-M365SnapshotClient -Tenant $TenantId -EnableAppRegistrationRead:$IncludeAppRegistrations -EnableSharedMailboxRead:$IncludeSharedMailboxes
                $ClientId = $registered.ClientId
                Save-TenantClientId -Tenant $TenantId -ResolvedClientId $ClientId -RegisteredAppDisplayName $registered.DisplayName -ContractVersion $currentContract.Version -ContractHash $currentContract.Hash -DelegatedScopes $currentContract.DelegatedScopes
            }
            catch {
                Write-Host "ERROR: App re-registration failed after contract change. $_" -ForegroundColor Red
                Write-Host "Run again with -ClientId <existing-client-id> or -RegisterClient with sufficient admin rights." -ForegroundColor Yellow
                exit 1
            }
        }
        else {
            Write-Host "Using stored ClientId for tenant '$TenantId'." -ForegroundColor Cyan
            $ClientId = $storedClientId
        }
    }
    else {
        Write-Host "No ClientId provided and none stored for '$TenantId'. Attempting app registration..." -ForegroundColor Yellow
        try {
            $registered = Register-M365SnapshotClient -Tenant $TenantId -EnableAppRegistrationRead:$IncludeAppRegistrations -EnableSharedMailboxRead:$IncludeSharedMailboxes
            $ClientId = $registered.ClientId
            Save-TenantClientId -Tenant $TenantId -ResolvedClientId $ClientId -RegisteredAppDisplayName $registered.DisplayName -ContractVersion $currentContract.Version -ContractHash $currentContract.Hash -DelegatedScopes $currentContract.DelegatedScopes
        }
        catch {
            Write-Host "ERROR: Automatic app registration failed. $_" -ForegroundColor Red
            Write-Host "Run again with -ClientId <existing-client-id> or -RegisterClient with sufficient admin rights." -ForegroundColor Yellow
            exit 1
        }
    }
}

if (-not (Test-IsGuid -Value $ClientId)) {
    Write-Host "ERROR: Resolved ClientId is not a GUID. Value: '$ClientId'" -ForegroundColor Red
    Write-Host "Provide a valid appId GUID via -ClientId or re-run with -RegisterClient." -ForegroundColor Yellow
    exit 1
}

# Import the module
$modulePath = Join-Path $PSScriptRoot "Syskit.Discovery.psd1"

if (-not (Test-Path $modulePath)) {
    Write-Host "ERROR: $($script:AppName) module manifest not found." -ForegroundColor Red
    Write-Host "Expected: $modulePath" -ForegroundColor Yellow
    Write-Host "Please ensure Syskit.Discovery.psd1 and M365Snapshot.psm1 are in the same directory." -ForegroundColor Yellow
    exit 1
}

try {
    Import-Module $modulePath -Force -ErrorAction Stop
}
catch {
    Write-Host "ERROR: Failed to import module from '$modulePath'." -ForegroundColor Red
    Write-Host "Details: $($_.Exception.Message)" -ForegroundColor Yellow
    Write-Host "Ensure prerequisite modules are installed: MSAL.PS, PnP.PowerShell, Microsoft.Graph.Authentication, Microsoft.Graph.Applications, and Microsoft.Graph.Identity.SignIns." -ForegroundColor Yellow
    exit 1
}

# Call the module function
$snapshotParams = @{
    TenantId = $TenantId
    ClientId = $ClientId
    MaxUsers = $MaxUsers
    MaxGroups = $MaxGroups
    MaxSites = $MaxSites
    IncludeAppRegistrations = $IncludeAppRegistrations
    MaxAppRegistrations = $MaxAppRegistrations
    IncludeSharedMailboxes = $IncludeSharedMailboxes
    MaxSharedMailboxes = $MaxSharedMailboxes
    IncludeSecurityDistributionGroups = $IncludeSecurityDistributionGroups
    MaxSecurityDistributionGroups = $MaxSecurityDistributionGroups
}

if ($ConfidentialSensitivityLabels.Count -gt 0) {
    $snapshotParams["ConfidentialSensitivityLabels"] = $ConfidentialSensitivityLabels
}

if ($LoadAllUsers) {
    $snapshotParams["LoadAllUsers"] = $true
}
if ($LoadAllGroups) {
    $snapshotParams["LoadAllGroups"] = $true
}
if ($LoadAllSites) {
    $snapshotParams["LoadAllSites"] = $true
}
if ($LoadAllAppRegistrations) {
    $snapshotParams["LoadAllAppRegistrations"] = $true
}

if ($GenerateHTMLReport) {
    $snapshotParams["GenerateHTMLReport"] = $true
    $snapshotParams["ReportPath"] = $ReportPath
    Get-M365Snapshot @snapshotParams
}
else {
    Get-M365Snapshot @snapshotParams
}