# ============================================================
# CHANGELOG
# v1.1 - Added -AllUsers switch to deploy to all user desktops
#        instead of only the public desktop.
#        Added Get-DesktopPaths helper function.
# v1.0 - Initial version, deploys to public desktop only.
# ============================================================

param(
    [switch]$AllUsers
)

# Setup logging
$baseFolder = Join-Path $env:SystemDrive "Intune"
$logFile = Join-Path $baseFolder "DesktopShortcuts.log"

# Check if Sysnative exists (means script is running in 32-bit mode on x64 OS)
if (Test-Path "$($env:windir)\Sysnative") {
    $rdpSignPath = "$($env:windir)\Sysnative\rdpsign.exe"
} else {
    $rdpSignPath = "$($env:windir)\System32\rdpsign.exe"
}

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

function Write-Log {
    param($Message)
    $logMessage = "$(Get-Date -Format 'dd-MM-yyyy HH:mm'): $Message"
    Add-Content -Path $logFile -Value $logMessage
    Write-Host $logMessage
}

function Get-DesktopPaths {
    if ($AllUsers) {
        # Return desktop path for each user profile found
        Get-ChildItem "$env:SystemDrive\Users" -Directory | ForEach-Object {
            $path = Join-Path $_.FullName "Desktop"
            if (Test-Path $path) { $path }
        }
    } else {
        [Environment]::GetFolderPath("CommonDesktopDirectory")
    }
}

try {
    $desktopPaths = Get-DesktopPaths
    Write-Log "Running in mode: $(if ($AllUsers) { 'All Users' } else { 'Public Desktop' })"

    foreach ($desktopPath in $desktopPaths) {
        Write-Log "Processing desktop: $desktopPath"

        # Copy RDP files to desktop
        Get-ChildItem -Path ".\rdp-files\*.rdp" | ForEach-Object {
            try {
                Copy-Item -Path $_.FullName -Destination $desktopPath -Force
            } catch {
                Write-Log "Error copying file: $_"
                exit 1
            }
        }
    }

    Start-Sleep -Seconds 2 # Wait for files to be fully copied before signing
    Write-Log "Copied .rdp files to desktop(s)"

    # Certificate setup
    $certSubjectName = "RDPSign"
    $certSubject = "CN=$certSubjectName"

    Write-Log "Searching for existing certificate: $certSubjectName..."

    $existingCert = Get-ChildItem Cert:\LocalMachine\My |
        Where-Object { $_.Subject -eq $certSubject } |
        Select-Object -First 1

    if ($existingCert) {
        Write-Log "Found existing certificate with Thumbprint: $($existingCert.Thumbprint)"
        $thumbprint = $existingCert.Thumbprint
    } else {
        Write-Log "No existing certificate found. Creating new one..."

        $cert = New-SelfSignedCertificate `
            -Subject $certSubject `
            -CertStoreLocation "Cert:\LocalMachine\My" `
            -Type CodeSigningCert `
            -KeyExportPolicy NonExportable `
            -NotAfter (Get-Date).AddYears(25)

        $thumbprint = $cert.Thumbprint

        # Add to Trusted Root
        $rootStore = New-Object System.Security.Cryptography.X509Certificates.X509Store("Root", "LocalMachine")
        $rootStore.Open("ReadWrite")
        $rootStore.Add($cert)
        $rootStore.Close()

        # Add to Trusted Publishers
        $pubStore = New-Object System.Security.Cryptography.X509Certificates.X509Store("TrustedPublisher", "LocalMachine")
        $pubStore.Open("ReadWrite")
        $pubStore.Add($cert)
        $pubStore.Close()

        # Add thumbprint to RDP trusted certs in registry
        $gpoPath = "HKLM:\SOFTWARE\Policies\Microsoft\Windows NT\Terminal Services"
        $keyName = "TrustedCertThumbprints"

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

        $currentValue = (Get-ItemProperty -Path $gpoPath -Name $keyName -ErrorAction SilentlyContinue).$keyName

        if ($currentValue -notmatch [regex]::Escape($thumbprint)) {
            $newValue = if ([string]::IsNullOrWhiteSpace($currentValue)) { $thumbprint } else { "$currentValue,$thumbprint" }

            if ($null -eq $currentValue) {
                New-ItemProperty -Path $gpoPath -Name $keyName -Value $newValue -PropertyType String -Force | Out-Null
            } else {
                Set-ItemProperty -Path $gpoPath -Name $keyName -Value $newValue
            }
            Write-Log "Updated TrustedCertThumbprints registry value."
        } else {
            Write-Log "Thumbprint already exists in registry. Skipping update."
        }

        Write-Log "Certificate created and trusted successfully."
    }

    # Sign RDP files on all desktop paths
    $successCount = 0
    $failCount = 0

    foreach ($desktopPath in $desktopPaths) {
        $rdpFiles = Get-ChildItem -Path $desktopPath -Filter "*.rdp" -Recurse -ErrorAction SilentlyContinue

        if ($rdpFiles.Count -eq 0) {
            Write-Log "No .rdp files found on desktop: $desktopPath"
            continue
        }

        Write-Log "Found $($rdpFiles.Count) file(s) to sign on: $desktopPath"

        foreach ($rdpFile in $rdpFiles) {
            Write-Log "Signing: $($rdpFile.Name)"
            & $rdpSignPath /sha256 $thumbprint $rdpFile.FullName

            if ($LASTEXITCODE -eq 0) {
                Write-Log "  [OK] $($rdpFile.Name)"
                $successCount++
            } else {
                Write-Log "  [FAILED] $($rdpFile.Name) - rdpsign exit code: $LASTEXITCODE"
                $failCount++
            }
        }
    }

    # Set registry key to suppress RDP warning dialogs
    reg add "HKLM\Software\Policies\Microsoft\Windows NT\Terminal Services\Client" /v RedirectionWarningDialogVersion /t REG_DWORD /d 1 /f

    if ($failCount -eq 0) {
        Write-Log "All $successCount file(s) signed successfully."
    } else {
        Write-Log "$successCount file(s) signed successfully, $failCount file(s) failed."
        exit 1
    }
    exit 0
} catch {
    Write-Log "Error: $_"
    exit 1
}