# ensure-vnc-password.ps1 - keep the UltraVNC service VNC password in sync
# with the password the user types in run.bat. Called from run.bat.
#
# It only WRITES when the password actually differs (so a normal launch with
# the same password does nothing and asks for NO admin rights). When it must
# change the password it re-launches itself elevated (one UAC prompt) and
# restarts the uvnc_service so the new password takes effect.
#
# Usage: powershell -NoProfile -ExecutionPolicy Bypass -File ensure-vnc-password.ps1 -Password "<pw>"
# Exit codes: 0 = password matches (or was 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)) {
    # No installed UltraVNC service config (e.g. portable copy / TightVNC) - nothing to sync.
    Write-Info "no %ProgramData%\UltraVNC\ultravnc.ini - nothing to sync."
    exit 0
}

# --- VNC password obfuscation: DES-ECB(key=bitrev([23,82,107,6,35,78,88,7]), pw8) + 1-byte checksum ---
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 }
function Get-ExpectedPasswd([string]$pw){
    $fixedkey = [byte[]](23,82,107,6,35,78,88,7)
    $kr = New-Object byte[] 8; for($i=0;$i -lt 8;$i++){ $kr[$i] = Invoke-BitRev $fixedkey[$i] }
    $pb = [System.Text.Encoding]::ASCII.GetBytes($pw)
    if ($pb.Length -gt 8) { $pb = $pb[0..7] }
    $plain = New-Object byte[] 8; for($i=0;$i -lt 8;$i++){ if($i -lt $pb.Length){ $plain[$i] = $pb[$i] } }
    $des = New-Object System.Security.Cryptography.DESCryptoServiceProvider
    $des.Mode = [System.Security.Cryptography.CipherMode]::ECB
    $des.Padding = [System.Security.Cryptography.PaddingMode]::None
    $des.Key = $kr
    $cipher = $des.CreateEncryptor().TransformFinalBlock($plain,0,8)
    $sum = 0; foreach($b in $cipher){ $sum = ($sum + $b) -band 0xFF }
    $h = ($cipher | ForEach-Object { $_.ToString("X2") }) -join ""
    return $h + $sum.ToString("X2")
}

function Get-CurrentPasswd{
    $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()
}

$expected = Get-ExpectedPasswd $Password
$current  = Get-CurrentPasswd
if ($current -ieq $expected) {
    Write-Info "UltraVNC password already matches - nothing to do."
    exit 0
}

Write-Info "UltraVNC password differs (ini=$current, expected=$expected). Need to set it (admin required)."

# 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 }

# Are we admin?
$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 }
    $now = Get-CurrentPasswd
    if ($now -ieq $expected) { Write-Info "password set OK."; exit 0 }
    Write-Info "password still not updated after elevated run (ini=$now)."; exit 1
}

# --- admin path: actually set the password ---
function Set-Tool{
    param([string]$exe)
    if (-not (Test-Path $exe)) { return $false }
    try {
        $psi = New-Object System.Diagnostics.ProcessStartInfo
        $psi.FileName = $exe
        $psi.Arguments = "`"$Password`""
        $psi.UseShellExecute = $false
        $p = [System.Diagnostics.Process]::Start($psi)
        if (-not $p.WaitForExit(5000)) { $p.Kill() }
        return ($p.ExitCode -eq 0)
    } catch { return $false }
}

$ok = $false
if (Set-Tool (Join-Path $uvnc "createpassword.exe")) { $ok = $true }
$now = Get-CurrentPasswd
if ($now -ine $expected -and (Test-Path (Join-Path $uvnc "setpasswd.exe"))) {
    # fallback to setpasswd.exe "<pw>"
    try {
        $psi = New-Object System.Diagnostics.ProcessStartInfo
        $psi.FileName = Join-Path $uvnc "setpasswd.exe"
        $psi.Arguments = "`"$Password`""
        $psi.UseShellExecute = $false
        $p = [System.Diagnostics.Process]::Start($psi)
        if (-not $p.WaitForExit(5000)) { $p.Kill() }
    } catch {}
    $now = Get-CurrentPasswd
}

if ($now -ine $expected) {
    Write-Info "tools did not update the password (ini=$now). Set it manually via UltraVNC Admin Properties as admin."
    exit 1
}
Write-Info "password updated in $ini."

# Restart the service so it reloads the new password.
$svc = "uvnc_service"
$svcExists = $false
try { $null = (& sc.exe query $svc 2>$null); $svcExists = ($LASTEXITCODE -eq 0) } catch {}
if ($svcExists) {
    Write-Info "restarting service $svc ..."
    & net.exe stop $svc 2>$null | Out-Null
    & net.exe start $svc 2>$null | Out-Null
}
exit 0