# ============================================================
# SIGNONLY + DETECT - Sandbox test script
# ============================================================
# CHANGELOG
# v1.0 - Initial
# ============================================================

param(
    [switch]$AllUsers
)
Start-Sleep -Seconds 10
# Setup logging
$baseFolder = Join-Path $env:SystemDrive "Intune"
$logFile = Join-Path $baseFolder "RDPSign.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) {
        Get-ChildItem "$env:SystemDrive\Users" -Directory -Exclude "Public","Default" | ForEach-Object {
            Get-ChildItem $_.FullName -Directory -Filter "Desktop" -Depth 2 -ErrorAction SilentlyContinue |
                Select-Object -ExpandProperty FullName
        }
    } else {
        [Environment]::GetFolderPath("CommonDesktopDirectory")
    }
}

# ============================================================
# SIGNONLY
# ============================================================
Write-Log "=== START SIGNONLY ==="
Write-Log "Running in mode: $(if ($AllUsers) { 'All Users' } else { 'Public Desktop' })"

try {
    $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."
    }

    $successCount = 0
    $failCount = 0

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

        if ($rdpFiles.Count -eq 0) {
            Write-Log "Skip: No .rdp files on $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."
    }

} catch {
    Write-Log "SIGNONLY error: $_"
}

# ============================================================
# DETECT
# ============================================================
Write-Log "=== START DETECT ==="

try {
    $certSubject = "CN=RDPSign"

    Write-Log "Checking certificate: $certSubject..."
    $certInMy   = Get-ChildItem Cert:\LocalMachine\My | Where-Object { $_.Subject -eq $certSubject }
    $certInRoot = Get-ChildItem Cert:\LocalMachine\Root | Where-Object { $_.Subject -eq $certSubject }
    $certInPub  = Get-ChildItem Cert:\LocalMachine\TrustedPublisher | Where-Object { $_.Subject -eq $certSubject }

    $certOk = $true
    if (-not $certInMy) {
        Write-Log "FAIL: Certificate not found in LocalMachine\My"
        $certOk = $false
    }
    if (-not $certInRoot) {
        Write-Log "FAIL: Certificate not found in LocalMachine\Root"
        $certOk = $false
    }
    if (-not $certInPub) {
        Write-Log "FAIL: Certificate not found in LocalMachine\TrustedPublisher"
        $certOk = $false
    }
    if (-not $certOk) {
        Write-Log "DETECT [FAIL]: Certificate missing"
    } else {
        Write-Log "OK: Certificate found in all stores"
    }

    # Check registry thumbprint
    Write-Log "Checking registry thumbprint..."
    $gpoPath = "HKLM:\SOFTWARE\Policies\Microsoft\Windows NT\Terminal Services"
    $thumbprintValue = (Get-ItemProperty -Path $gpoPath -Name "TrustedCertThumbprints" -ErrorAction SilentlyContinue).TrustedCertThumbprints
    if ([string]::IsNullOrWhiteSpace($thumbprintValue)) {
        Write-Log "FAIL: TrustedCertThumbprints not found in registry"
        $certOk = $false
    } else {
        Write-Log "OK: Thumbprint found in registry"
    }

    # Check that all RDP files on desktop(s) are signed
    Write-Log "Checking RDP file signatures..."
    $desktopsChecked = 0
    $filesChecked = 0
    $allSigned = $true

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

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

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

        foreach ($file in $rdpFiles) {
            $filesChecked++
            $content = Get-Content $file.FullName -ErrorAction SilentlyContinue
            $hasSig = $content | Where-Object { $_ -match "^signature:s:" }
            if (-not $hasSig) {
                Write-Log "FAIL: $($file.Name) on $desktopPath is not signed"
                $allSigned = $false
            }
        }
        $desktopsChecked++
    }

    if ($allSigned -and $certOk) {
        Write-Log "DETECT [OK]: $filesChecked file(s) across $desktopsChecked desktop(s) - all signed and configured"
        exit 0
    } else {
        Write-Log "DETECT [FAIL]: Check log for details"
        exit 1
    }
} catch {
    Write-Log "Detect error: $_"
    exit 1
}