# ensure-vnc-password.ps1 - keep the UltraVNC VNC password in sync with the
# password the user types in run.bat. Called from run.bat.
#
# It checks whether the running VNC server already ACCEPTS the given password
# (a real RFB VNC-Auth handshake against 127.0.0.1:5900 - the same mechanism
# noVNC uses). If it does, nothing happens and NO admin rights are needed.
# If it does not (or no server is up), it sets the UltraVNC service password
# to the given value (one UAC prompt) and restarts uvnc_service so it reloads.
#
# Usage: powershell -NoProfile -ExecutionPolicy Bypass -File ensure-vnc-password.ps1 -Password "<pw>"
# Exit codes: 0 = password already accepted (or set OK); 1 = could not set it.

param([Parameter(Mandatory=$true)][string]$Password)

$ErrorActionPreference = "Stop"
$ini = Join-Path $env:ProgramData "UltraVNC\ultravnc.ini"
function Write-Info($m){ Write-Host "ensure-vnc-password: $m" }

if (-not (Test-Path $ini)) {
    Write-Info "no %ProgramData%\UltraVNC\ultravnc.ini - nothing to sync."
    exit 0
}

# ---- bit-reverse a byte (VNC key bytes are bit-reversed) ----
function Invoke-BitRev([byte]$b){ $r=0; for($i=0;$i -lt 8;$i++){ $r = $r -bor ((($b -shr $i) -band 1) -shl (7-$i)) }; return [byte]$r }

# ---- real RFB VNC-Auth probe. Returns 'ok','rejected','unreachable','noserver' ----
function Test-VncAuth([string]$pw,[string]$target,[int]$port){
    $c = New-Object System.Net.Sockets.TcpClient
    try { $c.Connect($target,$port) } catch { return 'unreachable' }
    try {
        $s = $c.GetStream(); $s.ReadTimeout = 8000; $s.WriteTimeout = 8000
        $rb = New-Object byte[] 64
        $got = 0
        while($got -lt 12){ $n = $s.Read($rb,$got,12-$got); if($n -le 0){ return 'unreachable' }; $got += $n }
        if (-not ([System.Text.Encoding]::ASCII.GetString($rb,0,12).StartsWith("RFB "))) { return 'noserver' }
        $s.Write([System.Text.Encoding]::ASCII.GetBytes("RFB 003.008`n"),0,12)
        $n = $s.Read($rb,0,64); if($n -lt 1){ return 'unreachable' }
        $numTypes = $rb[0]
        if($numTypes -eq 0 -or $n -lt (1+$numTypes)){ return 'noserver' }
        $types = $rb[1..($numTypes)]
        if(-not ($types -contains 2)){ return 'noserver' }   # VNC-Auth not offered (e.g. MS-Logon only)
        $s.Write([byte[]](2),0,1)                              # choose VNC-Auth
        $got = 0
        while($got -lt 16){ $n = $s.Read($rb,$got,16-$got); if($n -le 0){ return 'unreachable' }; $got += $n }
        $challenge = $rb[0..15]
        $pb = [System.Text.Encoding]::ASCII.GetBytes($pw)
        if($pb.Length -gt 8){ $pb = $pb[0..7] }
        $key = New-Object byte[] 8
        for($i=0;$i -lt 8;$i++){ $kb = if($i -lt $pb.Length){[int]$pb[$i]}else{0}; $key[$i] = Invoke-BitRev ([byte]$kb) }
        $des = New-Object System.Security.Cryptography.DESCryptoServiceProvider
        $des.Mode = [System.Security.Cryptography.CipherMode]::ECB
        $des.Padding = [System.Security.Cryptography.PaddingMode]::None
        $des.Key = $key
        $resp = $des.CreateEncryptor().TransformFinalBlock($challenge,0,16)
        $s.Write($resp,0,16)
        $n = $s.Read($rb,0,4); if($n -lt 4){ return 'rejected' }
        $status = ($rb[0]*16777216)+($rb[1]*65536)+($rb[2]*256)+$rb[3]
        if($status -eq 0){ return 'ok' } else { return 'rejected' }
    } finally { $c.Close() }
}

$probe = Test-VncAuth $Password "127.0.0.1" 5900
if($probe -eq 'ok'){
    Write-Info "UltraVNC already accepts this password - nothing to do."
    exit 0
}
Write-Info "UltraVNC password needs setting (probe=$probe). Need admin."

# Locate UltraVNC install + password tools.
$uvncDirs = @(
    "$env:ProgramFiles\uvnc bvba\UltraVNC",
    "$env:ProgramFiles\UltraVNC",
    "${env:ProgramFiles(x86)}\uvnc bvba\UltraVNC",
    "${env:ProgramFiles(x86)}\UltraVNC"
)
$uvnc = $uvncDirs | Where-Object { Test-Path "$_\winvnc.exe" } | Select-Object -First 1
if(-not $uvnc){ Write-Info "UltraVNC install not found - cannot set password."; exit 1 }

function Get-IniPasswd{
    $line = Get-Content $ini -ErrorAction SilentlyContinue | Where-Object { $_ -match '^\s*passwd\s*=' -and $_ -notmatch 'passwd2' } | Select-Object -First 1
    if(-not $line){ return "" }
    return ($line -split '=',2)[1].Trim()
}

function Invoke-SetTool([string]$exe,[string]$pw){
    if(-not (Test-Path $exe)){ return $false }
    try {
        $psi = New-Object System.Diagnostics.ProcessStartInfo
        $psi.FileName = $exe; $psi.Arguments = "`"$pw`""; $psi.UseShellExecute = $false
        $p = [System.Diagnostics.Process]::Start($psi)
        if(-not $p.WaitForExit(5000)){ $p.Kill() }
        return ($p.ExitCode -eq 0)
    } catch { return $false }
}

function Do-Set([string]$pw){
    $before = Get-IniPasswd
    [void](Invoke-SetTool (Join-Path $uvnc "createpassword.exe") $pw)
    $after = Get-IniPasswd
    if($after -ieq $before){
        [void](Invoke-SetTool (Join-Path $uvnc "setpasswd.exe") $pw)
        $after = Get-IniPasswd
    }
    if($after -ieq $before){ Write-Info "password tools did not change $ini (still $after)."; return $false }
    # restart the service so the running server reloads (only if it was running)
    try { $q = (& sc.exe query uvnc_service 2>$null); $running = ($q -match 'RUNNING') } catch { $running = $false }
    if($running){
        Write-Info "restarting uvnc_service ..."
        & net.exe stop uvnc_service 2>$null | Out-Null
        & net.exe start uvnc_service 2>$null | Out-Null
    }
    return $true
}

$isAdmin = ([Security.Principal.WindowsPrincipal][Security.Principal.WindowsIdentity]::GetCurrent()).IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)

if(-not $isAdmin){
    Write-Info "re-launching elevated to set the password..."
    $p = Start-Process -FilePath "powershell.exe" -Verb RunAs -Wait -PassThru -ArgumentList @("-NoProfile","-ExecutionPolicy","Bypass","-File","`"$PSCommandPath`"","-Password","`"$Password`"")
    if($p.ExitCode -ne 0){ Write-Info "elevated run failed (exit $($p.ExitCode))."; exit 1 }
    # confirm by re-probing (if a server is up now)
    $again = Test-VncAuth $Password "127.0.0.1" 5900
    if($again -eq 'ok'){ Write-Info "password set OK."; exit 0 }
    if($again -eq 'unreachable'){ Write-Info "password written; no server up to verify (will be used on next start)."; exit 0 }
    Write-Info "password set but server still rejects it (probe=$again)."; exit 1
}

# admin path
if(Do-Set $Password){ Write-Info "password set OK."; exit 0 } else { exit 1 }