#requires -RunAsAdministrator
<#
.SYNOPSIS
  Windows Server RDP 运维配置脚本

.DESCRIPTION
  功能：
  1. 修改主 RDP 监听端口
  2. 为主 RDP 端口放行 TCP/UDP 防火墙规则
  3. 可选创建额外代理端口（portproxy）转发到主 RDP 端口
  4. 支持按端口设置 IP 白名单
  5. 可选禁用系统内置的旧版远程桌面防火墙规则
  6. 可选重启远程桌面服务
  7. 输出最终状态摘要

.PARAMETER Help
  显示帮助信息

.PARAMETER MainRdpPort
  主 RDP 监听端口，默认 33893

.PARAMETER ExtraProxyPorts
  额外代理端口，例如 3390,3391

.PARAMETER PortAccessMap
  按端口配置允许访问的远程地址
  示例：
  @{
      33893 = @("1.2.3.4","5.6.7.0/24")
      3390  = @("8.8.8.8")
      3391  = "Any"
  }

.PARAMETER RestartTermService
  配置完成后重启远程桌面服务

.PARAMETER DisableLegacy3389Rules
  禁用系统内置旧版远程桌面入站防火墙规则

.PARAMETER ClearOldPortProxy
  创建新 portproxy 规则前先清空现有规则

.NOTES
  修改 RDP 设置前，请确保保留控制台/KVM 访问能力。
#>

[CmdletBinding()]
param(
    [switch]$Help,

    [ValidateRange(1,65535)]
    [int]$MainRdpPort = 33893,

    [int[]]$ExtraProxyPorts = @(),

    [hashtable]$PortAccessMap = @{},

    [switch]$RestartTermService,

    [switch]$DisableLegacy3389Rules,

    [switch]$ClearOldPortProxy
)

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

function Write-Info {
    param([string]$Message)
    Write-Host ("[{0}] [INFO ] {1}" -f (Get-Date -Format 'yyyy-MM-dd HH:mm:ss'), $Message) -ForegroundColor Cyan
}

function Write-Ok {
    param([string]$Message)
    Write-Host ("[{0}] [ OK  ] {1}" -f (Get-Date -Format 'yyyy-MM-dd HH:mm:ss'), $Message) -ForegroundColor Green
}

function Write-WarnMsg {
    param([string]$Message)
    Write-Host ("[{0}] [WARN ] {1}" -f (Get-Date -Format 'yyyy-MM-dd HH:mm:ss'), $Message) -ForegroundColor Yellow
}

function Write-ErrMsg {
    param([string]$Message)
    Write-Host ("[{0}] [ERROR] {1}" -f (Get-Date -Format 'yyyy-MM-dd HH:mm:ss'), $Message) -ForegroundColor Red
}

function Show-Usage {
    Write-Host ""
    Write-Host "Set-RdpOpsConfig.ps1 使用说明" -ForegroundColor White
    Write-Host ""
    Write-Host "参数：" -ForegroundColor White
    Write-Host "  -Help"
    Write-Host "  -MainRdpPort <port>             主 RDP 监听端口。默认：33893"
    Write-Host "  -ExtraProxyPorts <port[]>       额外代理端口。示例：3390,3391"
    Write-Host "  -PortAccessMap <Hashtable>      按端口设置白名单"
    Write-Host "  -RestartTermService             应用变更后重启远程桌面服务"
    Write-Host "  -DisableLegacy3389Rules         仅禁用系统内置远程桌面入站规则"
    Write-Host "  -ClearOldPortProxy              先清空所有已有的 portproxy 规则"
    Write-Host ""
    Write-Host "PortAccessMap 示例：" -ForegroundColor White
    Write-Host '  $PortMap = @{'
    Write-Host '      33893 = @("1.2.3.4","5.6.7.0/24")'
    Write-Host '      3390  = @("8.8.8.8")'
    Write-Host '      3391  = "Any"'
    Write-Host '  }'
    Write-Host ""
    Write-Host "本地执行示例：" -ForegroundColor White
    Write-Host '  powershell -ExecutionPolicy Bypass -File .\Set-RdpOpsConfig.ps1 -MainRdpPort 33893 -RestartTermService'
    Write-Host ""
    Write-Host "远程下载执行示例：" -ForegroundColor White
    Write-Host '  $script = irm "https://mirrors.dtops.co/software/Windows/Set-RdpOpsConfig.ps1"'
    Write-Host '  & ([scriptblock]::Create($script)) -MainRdpPort 33893 -RestartTermService'
    Write-Host ""
    Write-WarnMsg "断开当前会话前，请先在新的 RDP 窗口测试新端口。"
    Write-Host ""
}

if ($Help -or $args -contains '--help' -or $args -contains '-help') {
    Show-Usage
    return
}

function Test-IsAdmin {
    $id = [Security.Principal.WindowsIdentity]::GetCurrent()
    $p  = New-Object Security.Principal.WindowsPrincipal($id)
    return $p.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}

function Backup-RdpRegistry {
    $backupDir = Join-Path $env:SystemDrive "RdpOpsBackup"
    if (-not (Test-Path $backupDir)) {
        New-Item -Path $backupDir -ItemType Directory -Force | Out-Null
    }

    $stamp = Get-Date -Format "yyyyMMdd_HHmmss"
    $regFile = Join-Path $backupDir "RDP-Tcp_$stamp.reg"
    & reg.exe export "HKLM\SYSTEM\CurrentControlSet\Control\Terminal Server\WinStations\RDP-Tcp" $regFile /y | Out-Null
    Write-Ok "注册表备份已保存：$regFile"
}

function Set-MainRdpPortValue {
    $regPath = "HKLM:\SYSTEM\CurrentControlSet\Control\Terminal Server\WinStations\RDP-Tcp"
    $oldPort = (Get-ItemProperty -Path $regPath -Name PortNumber).PortNumber
    Write-Info "当前主 RDP 端口：$oldPort"
    Set-ItemProperty -Path $regPath -Name PortNumber -Value $MainRdpPort
    Write-Ok "主 RDP 端口已设置为：$MainRdpPort"
}

function Remove-FirewallRuleIfExists {
    param([string]$DisplayName)

    $rules = @(Get-NetFirewallRule -DisplayName $DisplayName -ErrorAction SilentlyContinue)
    if ($rules.Count -gt 0) {
        $rules | Remove-NetFirewallRule
        Write-Info "已删除旧防火墙规则：$DisplayName"
    }
}

function Normalize-RemoteAddressList {
    param([object]$Value)

    $result = New-Object System.Collections.Generic.List[string]

    if ($null -eq $Value) {
        return @()
    }

    foreach ($item in @($Value)) {
        if ($null -eq $item) { continue }

        $raw = ([string]$item).Trim()
        if ([string]::IsNullOrWhiteSpace($raw)) { continue }

        # 同时支持数组输入，以及逗号/分号分隔的字符串输入。
        foreach ($s in ($raw -split '\s*[,;]\s*')) {
            if ([string]::IsNullOrWhiteSpace($s)) { continue }

            if ($s -match '^(?i:any|\*|0\.0\.0\.0/0)$') {
                return @()
            }

            $result.Add($s)
        }
    }

    return @($result.ToArray())
}

function Get-AllowedAddressesForPort {
    param([int]$Port)

    if ($PortAccessMap.ContainsKey($Port)) {
        return @(Normalize-RemoteAddressList -Value $PortAccessMap[$Port])
    }

    $portKey = [string]$Port
    if ($PortAccessMap.ContainsKey($portKey)) {
        return @(Normalize-RemoteAddressList -Value $PortAccessMap[$portKey])
    }

    return @()
}

function Add-OrReplaceFirewallRule {
    param(
        [string]$DisplayName,
        [string]$Protocol,
        [int]$Port,
        [string[]]$RemoteAddress = @()
    )

    Remove-FirewallRuleIfExists -DisplayName $DisplayName

    $RemoteAddress = @($RemoteAddress | Where-Object {
        $_ -ne $null -and -not [string]::IsNullOrWhiteSpace(([string]$_))
    })

    if ($RemoteAddress.Count -gt 0) {
        # 直接传入 string[]，提升不同 Windows Server 版本下的兼容性。
        New-NetFirewallRule -DisplayName $DisplayName -Direction Inbound -Action Allow -Enabled True -Profile Any -Protocol $Protocol -LocalPort $Port -RemoteAddress $RemoteAddress -ErrorAction Stop | Out-Null

        Write-Ok "防火墙规则已启用：$DisplayName；白名单：$($RemoteAddress -join ', ')"
    }
    else {
        New-NetFirewallRule -DisplayName $DisplayName -Direction Inbound -Action Allow -Enabled True -Profile Any -Protocol $Protocol -LocalPort $Port -ErrorAction Stop | Out-Null

        Write-Ok "防火墙规则已启用：$DisplayName；白名单：Any"
    }

    Enable-NetFirewallRule -DisplayName $DisplayName -ErrorAction Stop | Out-Null
}

function Set-RdpFirewall {
    $mainAllowed = @(Get-AllowedAddressesForPort -Port $MainRdpPort)

    Add-OrReplaceFirewallRule -DisplayName "OPS-RDP-TCP-$MainRdpPort" -Protocol TCP -Port $MainRdpPort -RemoteAddress $mainAllowed
    Add-OrReplaceFirewallRule -DisplayName "OPS-RDP-UDP-$MainRdpPort" -Protocol UDP -Port $MainRdpPort -RemoteAddress $mainAllowed

    foreach ($proxyPort in @($ExtraProxyPorts)) {
        $proxyAllowed = @(Get-AllowedAddressesForPort -Port $proxyPort)
        Add-OrReplaceFirewallRule -DisplayName "OPS-RDP-PROXY-TCP-$proxyPort" -Protocol TCP -Port $proxyPort -RemoteAddress $proxyAllowed
    }
}

function Ensure-IpHelperRunning {
    $svc = Get-Service -Name iphlpsvc -ErrorAction Stop
    if ($svc.Status -ne 'Running') {
        Start-Service -Name iphlpsvc
        Write-Ok "已启动 IP Helper 服务（iphlpsvc）"
    }
    else {
        Write-Info "IP Helper 服务已在运行"
    }
}

function Clear-PortProxyRulesInternal {
    Write-Info "正在清空所有现有 portproxy 规则..."
    & netsh interface portproxy reset | Out-Null
    Write-Ok "已清空所有现有 portproxy 规则"
}

function Remove-PortProxyIfExists {
    param([int]$ListenPort)

    $output = & netsh interface portproxy show v4tov4
    if ($output -match "(?m)^\s*0\.0\.0\.0\s+$ListenPort\s+") {
        & netsh interface portproxy delete v4tov4 listenaddress=0.0.0.0 listenport=$ListenPort | Out-Null
        Write-Info "已删除现有 portproxy：0.0.0.0:$ListenPort"
    }
}

function Add-PortProxyRule {
    param(
        [int]$ListenPort,
        [int]$ConnectPort
    )

    Remove-PortProxyIfExists -ListenPort $ListenPort
    & netsh interface portproxy add v4tov4 listenaddress=0.0.0.0 listenport=$ListenPort connectaddress=127.0.0.1 connectport=$ConnectPort | Out-Null
    Write-Ok "已添加端口代理：0.0.0.0:$ListenPort -> 127.0.0.1:$ConnectPort"
}

function Configure-PortProxy {
    if (@($ExtraProxyPorts).Count -eq 0) {
        Write-Info "未配置额外代理端口"
        return
    }

    Ensure-IpHelperRunning

    if ($ClearOldPortProxy) {
        Clear-PortProxyRulesInternal
    }

    foreach ($proxyPort in @($ExtraProxyPorts)) {
        if ($proxyPort -eq $MainRdpPort) {
            Write-WarnMsg "已跳过代理端口 $proxyPort，因为它与主 RDP 端口相同。"
            continue
        }
        Add-PortProxyRule -ListenPort $proxyPort -ConnectPort $MainRdpPort
    }
}

function Disable-Legacy3389FirewallRules {
    $legacyRules = @(Get-NetFirewallRule -ErrorAction SilentlyContinue | Where-Object {
        $_.Enabled -eq 'True' -and
        $_.Direction -eq 'Inbound' -and
        $_.DisplayName -notlike 'OPS-RDP-*' -and
        (
            $_.DisplayName -like 'Remote Desktop*' -or
            $_.Group -like '*Remote Desktop*'
        )
    })

    if ($legacyRules.Count -gt 0) {
        $legacyRules | Disable-NetFirewallRule | Out-Null
        Write-Ok "已禁用系统内置旧版远程桌面入站防火墙规则"
    }
    else {
        Write-Info "未找到需要禁用的系统内置旧版远程桌面入站规则"
    }
}

function Restart-RdpServiceInternal {
    if (-not $RestartTermService) {
        Write-WarnMsg "未指定 RestartTermService。新端口可能需要重启服务或系统后才会完全生效。"
        return
    }

    try {
        Restart-Service -Name TermService -Force -ErrorAction Stop
        Write-Ok "已重启 TermService"
    }
    catch {
        Write-WarnMsg "重启 TermService 失败：$($_.Exception.Message)"
        Write-WarnMsg "在活跃的 RDP 会话中这是常见现象，必要时请重启系统。"
    }
}

function Test-RdpPortListening {
    param([int]$Port)

    try {
        $listeners = @(Get-NetTCPConnection -State Listen -LocalPort $Port -ErrorAction Stop)
        return ($listeners.Count -gt 0)
    }
    catch {
        $netstat = & netstat -ano
        return ($netstat -match "(?m)^\s*TCP\s+\S+:$Port\s+\S+\s+LISTENING\s+\d+\s*$")
    }
}

function Show-Summary {
    $regPath = "HKLM:\SYSTEM\CurrentControlSet\Control\Terminal Server\WinStations\RDP-Tcp"
    $currentPort = (Get-ItemProperty -Path $regPath -Name PortNumber).PortNumber

    Write-Host ""
    Write-Host "================ 最终结果 ================" -ForegroundColor White
    Write-Host ("主 RDP 端口：{0}" -f $currentPort) -ForegroundColor Green

    Write-Host ""
    Write-Host "端口访问映射：" -ForegroundColor White
    if (@($PortAccessMap.Keys).Count -eq 0) {
        Write-Host "  未指定；未配置端口默认 = Any" -ForegroundColor Yellow
    }
    else {
        foreach ($key in ($PortAccessMap.Keys | Sort-Object {[int]$_})) {
            $value = @(Normalize-RemoteAddressList -Value $PortAccessMap[$key])
            $show = if ($value.Count -gt 0) { $value -join ', ' } else { 'Any' }
            Write-Host ("  {0} -> {1}" -f $key, $show) -ForegroundColor Gray
        }
    }

    Write-Host ""
    Write-Host "防火墙规则：" -ForegroundColor White
    @(Get-NetFirewallRule -DisplayName "OPS-RDP*" -ErrorAction SilentlyContinue) |
        Sort-Object DisplayName |
        ForEach-Object {
            $pf = $_ | Get-NetFirewallPortFilter
            $af = $_ | Get-NetFirewallAddressFilter
            [PSCustomObject]@{
                Name          = $_.DisplayName
                Enabled       = $_.Enabled
                Action        = $_.Action
                Protocol      = $pf.Protocol
                LocalPort     = $pf.LocalPort
                RemoteAddress = $af.RemoteAddress
            }
        } | Format-Table -AutoSize

    Write-Host ""
    Write-Host "PortProxy 规则：" -ForegroundColor White
    & netsh interface portproxy show v4tov4

    Write-Host ""
    Write-Host "监听状态检查：" -ForegroundColor White
    $mainListening = Test-RdpPortListening -Port $currentPort
    Write-Host ("  {0} : {1}" -f $currentPort, $(if ($mainListening) { '正在监听' } else { '尚未监听' })) -ForegroundColor Yellow

    Write-Host ""
    Write-Host "连接测试示例：" -ForegroundColor White
    Write-Host ("  mstsc /v:<SERVER-IP>:{0}" -f $currentPort) -ForegroundColor Yellow
    foreach ($proxyPort in @($ExtraProxyPorts)) {
        Write-Host ("  mstsc /v:<SERVER-IP>:{0}" -f $proxyPort) -ForegroundColor Yellow
    }

    Write-Host ""
    Write-WarnMsg "断开当前会话前，请先在新的 RDP 窗口测试新端口。"
}

try {
    if (-not (Test-IsAdmin)) {
        throw "请以管理员身份运行 PowerShell。"
    }

    Write-Info "开始执行 Windows Server RDP 运维配置"
    Backup-RdpRegistry
    Set-MainRdpPortValue
    Set-RdpFirewall
    Configure-PortProxy

    if ($DisableLegacy3389Rules) {
        Disable-Legacy3389FirewallRules
    }

    Restart-RdpServiceInternal
    Show-Summary
    Write-Ok "执行完成"
}
catch {
    Write-ErrMsg $_.Exception.Message
    throw
}