[CmdletBinding(SupportsShouldProcess)]
param(
    [Parameter(Mandatory)]
    [ValidatePattern('^[A-Za-z0-9][A-Za-z0-9.-]*$')]
    [string]$ServerAddress,

    [Parameter(Mandatory)]
    [ValidatePattern('^[A-Za-z0-9+/=]+$')]
    [string]$ServerPublicKey,

    [securestring]$AccessPassword,
    [string]$InstallRoot = "$env:ProgramFiles\RustDesk",
    [string]$StateRoot = "$env:ProgramData\SGU\RustDesk\Client",
    [string]$ClientVersion = '1.4.9',
    [uri]$DownloadUri = 'https://github.com/rustdesk/rustdesk/releases/download/1.4.9/rustdesk-1.4.9-x86_64.msi',
    [ValidatePattern('^[A-Fa-f0-9]{64}$')]
    [string]$ExpectedSha256 = 'C87D2F4CEF2A5ACD6003B6507DCFBF5D5168A256DB082CD90B54D35193224AAA'
)

$ErrorActionPreference = 'Stop'
$ProgressPreference = 'SilentlyContinue'
$downloadRoot = Join-Path $env:ProgramData 'SGU\RustDesk\Downloads'
$installerPath = Join-Path $downloadRoot "rustdesk-$ClientVersion-x86_64.msi"
$installerLogPath = Join-Path $downloadRoot "rustdesk-$ClientVersion-install.log"
$secretPath = Join-Path $StateRoot 'access.secret'
$devicePath = Join-Path $StateRoot 'device.json'

function Assert-Administrator {
    $identity = [Security.Principal.WindowsIdentity]::GetCurrent()
    $principal = [Security.Principal.WindowsPrincipal]::new($identity)
    if (-not $principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)) {
        throw 'Run this script from an elevated Windows PowerShell session.'
    }
}

function Initialize-DataProtection {
    if (-not ('SguRustDeskDataProtection' -as [type])) {
        Add-Type -TypeDefinition @'
using System;
using System.ComponentModel;
using System.Runtime.InteropServices;

public static class SguRustDeskDataProtection {
    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
    private struct DataBlob { public int cbData; public IntPtr pbData; }

    [DllImport("crypt32.dll", SetLastError = true, CharSet = CharSet.Unicode)]
    private static extern bool CryptProtectData(ref DataBlob input, string description,
        IntPtr optionalEntropy, IntPtr reserved, IntPtr prompt, int flags, out DataBlob output);

    [DllImport("crypt32.dll", SetLastError = true, CharSet = CharSet.Unicode)]
    private static extern bool CryptUnprotectData(ref DataBlob input, IntPtr description,
        IntPtr optionalEntropy, IntPtr reserved, IntPtr prompt, int flags, out DataBlob output);

    [DllImport("kernel32.dll", SetLastError = true)]
    private static extern IntPtr LocalFree(IntPtr memory);

    private const int CryptProtectLocalMachine = 0x4;

    private static DataBlob ToBlob(byte[] value) {
        var blob = new DataBlob { cbData = value.Length, pbData = IntPtr.Zero };
        if (value.Length > 0) {
            blob.pbData = Marshal.AllocHGlobal(value.Length);
            Marshal.Copy(value, 0, blob.pbData, value.Length);
        }
        return blob;
    }

    private static byte[] FromBlob(DataBlob blob) {
        var value = new byte[blob.cbData];
        if (blob.cbData > 0) Marshal.Copy(blob.pbData, value, 0, blob.cbData);
        return value;
    }

    public static byte[] Protect(byte[] value) {
        var input = ToBlob(value); var output = new DataBlob();
        try {
            if (!CryptProtectData(ref input, null, IntPtr.Zero, IntPtr.Zero, IntPtr.Zero,
                    CryptProtectLocalMachine, out output)) {
                throw new Win32Exception(Marshal.GetLastWin32Error());
            }
            return FromBlob(output);
        } finally {
            if (input.pbData != IntPtr.Zero) Marshal.FreeHGlobal(input.pbData);
            if (output.pbData != IntPtr.Zero) LocalFree(output.pbData);
        }
    }

    public static byte[] Unprotect(byte[] value) {
        var input = ToBlob(value); var output = new DataBlob();
        try {
            if (!CryptUnprotectData(ref input, IntPtr.Zero, IntPtr.Zero, IntPtr.Zero, IntPtr.Zero,
                    0, out output)) {
                throw new Win32Exception(Marshal.GetLastWin32Error());
            }
            return FromBlob(output);
        } finally {
            if (input.pbData != IntPtr.Zero) Marshal.FreeHGlobal(input.pbData);
            if (output.pbData != IntPtr.Zero) LocalFree(output.pbData);
        }
    }
}
'@ -ErrorAction Stop
    }
}

function Set-PrivateDirectoryAcl {
    param([Parameter(Mandatory)][string]$Path)

    New-Item -ItemType Directory -Path $Path -Force | Out-Null
    $acl = New-Object Security.AccessControl.DirectorySecurity
    $acl.SetAccessRuleProtection($true, $false)
    $inheritance = [Security.AccessControl.InheritanceFlags]'ContainerInherit,ObjectInherit'
    $allow = [Security.AccessControl.AccessControlType]::Allow
    foreach ($sid in @('S-1-5-18', 'S-1-5-32-544')) {
        $acl.AddAccessRule([Security.AccessControl.FileSystemAccessRule]::new(
            [Security.Principal.SecurityIdentifier]::new($sid),
            [Security.AccessControl.FileSystemRights]::FullControl,
            $inheritance,
            [Security.AccessControl.PropagationFlags]::None,
            $allow))
    }
    Set-Acl -LiteralPath $Path -AclObject $acl
}

function ConvertTo-PlainText {
    param([Parameter(Mandatory)][securestring]$SecureString)

    $pointer = [Runtime.InteropServices.Marshal]::SecureStringToBSTR($SecureString)
    try {
        return [Runtime.InteropServices.Marshal]::PtrToStringBSTR($pointer)
    }
    finally {
        [Runtime.InteropServices.Marshal]::ZeroFreeBSTR($pointer)
    }
}

function New-RandomAccessPassword {
    $characters = 'ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz23456789!@#$%*+-_'.ToCharArray()
    $bytes = New-Object byte[] 24
    $rng = [Security.Cryptography.RandomNumberGenerator]::Create()
    try {
        $rng.GetBytes($bytes)
    }
    finally {
        $rng.Dispose()
    }
    $value = -join ($bytes | ForEach-Object { $characters[$_ % $characters.Length] })
    return (ConvertTo-SecureString -String $value -AsPlainText -Force)
}

function Save-AccessPassword {
    param([Parameter(Mandatory)][securestring]$Password)

    $plainText = ConvertTo-PlainText -SecureString $Password
    try {
        $cipherText = [SguRustDeskDataProtection]::Protect(
            [Text.Encoding]::UTF8.GetBytes($plainText))
        [IO.File]::WriteAllBytes($secretPath, $cipherText)
    }
    finally {
        $plainText = $null
    }
}

function Get-SavedAccessPassword {
    if (-not (Test-Path -LiteralPath $secretPath -PathType Leaf)) {
        return $null
    }
    $plainText = [Text.Encoding]::UTF8.GetString(
        [SguRustDeskDataProtection]::Unprotect(
            [IO.File]::ReadAllBytes($secretPath)))
    try {
        return (ConvertTo-SecureString -String $plainText -AsPlainText -Force)
    }
    finally {
        $plainText = $null
    }
}

function Assert-FileHash {
    param(
        [Parameter(Mandatory)][string]$Path,
        [Parameter(Mandatory)][string]$ExpectedHash
    )

    $actualHash = (Get-FileHash -LiteralPath $Path -Algorithm SHA256).Hash
    if (-not $actualHash.Equals($ExpectedHash, [StringComparison]::OrdinalIgnoreCase)) {
        throw "SHA-256 verification failed for $Path."
    }
}

function Test-TcpConnection {
    param([Parameter(Mandatory)][string]$HostName, [Parameter(Mandatory)][int]$Port)

    $client = [Net.Sockets.TcpClient]::new()
    try {
        $connect = $client.BeginConnect($HostName, $Port, $null, $null)
        if (-not $connect.AsyncWaitHandle.WaitOne(5000)) {
            return $false
        }
        $client.EndConnect($connect)
        return $true
    }
    catch {
        return $false
    }
    finally {
        $client.Dispose()
    }
}

Assert-Administrator
Initialize-DataProtection
if (-not $PSCmdlet.ShouldProcess($env:COMPUTERNAME, 'Install and configure the managed RustDesk client')) {
    return
}

Set-PrivateDirectoryAcl -Path $StateRoot
New-Item -ItemType Directory -Path $downloadRoot -Force | Out-Null
$rustDeskExecutable = Join-Path $InstallRoot 'RustDesk.exe'
$installedVersion = if (Test-Path -LiteralPath $rustDeskExecutable -PathType Leaf) {
    [string](Get-Item -LiteralPath $rustDeskExecutable).VersionInfo.ProductVersion
}
else {
    ''
}
if (-not $installedVersion.StartsWith($ClientVersion, [StringComparison]::OrdinalIgnoreCase)) {
    if (-not (Test-Path -LiteralPath $installerPath -PathType Leaf) -or
        (Get-FileHash -LiteralPath $installerPath -Algorithm SHA256).Hash -ne $ExpectedSha256) {
        Invoke-WebRequest -Uri $DownloadUri -OutFile $installerPath -UseBasicParsing
    }
    Assert-FileHash -Path $installerPath -ExpectedHash $ExpectedSha256
    # The vendor's MSI is the supported path for managed, silent Windows
    # deployment.  Unlike the GUI-oriented EXE it does not require an
    # interactive desktop, which matters for startup/bootstrap execution.
    $msiArguments = "/i `"$installerPath`" /qn /norestart " +
        "INSTALLFOLDER=`"$InstallRoot`" CREATESTARTMENUSHORTCUTS=`"N`" " +
        "CREATEDESKTOPSHORTCUTS=`"N`" INSTALLPRINTER=`"N`" /l*v `"$installerLogPath`""
    $installer = Start-Process -FilePath (Join-Path $env:WINDIR 'System32\msiexec.exe') `
        -ArgumentList $msiArguments -Wait -PassThru
    if ($installer.ExitCode -notin @(0, 3010)) {
        throw "RustDesk MSI installation failed with exit code $($installer.ExitCode). See $installerLogPath."
    }
}

if (-not (Test-Path -LiteralPath $rustDeskExecutable -PathType Leaf)) {
    throw "RustDesk installation did not create $rustDeskExecutable."
}

$rustDeskService = Get-Service -Name 'RustDesk' -ErrorAction SilentlyContinue
if (-not $rustDeskService) {
    $serviceInstaller = Start-Process -FilePath $rustDeskExecutable -ArgumentList '--install-service' `
        -Wait -PassThru
    if ($serviceInstaller.ExitCode -ne 0) {
        throw "RustDesk service installation failed with exit code $($serviceInstaller.ExitCode)."
    }
    Start-Sleep -Seconds 2
    $rustDeskService = Get-Service -Name 'RustDesk' -ErrorAction SilentlyContinue
}
if (-not $rustDeskService) {
    throw 'RustDesk did not register its Windows service.'
}

Set-Service -Name $rustDeskService.Name -StartupType Automatic
if ($rustDeskService.Status -ne 'Stopped') {
    Stop-Service -Name $rustDeskService.Name -Force
    $rustDeskService.WaitForStatus('Stopped', (New-TimeSpan -Seconds 20))
}

$rendezvousAddress = "$ServerAddress`:21116"
$relayAddress = "$ServerAddress`:21117"
$configuration = @"
rendezvous_server = '$rendezvousAddress'
nat_type = 1
serial = 0

[options]
custom-rendezvous-server = '$rendezvousAddress'
relay-server = '$relayAddress'
key = '$ServerPublicKey'
"@
$configurationPaths = @(
    (Join-Path $env:ProgramData 'RustDesk\config\RustDesk2.toml'),
    (Join-Path $env:WINDIR 'ServiceProfiles\LocalService\AppData\Roaming\RustDesk\config\RustDesk2.toml'),
    (Join-Path $env:WINDIR 'System32\config\systemprofile\AppData\Roaming\RustDesk\config\RustDesk2.toml'),
    (Join-Path $env:SystemDrive 'Users\Default\AppData\Roaming\RustDesk\config\RustDesk2.toml')
)
foreach ($configurationPath in $configurationPaths) {
    New-Item -ItemType Directory -Path (Split-Path $configurationPath -Parent) -Force | Out-Null
    [IO.File]::WriteAllText($configurationPath, $configuration, [Text.UTF8Encoding]::new($false))
}

$existingPassword = Get-SavedAccessPassword
if ($AccessPassword) {
    $managedPassword = $AccessPassword
    Save-AccessPassword -Password $managedPassword
    $passwordWasGenerated = $false
}
elseif ($existingPassword) {
    $managedPassword = $existingPassword
    $passwordWasGenerated = $false
}
else {
    $managedPassword = New-RandomAccessPassword
    Save-AccessPassword -Password $managedPassword
    $passwordWasGenerated = $true
}

Start-Service -Name $rustDeskService.Name
$rustDeskService = Get-Service -Name $rustDeskService.Name
$rustDeskService.WaitForStatus('Running', (New-TimeSpan -Seconds 20))

$plainPassword = ConvertTo-PlainText -SecureString $managedPassword
try {
    # RustDesk on Windows only reliably treats its CLI output path as a command
    # invocation when stdout is consumed.  Without the pipeline it can attach
    # to the GUI instance and leave a non-interactive bootstrap waiting.
    $null = & $rustDeskExecutable --password $plainPassword | Out-String
    if ($LASTEXITCODE -ne 0) {
        throw "RustDesk could not set the managed access password (exit code $LASTEXITCODE)."
    }
}
finally {
    $plainPassword = $null
}

$rustDeskId = ((& $rustDeskExecutable --get-id | Out-String).Trim() -split "`r?`n" |
    Select-Object -Last 1).Trim()
if ($rustDeskId -notmatch '^\d+$') {
    throw "RustDesk returned an invalid device ID: $rustDeskId"
}
if (-not (Test-TcpConnection -HostName $ServerAddress -Port 21116)) {
    throw "The RustDesk rendezvous server $rendezvousAddress is not reachable from this client."
}

$device = [ordered]@{
    ComputerName = $env:COMPUTERNAME
    RustDeskId = $rustDeskId
    ServerAddress = $ServerAddress
    ServerPublicKeySha256 = ([Security.Cryptography.SHA256]::Create().ComputeHash(
        [Text.Encoding]::UTF8.GetBytes($ServerPublicKey)) | ForEach-Object ToString x2) -join ''
    ConfiguredAt = (Get-Date).ToString('o')
}
[IO.File]::WriteAllText($devicePath, ($device | ConvertTo-Json), [Text.UTF8Encoding]::new($false))

[pscustomobject]@{
    RustDeskId = $rustDeskId
    ServerAddress = $ServerAddress
    ServiceName = $rustDeskService.Name
    ServiceStatus = (Get-Service -Name $rustDeskService.Name).Status.ToString()
    RendezvousReachable = $true
    AccessPassword = $managedPassword
    AccessPasswordWasGenerated = $passwordWasGenerated
    DevicePath = $devicePath
}