RemoteDesk/packaging/windows-host/RemoteDeskHost.psm1

Set-StrictMode -Version Latest
$ErrorActionPreference = 'Stop'

$script:TerminalServerKey = 'HKLM:\SYSTEM\CurrentControlSet\Control\Terminal Server'
$script:RdpTcpKey = 'HKLM:\SYSTEM\CurrentControlSet\Control\Terminal Server\WinStations\RDP-Tcp'
$script:CurrentVersionKey = 'HKLM:\SOFTWARE\Microsoft\Windows NT\CurrentVersion'
$script:FirewallRuleNames = @(
    'RemoteDesktop-UserMode-In-TCP',
    'RemoteDesktop-UserMode-In-UDP'
)

function Get-RemoteDeskHostStatePath {
    [CmdletBinding()]
    param()

    $programData = [Environment]::GetFolderPath([Environment+SpecialFolder]::CommonApplicationData)
    Join-Path $programData 'RemoteDesk\Host\original-state.json'
}

function Test-RemoteDeskAdministrator {
    [CmdletBinding()]
    param()

    $identity = [Security.Principal.WindowsIdentity]::GetCurrent()
    $principal = [Security.Principal.WindowsPrincipal]::new($identity)
    $principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}

function Get-RegistryValueSnapshot {
    param(
        [Parameter(Mandatory)][string]$Path,
        [Parameter(Mandatory)][string]$Name
    )

    $key = Get-Item -LiteralPath $Path -ErrorAction Stop
    $exists = $null -ne $key.GetValue($Name, $null, [Microsoft.Win32.RegistryValueOptions]::DoNotExpandEnvironmentNames)
    [ordered]@{
        path = $Path
        name = $Name
        exists = $exists
        value = if ($exists) { [int]$key.GetValue($Name) } else { $null }
    }
}

function Get-RemoteDesktopFirewallRules {
    $rules = foreach ($name in $script:FirewallRuleNames) {
        $rule = Get-NetFirewallRule -Name $name -ErrorAction SilentlyContinue
        if ($null -ne $rule) {
            [ordered]@{
                name = $rule.Name
                enabled = [string]$rule.Enabled
            }
        }
    }
    @($rules)
}

function Get-RemoteDeskHostStatus {
    [CmdletBinding()]
    param()

    if ($env:OS -ne 'Windows_NT') {
        throw 'RemoteDesk Windows Host can only run on Windows.'
    }

    $currentVersion = Get-ItemProperty -LiteralPath $script:CurrentVersionKey
    $editionId = [string]$currentVersion.EditionID
    $productName = [string]$currentVersion.ProductName
    $unsupportedEdition = $editionId -match '^(Core|Starter|Cloud)'
    $rdpEnabled = (Get-ItemPropertyValue -LiteralPath $script:TerminalServerKey -Name fDenyTSConnections) -eq 0
    $nlaEnabled = (Get-ItemPropertyValue -LiteralPath $script:RdpTcpKey -Name UserAuthentication) -eq 1
    $service = Get-CimInstance Win32_Service -Filter "Name='TermService'"
    $firewallRules = @(Get-RemoteDesktopFirewallRules)
    $firewallReady = $firewallRules.Count -gt 0 -and @($firewallRules | Where-Object enabled -ne 'True').Count -eq 0
    $statePath = Get-RemoteDeskHostStatePath

    [pscustomobject][ordered]@{
        supported = -not $unsupportedEdition
        productName = $productName
        editionId = $editionId
        administrator = Test-RemoteDeskAdministrator
        rdpEnabled = $rdpEnabled
        nlaEnabled = $nlaEnabled
        termServiceState = [string]$service.State
        termServiceStartMode = [string]$service.StartMode
        firewallReady = $firewallReady
        firewallRules = $firewallRules
        backupPresent = Test-Path -LiteralPath $statePath -PathType Leaf
        ready = (-not $unsupportedEdition) -and $rdpEnabled -and $nlaEnabled -and
            $service.State -eq 'Running' -and $firewallReady
    }
}

function Save-RemoteDeskOriginalState {
    [CmdletBinding()]
    param()

    $statePath = Get-RemoteDeskHostStatePath
    if (Test-Path -LiteralPath $statePath -PathType Leaf) {
        return $statePath
    }

    $service = Get-CimInstance Win32_Service -Filter "Name='TermService'"
    $state = [ordered]@{
        schema = 1
        capturedAtUtc = [DateTime]::UtcNow.ToString('o')
        rdp = Get-RegistryValueSnapshot -Path $script:TerminalServerKey -Name 'fDenyTSConnections'
        nla = Get-RegistryValueSnapshot -Path $script:RdpTcpKey -Name 'UserAuthentication'
        termService = [ordered]@{
            startMode = [string]$service.StartMode
            wasRunning = $service.State -eq 'Running'
        }
        firewallRules = Get-RemoteDesktopFirewallRules
    }

    $stateDirectory = Split-Path -Parent $statePath
    New-Item -ItemType Directory -Path $stateDirectory -Force | Out-Null
    $temporaryPath = "$statePath.$PID.tmp"
    try {
        $state | ConvertTo-Json -Depth 6 | Set-Content -LiteralPath $temporaryPath -Encoding utf8
        Move-Item -LiteralPath $temporaryPath -Destination $statePath -Force
    } finally {
        Remove-Item -LiteralPath $temporaryPath -Force -ErrorAction SilentlyContinue
    }
    $statePath
}

function Enable-RemoteDeskHost {
    [CmdletBinding(SupportsShouldProcess, ConfirmImpact = 'High')]
    param()

    if (-not (Test-RemoteDeskAdministrator)) {
        throw 'Administrator privileges are required to configure Remote Desktop.'
    }

    $status = Get-RemoteDeskHostStatus
    if (-not $status.supported) {
        throw "Windows edition '$($status.editionId)' does not support hosting Microsoft Remote Desktop."
    }
    if (-not $PSCmdlet.ShouldProcess($status.productName, 'Enable Remote Desktop, NLA, firewall rules, and TermService')) {
        return $status
    }

    [void](Save-RemoteDeskOriginalState)
    Set-ItemProperty -LiteralPath $script:TerminalServerKey -Name fDenyTSConnections -Type DWord -Value 0
    Set-ItemProperty -LiteralPath $script:RdpTcpKey -Name UserAuthentication -Type DWord -Value 1

    $service = Get-CimInstance Win32_Service -Filter "Name='TermService'"
    if ($service.StartMode -eq 'Disabled') {
        Set-Service -Name TermService -StartupType Manual
    }
    if ((Get-Service -Name TermService).Status -ne 'Running') {
        Start-Service -Name TermService
    }

    foreach ($name in $script:FirewallRuleNames) {
        if ($null -ne (Get-NetFirewallRule -Name $name -ErrorAction SilentlyContinue)) {
            Enable-NetFirewallRule -Name $name
        }
    }

    Get-RemoteDeskHostStatus
}

function Set-RegistryValueFromSnapshot {
    param([Parameter(Mandatory)]$Snapshot)

    if ([bool]$Snapshot.exists) {
        Set-ItemProperty -LiteralPath ([string]$Snapshot.path) -Name ([string]$Snapshot.name) -Type DWord -Value ([int]$Snapshot.value)
    } else {
        Remove-ItemProperty -LiteralPath ([string]$Snapshot.path) -Name ([string]$Snapshot.name) -ErrorAction SilentlyContinue
    }
}

function Restore-RemoteDeskHost {
    [CmdletBinding(SupportsShouldProcess, ConfirmImpact = 'High')]
    param()

    if (-not (Test-RemoteDeskAdministrator)) {
        throw 'Administrator privileges are required to restore Remote Desktop settings.'
    }

    $statePath = Get-RemoteDeskHostStatePath
    if (-not (Test-Path -LiteralPath $statePath -PathType Leaf)) {
        return Get-RemoteDeskHostStatus
    }
    $state = Get-Content -LiteralPath $statePath -Raw | ConvertFrom-Json
    if ([int]$state.schema -ne 1) {
        throw "Unsupported RemoteDesk Host state schema: $($state.schema)"
    }
    if (-not $PSCmdlet.ShouldProcess($statePath, 'Restore the pre-RemoteDesk RDP configuration')) {
        return Get-RemoteDeskHostStatus
    }

    Set-RegistryValueFromSnapshot -Snapshot $state.rdp
    Set-RegistryValueFromSnapshot -Snapshot $state.nla
    foreach ($savedRule in @($state.firewallRules)) {
        $rule = Get-NetFirewallRule -Name ([string]$savedRule.name) -ErrorAction SilentlyContinue
        if ($null -ne $rule) {
            if ([string]$savedRule.enabled -eq 'True') {
                Enable-NetFirewallRule -Name $rule.Name
            } else {
                Disable-NetFirewallRule -Name $rule.Name
            }
        }
    }

    switch ([string]$state.termService.startMode) {
        'Auto' { Set-Service -Name TermService -StartupType Automatic }
        'Manual' { Set-Service -Name TermService -StartupType Manual }
        'Disabled' { Set-Service -Name TermService -StartupType Disabled }
        default { throw "Unsupported saved TermService start mode: $($state.termService.startMode)" }
    }
    if ([bool]$state.termService.wasRunning) {
        if ((Get-Service -Name TermService).Status -ne 'Running') {
            Start-Service -Name TermService
        }
    } elseif ((Get-Service -Name TermService).Status -eq 'Running') {
        Stop-Service -Name TermService -Force
    }

    Remove-Item -LiteralPath $statePath -Force
    Get-RemoteDeskHostStatus
}

Export-ModuleMember -Function @(
    'Enable-RemoteDeskHost',
    'Get-RemoteDeskHostStatePath',
    'Get-RemoteDeskHostStatus',
    'Restore-RemoteDeskHost',
    'Test-RemoteDeskAdministrator'
)