#Requires -Version 5.1
Set-StrictMode -Version Latest

$script:RdpRepairDataRoot = Join-Path $env:ProgramData 'RdpRepair'
$script:RdpRepairRdpTcpKey = 'HKLM:\SYSTEM\CurrentControlSet\Control\Terminal Server\WinStations\RDP-Tcp'
$script:RdpRepairTsKey = 'HKLM:\SYSTEM\CurrentControlSet\Control\Terminal Server'

function Initialize-RdpRepairDataRoot {
    if (-not (Test-Path -LiteralPath $script:RdpRepairDataRoot)) {
        New-Item -ItemType Directory -Path $script:RdpRepairDataRoot -Force | Out-Null
    }
    return $script:RdpRepairDataRoot
}

function Write-RdpRepairLog {
    param(
        [string]$Message,
        [ValidateSet('INFO', 'WARN', 'ERROR', 'OK', 'STEP')]
        [string]$Level = 'INFO',
        [string]$LogPath
    )
    $ts = Get-Date -Format 'yyyy-MM-dd HH:mm:ss'
    $line = "[$ts] [$Level] $Message"
    if ($LogPath) {
        Add-Content -LiteralPath $LogPath -Value $line -Encoding UTF8
    }
    switch ($Level) {
        'ERROR' { Write-Host $line -ForegroundColor Red }
        'WARN'  { Write-Host $line -ForegroundColor Yellow }
        'OK'    { Write-Host $line -ForegroundColor Green }
        'STEP'  { Write-Host $line -ForegroundColor Cyan }
        default { Write-Host $line }
    }
}

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

function Assert-RdpRepairAdministrator {
    if (-not (Test-RdpRepairIsAdministrator)) {
        throw 'Run PowerShell as Administrator.'
    }
}

function Get-RdpRepairComputerDnsNames {
    $cs = Get-CimInstance Win32_ComputerSystem -ErrorAction Stop
    $computer = $env:COMPUTERNAME
    $fqdn = "$computer.$($cs.Domain)".ToLower()
    $names = @($computer, $fqdn)
    if ($fqdn -notmatch '\.') {
        $names += "$computer.local"
    }
    return $names | Select-Object -Unique
}

function ConvertTo-RdpRepairThumbprintHex {
    param([object]$HashValue)
    if ($null -eq $HashValue) { return $null }

    if ($HashValue -is [byte[]]) {
        if ($HashValue.Length -eq 0) { return $null }
        return -join ($HashValue | ForEach-Object { '{0:X2}' -f $_ })
    }

    if ($HashValue -is [System.Array] -and $HashValue.Count -gt 0) {
        $allNumeric = $true
        foreach ($b in $HashValue) {
            if ($b -isnot [byte] -and $b -isnot [int] -and $b -isnot [short]) {
                $allNumeric = $false
                break
            }
        }
        if ($allNumeric -and ($HashValue.Count -eq 20 -or $HashValue.Count -eq 16)) {
            return -join ($HashValue | ForEach-Object { '{0:X2}' -f [byte]$_ })
        }
    }

    $text = [string]$HashValue
    if ([string]::IsNullOrWhiteSpace($text)) { return $null }
    $text = $text.Trim()

    if ($text -match '^\d+( \d+)+$') {
        $bytes = $text -split '\s+' | ForEach-Object { [byte]$_ }
        return -join ($bytes | ForEach-Object { '{0:X2}' -f $_ })
    }

    $hex = ($text -replace '\s', '').ToUpper()
    if ($hex -match '^[0-9A-F]+$') { return $hex }
    return $null
}

function Get-RdpRepairSslCertificateHash {
    $p = Get-ItemProperty -Path $script:RdpRepairRdpTcpKey -Name SSLCertificateSHA1Hash -ErrorAction SilentlyContinue
    if ($null -eq $p -or -not $p.PSObject.Properties['SSLCertificateSHA1Hash']) { return $null }
    return ConvertTo-RdpRepairThumbprintHex -HashValue $p.SSLCertificateSHA1Hash
}

function Get-RdpRepairSecurityLayer {
    $p = Get-ItemProperty -Path $script:RdpRepairRdpTcpKey -ErrorAction SilentlyContinue
    if ($null -eq $p.SecurityLayer) { return '(default)' }
    return [int]$p.SecurityLayer
}

function Get-RdpRepairRemoteDesktopCertificates {
    $list = New-Object System.Collections.Generic.List[object]
    $store = New-Object System.Security.Cryptography.X509Certificates.X509Store('Remote Desktop', 'LocalMachine')
    $store.Open([System.Security.Cryptography.X509Certificates.OpenFlags]::ReadOnly)
    try {
        foreach ($c in $store.Certificates) {
            $list.Add([PSCustomObject]@{
                Subject    = $c.Subject
                Thumbprint = $c.Thumbprint
                NotAfter   = $c.NotAfter
                HasKey     = $c.HasPrivateKey
            })
        }
    } finally {
        $store.Close()
    }
    return [object[]]$list.ToArray()
}

function Get-RdpRepairCertificateArray {
    param([object]$Certificates)
    if ($null -eq $Certificates) { return @() }
    $flat = New-Object System.Collections.Generic.List[object]
    $pending = New-Object System.Collections.Generic.Stack[object]
    foreach ($entry in @($Certificates)) { $pending.Push($entry) }
    while ($pending.Count -gt 0) {
        $item = $pending.Pop()
        if ($null -eq $item) { continue }
        if ($item -is [object[]]) {
            foreach ($inner in @($item)) { $pending.Push($inner) }
        } elseif ($null -ne $item.PSObject.Properties['Thumbprint']) {
            [void]$flat.Add($item)
        } else {
            [void]$flat.Add($item)
        }
    }
    return [object[]]$flat.ToArray()
}

function Get-RdpRepairFirstCertificate {
    param([object]$Certificates)
    $arr = @(Get-RdpRepairCertificateArray -Certificates $Certificates)
    if ($arr.Count -lt 1) { return $null }
    return $arr[0]
}

function Clear-RdpRepairCertificateStores {
    param([string]$LogPath)
    foreach ($sn in @('Remote Desktop', 'My')) {
        $store = New-Object System.Security.Cryptography.X509Certificates.X509Store($sn, 'LocalMachine')
        $store.Open([System.Security.Cryptography.X509Certificates.OpenFlags]::ReadWrite)
        try {
            foreach ($c in @($store.Certificates)) {
                Write-RdpRepairLog "Удаление cert из ${sn}: $($c.Thumbprint) ($($c.Subject))" 'INFO' $LogPath
                $store.Remove($c)
            }
        } finally {
            $store.Close()
        }
    }
    Remove-ItemProperty -Path $script:RdpRepairRdpTcpKey -Name SSLCertificateSHA1Hash -ErrorAction SilentlyContinue
}

function Get-RdpRepairPrivateKeyPath {
    param([System.Security.Cryptography.X509Certificates.X509Certificate2]$Certificate)
    $rsa = [System.Security.Cryptography.X509Certificates.RSACertificateExtensions]::GetRSAPrivateKey($Certificate)
    if (-not $rsa) { return $null }
    $name = $rsa.Key.UniqueName
    foreach ($base in @(
        Join-Path $env:ProgramData 'Microsoft\Crypto\RSA\MachineKeys'
        Join-Path $env:ProgramData 'Microsoft\Crypto\Keys'
    )) {
        $path = Join-Path $base $name
        if (Test-Path -LiteralPath $path) { return $path }
    }
    return $null
}

function Set-RdpRepairPrivateKeyAcl {
    param(
        [System.Security.Cryptography.X509Certificates.X509Certificate2]$Certificate,
        [string]$LogPath
    )
    $keyPath = Get-RdpRepairPrivateKeyPath -Certificate $Certificate
    if (-not $keyPath) {
        Write-RdpRepairLog 'Не найден файл приватного ключа' 'WARN' $LogPath
        return $false
    }
    & icacls.exe $keyPath /grant 'NETWORK SERVICE:(F)' 'NT AUTHORITY\SYSTEM:(F)' 'BUILTIN\Administrators:(F)' 'NT AUTHORITY\LOCAL SERVICE:(R)' /c | Out-Null
    Write-RdpRepairLog "ACL обновлён: $keyPath" 'OK' $LogPath
    return $true
}

function Add-RdpRepairCertificateToRemoteDesktopStore {
    param([System.Security.Cryptography.X509Certificates.X509Certificate2]$Certificate)
    $store = New-Object System.Security.Cryptography.X509Certificates.X509Store('Remote Desktop', 'LocalMachine')
    $store.Open([System.Security.Cryptography.X509Certificates.OpenFlags]::ReadWrite)
    try {
        $store.Add($Certificate)
    } finally {
        $store.Close()
    }
}

function Set-RdpRepairCertificateBinding {
    param(
        [string]$Thumbprint,
        [string]$LogPath
    )
    $thumb = $Thumbprint.ToUpper()
    $regPath = 'HKLM\SYSTEM\CurrentControlSet\Control\Terminal Server\WinStations\RDP-Tcp'
    & reg.exe add $regPath /v SSLCertificateSHA1Hash /t REG_SZ /d $thumb /f | Out-Null
    try {
        $ts = Get-CimInstance -Namespace root\cimv2\terminalservices -ClassName Win32_TSGeneralSetting -Filter "TerminalName='RDP-Tcp'" -ErrorAction Stop
        $ts | Set-CimInstance -Property @{ SSLCertificateSHA1Hash = $thumb } -ErrorAction Stop
        Write-RdpRepairLog "WMI: SSLCertificateSHA1Hash = $thumb" 'INFO' $LogPath
    } catch {
        Write-RdpRepairLog "WMI Set-CimInstance: $($_.Exception.Message) (reg уже записан)" 'WARN' $LogPath
    }
    return (Get-RdpRepairSslCertificateHash)
}

function Restart-RdpRepairTermService {
    param([int]$WaitSeconds = 8)
    Stop-Service TermService -Force -ErrorAction Stop
    Start-Sleep -Seconds 2
    Start-Service TermService -ErrorAction Stop
    Start-Sleep -Seconds $WaitSeconds
}

function Test-RdpRepairCertificateBindingPersisted {
    param(
        [string]$ExpectedThumbprint,
        [string]$LogPath
    )
    $before = Get-RdpRepairSslCertificateHash
    Write-RdpRepairLog "Hash BEFORE TermService start: $(if ($before) { $before } else { '(empty)' })" 'INFO' $LogPath
    Restart-RdpRepairTermService
    $after = Get-RdpRepairSslCertificateHash
    Write-RdpRepairLog "Hash AFTER TermService start:  $(if ($after) { $after } else { '(empty)' })" 'INFO' $LogPath
    $expected = ConvertTo-RdpRepairThumbprintHex -HashValue $ExpectedThumbprint
    $ok = ($after -eq $expected)
    if ($ok) {
        Write-RdpRepairLog 'Привязка cert принята TermService' 'OK' $LogPath
    } else {
        Write-RdpRepairLog 'TermService сбросил или отклонил привязку cert' 'ERROR' $LogPath
    }
    return [PSCustomObject]@{
        Before    = $before
        After     = $after
        Persisted = $ok
    }
}

function New-RdpRepairSelfSignedCertificate {
    param(
        [string[]]$DnsNames,
        [string]$LogPath
    )
    $providers = @(
        @{ Name = 'Microsoft Enhanced RSA and AES Cryptographic Provider'; KeySpec = 'KeyExchange' }
        @{ Name = $null; KeySpec = $null }
    )
    $notAfter = (Get-Date).AddYears(2)
    $subject = "CN=$($DnsNames | Select-Object -First 1)"
    foreach ($prov in $providers) {
        try {
            $params = @{
                Subject           = $subject
                DnsName           = $DnsNames
                CertStoreLocation = 'Cert:\LocalMachine\My'
                KeyAlgorithm      = 'RSA'
                KeyLength         = 2048
                KeyUsage          = 'DigitalSignature', 'KeyEncipherment'
                TextExtension     = @('2.5.29.37={text}1.3.6.1.5.5.7.3.1')
                NotAfter          = $notAfter
            }
            if ($prov.Name) {
                $params['Provider'] = $prov.Name
                $params['KeySpec'] = $prov.KeySpec
            }
            $cert = New-SelfSignedCertificate @params
            $keyPath = Get-RdpRepairPrivateKeyPath -Certificate $cert
            $provLabel = if ($prov.Name) { $prov.Name } else { 'default' }
            Write-RdpRepairLog "Создан cert $($cert.Thumbprint) провайдер=$provLabel key=$keyPath" 'OK' $LogPath
            return $cert
        } catch {
            $provLabel = if ($prov.Name) { $prov.Name } else { 'default' }
            Write-RdpRepairLog "New-SelfSignedCertificate ($provLabel): $($_.Exception.Message)" 'WARN' $LogPath
        }
    }
    throw 'Не удалось создать самозаверяющий сертификат ни одним провайдером.'
}

function Invoke-RdpRepairNativeReset {
    param(
        [string]$LogPath,
        [switch]$Reboot
    )
    Write-RdpRepairLog 'Нативный сброс RDP cert (без ручного cert)' 'STEP' $LogPath
    Stop-Service TermService -Force -ErrorAction SilentlyContinue
    Set-ItemProperty -Path $script:RdpRepairTsKey -Name fDenyTSConnections -Value 1
    Clear-RdpRepairCertificateStores -LogPath $LogPath
    Set-ItemProperty -Path $script:RdpRepairTsKey -Name fDenyTSConnections -Value 0
    if ($Reboot) {
        Write-RdpRepairLog 'Перезагрузка через 60 с для нативного создания cert...' 'WARN' $LogPath
        shutdown.exe /r /t 60 /f /c 'RdpRepair: native certificate rebuild'
        return [PSCustomObject]@{ RebootScheduled = $true }
    }
    Start-Service TermService -ErrorAction SilentlyContinue
    Start-Sleep -Seconds 5
    $certs = @(Get-RdpRepairRemoteDesktopCertificates)
    return [PSCustomObject]@{
        RebootScheduled = $false
        CertCount       = $certs.Count
        Certificates    = ,[object[]]$certs
    }
}

function Get-RdpRepairDiagnosticSnapshot {
    $polPath = 'HKLM:\SOFTWARE\Policies\Microsoft\Windows NT\Terminal Services'
    $pol = if (Test-Path $polPath) { Get-ItemProperty $polPath -ErrorAction SilentlyContinue } else { $null }
    $fips = (Get-ItemProperty 'HKLM:\SYSTEM\CurrentControlSet\Control\Lsa\FipsAlgorithmPolicy' -Name Enabled -ErrorAction SilentlyContinue).Enabled
    $certs = @(Get-RdpRepairRemoteDesktopCertificates)
    $ev1057 = @(Get-WinEvent -FilterHashtable @{
        LogName   = 'System'
        Id        = 1057, 36870
        StartTime = (Get-Date).AddHours(-24)
    } -ErrorAction SilentlyContinue).Count
    $os = Get-CimInstance Win32_OperatingSystem
    return [PSCustomObject]@{
        ComputerName   = $env:COMPUTERNAME
        OsCaption      = $os.Caption
        OsBuild        = $os.BuildNumber
        FipsEnabled    = $fips
        SecurityLayer  = Get-RdpRepairSecurityLayer
        SslHash        = Get-RdpRepairSslCertificateHash
        TermService    = (Get-Service TermService).Status
        RdpAllowed     = -not (Get-ItemProperty $script:RdpRepairTsKey).fDenyTSConnections
        CertCount      = $certs.Count
        Certificates   = ,[object[]]$certs
        GpoLicenseSrv  = $pol.LicenseServers
        GpoLicMode     = $pol.LicensingMode
        Events1057_24h = $ev1057
    }
}

function Export-RdpRepairDiagnosticReport {
    param([string]$LogPath)
    $snap = Get-RdpRepairDiagnosticSnapshot
    $reportPath = Join-Path (Initialize-RdpRepairDataRoot) ("report-{0:yyyyMMdd-HHmmss}.txt" -f (Get-Date))
    $lines = @(
        "=== RDP/TLS Diagnostic Report ==="
        "Generated: $(Get-Date -Format o)"
        ""
    )
    $snap.PSObject.Properties | ForEach-Object {
        if ($_.Name -eq 'Certificates') {
            $lines += 'Certificates:'
            foreach ($c in @(Get-RdpRepairCertificateArray $_.Value)) {
                $lines += ('  - {0}; Thumb={1}; NotAfter={2}; Key={3}' -f $c.Subject, $c.Thumbprint, $c.NotAfter, $c.HasKey)
            }
        } else {
            $lines += ('{0}: {1}' -f $_.Name, $_.Value)
        }
    }
    $text = $lines -join "`r`n"
    [System.IO.File]::WriteAllText($reportPath, $text, [System.Text.UTF8Encoding]::new($false))
    Write-RdpRepairLog "Отчёт: $reportPath" 'OK' $LogPath
    return $reportPath
}

function Set-RdpRepairSecurityLayer {
    param(
        [int]$Layer,
        [string]$LogPath
    )
    Stop-Service TermService -Force -ErrorAction Stop
    Set-ItemProperty -Path $script:RdpRepairRdpTcpKey -Name SecurityLayer -Value $Layer
    if ($Layer -eq 0 -or $Layer -eq 1) {
        Remove-ItemProperty -Path $script:RdpRepairRdpTcpKey -Name SSLCertificateSHA1Hash -ErrorAction SilentlyContinue
    }
    Start-Service TermService -ErrorAction Stop
    Write-RdpRepairLog "SecurityLayer установлен в $Layer" 'OK' $LogPath
}

function Confirm-RdpRepairContinue {
    param([string]$Prompt)
    if ($env:RDP_REPAIR_NONINTERACTIVE -eq '1') { return $true }
    $r = Read-Host "$Prompt [Y/n]"
    return ($r -eq '' -or $r -match '^[yY]$')
}