<#
.SYNOPSIS
    Install RustDesk if needed, then point it at your own ID/relay server.

.DESCRIPTION
    Run as Administrator (directly, or through GPO / Intune / your RMM tool).
    It is safe to re-run. Without -Installer, a machine that already has RustDesk is
    only re-pointed at your server; one without RustDesk gets the latest official release.

    Steps:
      1. Install RustDesk silently: the -Installer you give (your own build, fresh install
         or upgrade over any existing RustDesk), or, only if RustDesk is missing, the
         latest official x86_64 .exe from github.com/rustdesk/rustdesk.
      2. Apply your server settings with `rustdesk.exe --config <config string>`.
      3. Optionally set a permanent password for unattended access.
      4. Optionally enroll the device in the CortenDesk console (`rustdesk.exe --assign`):
         owner, device group, strategy, address book. Needs an enrollment token, which
         is passed in when the script runs and is never stored in a kit.
      5. Print the RustDesk ID so you can add the machine to your address book.

    Without a token the device still registers with the console on its own; a user who
    signs in to the console from the RustDesk client then owns it (with "Approve when
    the device's user signs in" on).

.PARAMETER IdServer
    Your hbbs host, e.g. id.example.com (add :port if not 21116).
.PARAMETER Key
    Contents of id_ed25519.pub on your server.
.PARAMETER RelayServer
    Optional. Only needed if hbbr runs on a different host than hbbs and hbbs was not
    started with -r. Ignored by RustDesk 1.2.x clients.
.PARAMETER ApiServer
    Optional. URL of an API/web console server, if you run one.
.PARAMETER ConfigString
    Use this instead of the four fields above: the string from helpdesk/tools/make-config.sh
    or from "Export Server Config" in the RustDesk network settings.
.PARAMETER Installer
    Optional. URL or local path of a RustDesk .exe or .msi to install, e.g. your own
    build. When given, it is installed even if RustDesk is already there (upgrade).
.PARAMETER Password
    Optional. Permanent password for unattended access.
.PARAMETER RandomPassword
    Generate a random 16-character permanent password and print it.
.PARAMETER EnrollToken
    Optional. CortenDesk API token (Settings > Security > API tokens) with Manage on
    devices, plus Manage on strategies / address books if you assign those. The
    CORTENDESK_ENROLL_TOKEN environment variable works too, for RMM tools that pass
    secrets that way.
.PARAMETER PromptEnrollToken
    Ask for the token (input hidden) and, if none of the fields below are given, for the
    owner and device group. Press Enter at the token prompt to skip enrollment.
.PARAMETER EnrollUser
    Console user who owns the device.
.PARAMETER DeviceGroup
    Device group to put it in.
.PARAMETER Strategy
    Strategy (device policy) to pin to it.
.PARAMETER AddressBook
    Shared address book to add it to.
.PARAMETER DeviceName
    Name to show for the device in the console.
.PARAMETER Note
    Note to store on the device.

.EXAMPLE
    .\Install-Helpdesk.ps1 -IdServer id.example.com -Key 'AbCd...='
.EXAMPLE
    .\Install-Helpdesk.ps1 -ConfigString 'Qfi0zdC1...' -RandomPassword
.EXAMPLE
    .\Install-Helpdesk.ps1 -ConfigString 'Qfi0zdC1...' -RandomPassword -PromptEnrollToken
.EXAMPLE
    # RMM: token from the tool's secret store, owner and group from its variables
    $env:CORTENDESK_ENROLL_TOKEN = '<token>'
    .\Install-Helpdesk.ps1 -ConfigString 'Qfi0zdC1...' -EnrollUser alice -DeviceGroup Reception
#>
[CmdletBinding(DefaultParameterSetName = 'Fields')]
param(
    [Parameter(Mandatory = $true, ParameterSetName = 'Fields')][string]$IdServer,
    [Parameter(Mandatory = $true, ParameterSetName = 'Fields')][string]$Key,
    [Parameter(ParameterSetName = 'Fields')][string]$RelayServer = '',
    [Parameter(ParameterSetName = 'Fields')][string]$ApiServer = '',
    [Parameter(Mandatory = $true, ParameterSetName = 'String')][string]$ConfigString,
    [string]$Installer = '',
    [string]$Password = '',
    [switch]$RandomPassword,
    [string]$EnrollToken = '',
    [switch]$PromptEnrollToken,
    [string]$EnrollUser = '',
    [string]$DeviceGroup = '',
    [string]$Strategy = '',
    [string]$AddressBook = '',
    [string]$DeviceName = '',
    [string]$Note = ''
)

$ErrorActionPreference = 'Stop'
$ProgressPreference = 'SilentlyContinue'   # Invoke-WebRequest is very slow with the progress bar on PS 5.1
[Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12

$ServiceName = 'RustDesk'
# RMM agents often run 32-bit PowerShell, where $env:ProgramFiles points at "Program Files (x86)".
$ProgramFiles64 = if ($env:ProgramW6432) { $env:ProgramW6432 } else { $env:ProgramFiles }
$RustDeskExe = Join-Path $ProgramFiles64 'RustDesk\rustdesk.exe'

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

# Same format as helpdesk/tools/make-config.sh: JSON -> base64url (no padding) -> reversed.
function ConvertTo-ConfigString([string]$IdHost, [string]$PubKey, [string]$Relay, [string]$Api) {
    $json = [ordered]@{ host = $IdHost; relay = $Relay; api = $Api; key = $PubKey } | ConvertTo-Json -Compress
    $b64 = [Convert]::ToBase64String([Text.Encoding]::UTF8.GetBytes($json)).TrimEnd('=').Replace('+', '-').Replace('/', '_')
    $chars = $b64.ToCharArray()
    [Array]::Reverse($chars)
    return -join $chars
}

# rustdesk.exe is a GUI-subsystem program: piping its output makes PowerShell wait for it.
function Invoke-RustDesk([string[]]$Arguments) {
    return (& $RustDeskExe @Arguments | Out-String).Trim()
}

function Get-LatestInstallerUrl {
    $release = Invoke-RestMethod -UseBasicParsing -Uri 'https://api.github.com/repos/rustdesk/rustdesk/releases/latest' `
        -Headers @{ 'User-Agent' = 'helpdesk-installer' }
    $asset = $release.assets | Where-Object { $_.name -match 'x86_64\.exe$' } | Select-Object -First 1
    if (-not $asset) { throw "No x86_64 .exe found in RustDesk release $($release.tag_name)." }
    return $asset.browser_download_url
}

function Wait-RustDeskService([int]$TimeoutSeconds = 180) {
    $deadline = (Get-Date).AddSeconds($TimeoutSeconds)
    while ((Get-Date) -lt $deadline) {
        $svc = Get-Service -Name $ServiceName -ErrorAction SilentlyContinue
        if ((Test-Path $RustDeskExe) -and $svc -and $svc.Status -eq 'Running') { return }
        Start-Sleep -Seconds 3
    }
    throw "RustDesk did not finish installing within $TimeoutSeconds seconds."
}

function Install-RustDesk {
    $source = $Installer
    if (-not $source) { $source = Get-LatestInstallerUrl }
    $ext = if ($source -match '\.msi(\?|$)') { '.msi' } else { '.exe' }
    $setupFile = Join-Path $env:TEMP ("rustdesk-setup-" + [guid]::NewGuid().ToString('N') + $ext)

    if (Test-Path -LiteralPath $source) {
        Copy-Item -LiteralPath $source -Destination $setupFile
    } else {
        Write-Host "Downloading $source"
        Invoke-WebRequest -UseBasicParsing -Uri $source -OutFile $setupFile
    }
    try {
        Write-Host 'Installing RustDesk...'
        if ($ext -eq '.msi') {
            $p = Start-Process -FilePath 'msiexec.exe' -ArgumentList @('/i', "`"$setupFile`"", '/qn', '/norestart') -Wait -PassThru
            if ($p.ExitCode -ne 0 -and $p.ExitCode -ne 3010) { throw "msiexec failed with exit code $($p.ExitCode)." }
        } else {
            # Wait for the setup process only (Start-Process -Wait would also wait for the
            # RustDesk processes it launches), then poll for the service below.
            $p = Start-Process -FilePath $setupFile -ArgumentList '--silent-install' -PassThru
            if (-not $p.WaitForExit(600000)) { throw 'RustDesk setup did not finish within 10 minutes.' }
            Start-Sleep -Seconds 5
        }
        Wait-RustDeskService
    } finally {
        Start-Sleep -Seconds 5   # the setup process may still hold the file briefly
        Remove-Item -LiteralPath $setupFile -Force -ErrorAction SilentlyContinue
    }
}

function Read-Secret([string]$Prompt) {
    $secure = Read-Host -Prompt $Prompt -AsSecureString
    $bstr = [Runtime.InteropServices.Marshal]::SecureStringToBSTR($secure)
    try { return [Runtime.InteropServices.Marshal]::PtrToStringBSTR($bstr) }
    finally { [Runtime.InteropServices.Marshal]::ZeroFreeBSTR($bstr) }
}

# `rustdesk.exe --assign` registers this device with the console's /api/devices/cli.
# The console replies with nothing on success (the client prints "Done!") and with a
# short reason otherwise, e.g. "User 'x' not found." or a missing token permission.
function Get-AssignArguments {
    $fields = [ordered]@{
        '--user_name' = $EnrollUser; '--device_group_name' = $DeviceGroup
        '--strategy_name' = $Strategy; '--address_book_name' = $AddressBook
        '--device_name' = $DeviceName; '--note' = $Note
    }
    $assign = @('--assign', '--token', $EnrollToken)
    foreach ($flag in $fields.Keys) {
        if ($fields[$flag]) { $assign += @($flag, $fields[$flag]) }
    }
    return , $assign
}

if (-not (Test-IsAdmin)) { throw 'Run this script as Administrator.' }

if ($PSCmdlet.ParameterSetName -eq 'Fields') {
    $ConfigString = ConvertTo-ConfigString $IdServer $Key $RelayServer $ApiServer
}

# Ask before installing, so whoever runs the kit answers once and can walk away.
if (-not $EnrollToken -and $env:CORTENDESK_ENROLL_TOKEN) { $EnrollToken = $env:CORTENDESK_ENROLL_TOKEN }
# No prompt without an interactive user: RMM agents run as a service and would hang here.
if ($PromptEnrollToken -and -not $EnrollToken -and [Environment]::UserInteractive -and -not [Console]::IsInputRedirected) {
    $EnrollToken = (Read-Secret 'CortenDesk enrollment token (Enter to skip enrollment)').Trim()
    if ($EnrollToken -and -not ($EnrollUser -or $DeviceGroup -or $Strategy -or $AddressBook -or $DeviceName -or $Note)) {
        $EnrollUser = (Read-Host 'Console user who owns this device (Enter for none)').Trim()
        $DeviceGroup = (Read-Host 'Device group (Enter for none)').Trim()
    }
}
if ($EnrollToken -and (Get-AssignArguments).Count -eq 3) {
    throw 'Enrollment needs at least one of -EnrollUser, -DeviceGroup, -Strategy, -AddressBook, -DeviceName or -Note.'
}

if ($Installer -or -not (Test-Path $RustDeskExe)) {
    Install-RustDesk
} else {
    Write-Host "RustDesk already installed at $RustDeskExe; re-pointing it at your server."
    Start-Service -Name $ServiceName -ErrorAction SilentlyContinue
    Wait-RustDeskService -TimeoutSeconds 30
}

$out = Invoke-RustDesk @('--config', $ConfigString)
if ($out -match 'required') { throw "Applying the server config failed: $out" }
Write-Host 'Server config applied.'

if ($RandomPassword -and -not $Password) {
    $alphabet = 'abcdefghijkmnpqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ23456789'
    $rng = New-Object Security.Cryptography.RNGCryptoServiceProvider
    $bytes = New-Object byte[] 16
    $rng.GetBytes($bytes)
    $Password = -join ($bytes | ForEach-Object { $alphabet[$_ % $alphabet.Length] })
}
if ($Password) {
    $out = Invoke-RustDesk @('--password', $Password)
    if ($out -notmatch 'Done') { throw "Setting the permanent password failed: $out" }
}

# Restart so the service re-registers with the new server right away.
Restart-Service -Name $ServiceName
$id = ''
for ($i = 0; $i -lt 10 -and -not $id; $i++) {
    Start-Sleep -Seconds 3
    $id = Invoke-RustDesk @('--get-id')
}

$enrolled = $false
if ($EnrollToken) {
    $out = Invoke-RustDesk (Get-AssignArguments)
    if ($out -notmatch 'Done') { throw "Enrolling the device in the console failed: $out" }
    $enrolled = $true
}

Write-Host ''
Write-Host "RustDesk ID: $id"
if ($Password) { Write-Host "Password:    $Password" }
if ($enrolled) {
    Write-Host 'Console:     enrolled'
} else {
    Write-Host 'Console:     not enrolled here. Sign in from RustDesk (Settings > Account) with a console account to claim it.'
}