#include-once
#include <File.au3>
#include <Inet.au3>

#Region ##### DECLARE GLOBALS #####
Global $g_RustDesk_RegKey = "Version"
Global $g_RustDesk_Tags = "https://github.com/rustdesk/rustdesk/tags.atom"

Global $g_RustDesk_CurrentVersion, $g_RustDesk_LatestVersion, $g_RustDesk_Download_URL, $g_RustDesk_InstalledPath
Global $g_RustDesk_TempInstallerPath, $g_RustDesk_InstalledPathDIR, $g_RustDesk_HKLM, $g_RustDesk_RegPath
Global $g_RustDesk_DebugPath = @DesktopDir & "\RD_DBG.txt"
Global $g_RustDesk_DebugToFile = True
#EndRegion ##### DECLARE GLOBALS #####
#Region ##### INIT GLOBALS #####
_DebugInit()
_Debug(">> Firing up RustDesk UDF!")
If @AutoItX64 Then
	_Debug("+> Running in an X64 Environment!")
	$g_RustDesk_RegPath = "HKLM\SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall\RustDesk"
	$g_RustDesk_HKLM = "HKLM"
Else
	_Debug("-> Running is an X86 Environment!")
	$g_RustDesk_RegPath = "HKLM64\SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall\RustDesk"
	$g_RustDesk_HKLM = "HKLM64"
EndIf
$g_RustDesk_TempInstallerPath = _TempFile(Default, Default, ".msi")
$g_RustDesk_CurrentVersion = _RustDesk_GetLocalVersion()
$g_RustDesk_LatestVersion = _RustDesk_GetLatestVersion()
$g_RustDesk_Download_URL = "https://github.com/rustdesk/rustdesk/releases/download/" & $g_RustDesk_LatestVersion & "/rustdesk-" & $g_RustDesk_LatestVersion & "-x86_64.msi"
If _RustDesk_isInstalled() Then
	$g_RustDesk_InstalledPath = RegRead($g_RustDesk_HKLM & "\SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall\RustDesk", "InstallLocation") & "RustDesk.exe"
	$g_RustDesk_InstalledPathDIR = RegRead($g_RustDesk_HKLM & "\SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall\RustDesk", "InstallLocation")
Else
	$g_RustDesk_InstalledPath = @ProgramFilesDir & "\RustDesk\RustDesk.exe"
	$g_RustDesk_InstalledPathDIR = @ProgramFilesDir & "\RustDesk"
EndIf
;_Debug($g_RustDesk_InstalledPath)
_Debug("+> RustDesk UDF Loaded!")
#EndRegion ##### INIT GLOBALS #####

#Region ##### RUSTDESK UDFs #####
Func _RustDesk_Deploy($pVisible = True, $pTimeout = 30)
	_Debug(">> _RustDesk_Deploy() : Init.")
	Local $iTimeout = 0
	Local $bShow = ($pVisible) ? @SW_SHOW : @SW_HIDE
	Local $sArguments = ' /i "' & $g_RustDesk_TempInstallerPath & '" /norestart /passive'
	If $pVisible = False Then $sArguments &= " /quiet /qn"
	_Download($g_RustDesk_Download_URL, $g_RustDesk_TempInstallerPath, "Rust Desk...")
	ShellExecuteWait("msiexec.exe", ' /i "' & $g_RustDesk_TempInstallerPath & $sArguments, @TempDir, Default, $bShow)
	; The Service, in all my tests, is created and immediately started after RustDesk finishes installing through the MSI.
	; But, a very small window does exist in-between RustDesk finishing the installation and the service starting. So we do have to wait/check.
	;	We need this service to function, because it handles getting the RD ID and setting the RD Password.
	While Not _ServiceExists("RustDesk")
		Sleep(3 * 1000)
		$iTimeout += 3
		If $iTimeout >= $pTimeout / 2 Then ;Only when we've waited long enough, we attempt a manual service install.
			; We do NOT use ShellExecuteWait here, because the STDOUT stream gets locked into permenance (RustDesk never stops the stream).
			ShellExecute($g_RustDesk_InstalledPath, "--install-service", $g_RustDesk_InstalledPathDIR, Default, $bShow) ; --install-service will also AUTO START the service after it is installed, per testing on 1.3.6
			_Debug("!> _RustDesk_Deploy(): TIMEOUT. Initiated manual service creation.")
			$iTimeout = 0
			ExitLoop
		EndIf
	WEnd
	While Not _ServiceStatus("RustDesk") = "RUNNING" ; Give it a chance to come online manually or skip to the bottom if we're running!
		Sleep(3 * 1000)
		$iTimeout += 3
		If $iTimeout >= $pTimeout / 2 Then
			_Debug("!> _RustDesk_Deploy(): TIMEOUT. Manual create/start failed!")
			Return SetError(10, 0, False)
		EndIf
	WEnd
	_Debug("+> _RustDesk_Deploy(): SUCCESS!")
EndFunc   ;==>_RustDesk_Deploy

Func _RustDesk_Start($pTimeout = 30)
	Local $iTimer = 0, $sReturn
	Local $sAppName = "RustDesk"
	_Debug(">> _RustDesk_Start(" & $pTimeout & "): Initialization.")
	If Not _ServiceExists("RustDesk") Then
		_Debug("!> _RustDesk_Start(): Service does not exist!")
		_RustDesk_InstallService($pTimeout)
	Else
		$sReturn = _ServiceStart($sAppName)
		If @error Then
			_Debug("!> _RustDesk_Start() : Failed to start Rust Desk Service.")
			Return SetError(10, 0, False)
		EndIf
		If @extended = 1 Then ;Service is starting.
			_Debug("+> _RustDesk_Start(): Service has starting status...")
			While Not _ServiceStatus($sAppName) = "RUNNING"
				Sleep(1000)
				$iTimer += 1
				If $iTimer >= $pTimeout Then
					_Debug("!> _RustDesk_Start(): TIMEOUT while waiting for the RustDesk Service to start.")
					Return SetError(2, 0, False)
				EndIf
			WEnd
			_Debug("+> _RustDesk_Start(): Service RUNNING!")
			Return True
		ElseIf @extended = 2 Then ; Service is already RUNNING
			_Debug("+> _RustDesk_Start(): Service RUNNING!")
			Return True
		EndIf
	EndIf
EndFunc   ;==>_RustDesk_Start


Func _RustDesk_GetLocalVersion()
	Local $sVersion = StringLeft(RegRead($g_RustDesk_RegPath, $g_RustDesk_RegKey), 5)
	If StringLen($sVersion) <> 5 Then
		_Debug("-> _RustDesk_GetLocalVersion(): " & "Not Installed!")
		Return False
	Else
		_Debug("+> _RustDesk_GetLocalVersion(): " & $sVersion)
		Return $sVersion
	EndIf
EndFunc   ;==>_RustDesk_GetLocalVersion

Func _RustDesk_InstallService($pTimeout = 30)
	Local $sReturn, $sProcess_Status, $iTimer = 0
	Local $sAppName = "RustDesk"
	_Debug(">> _RustDesk_InstallService(" & $pTimeout & ")")
	If Not _RustDesk_isInstalled() Then
		_Debug("!> _RustDesk_InstallService() : Tried to install the RustDesk SERVICE, before _RustDesk_ was properly installed.")
		Return SetError(10, 0, False)
	ElseIf Not _ServiceExists($sAppName) Then
		ShellExecute($g_RustDesk_InstalledPath, "--install-service", $g_RustDesk_InstalledPathDIR, Default, @SW_HIDE)
		While Not _ServiceExists("RustDesk")
			Sleep(1000)
			$iTimer += 1
			If $iTimer >= $pTimeout Then
				_Debug("!> _RustDesk_InstallService(): TIMEOUT waiting for manual service install!")
				Return SetError(10, 0, False)
			EndIf
		WEnd
		$sReturn = _ServiceStart($sAppName)
		If @error Then Return SetError(10, 0, False)
		If @extended = 1 Then ;Service is starting.
			While Not _ServiceStatus($sAppName) = "RUNNING"
				Sleep(1000)
				$iTimer += 1
				If $iTimer >= $pTimeout Then Return SetError(10, 0, False)
			WEnd
			Return True
		ElseIf @extended = 2 Then
			Return True ; Service is already RUNNING
		EndIf
	EndIf
EndFunc   ;==>_RustDesk_InstallService



Func _RuseDesk_RelayConfig($pConfig)
	ShellExecute($g_RustDesk_InstalledPath, "--config " & $pConfig, $g_RustDesk_InstalledPathDIR, Default, @SW_HIDE)
	;We can't ShellExecuteWait or RunWait this process because of the way RustDesk handles STDOUT.
EndFunc   ;==>_RuseDesk_RelayConfig



Func _RustDesk_GetLatestVersion()
	Local $sRAW, $aVersion
	$sRAW = _INetGetSource($g_RustDesk_Tags)
	If @error Or $sRAW = "" Then Return SetError(4, 0, False)
	$aVersion = StringRegExp($sRAW, "<entry>\n+.+\n+.+\n+.+\n.+<title>(.+)<\/title>", 1)
	If IsArray($aVersion) Then
		_Debug("+> _RustDesk_GetLatestVersion(): " & $aVersion[0])
		Return $aVersion[0]
	Else
		Return SetError(5, 0, False)
	EndIf
EndFunc   ;==>_RustDesk_GetLatestVersion

Func _RustDesk_isInstalled()
	RegRead($g_RustDesk_RegPath, $g_RustDesk_RegKey)
	If @error Then
		_Debug("-> _RustDesk_isInstalled(): " & "False")
		Return False
	Else
		_Debug("+> _RustDesk_isInstalled(): " & "True")
		Return True
	EndIf
EndFunc   ;==>_RustDesk_isInstalled

Func _RustDesk_GetID()
	Local $lReturn
	$lReturn = _Run($g_RustDesk_InstalledPath, "--get-id", $g_RustDesk_InstalledPathDIR)
	If @error Then
		Return SetError(9, 0, False)
	EndIf
	$lReturn = StringStripWS($lReturn, $STR_STRIPALL)
	_Debug(">> _Get_RustDesk_ID(): " & $lReturn)
	Return $lReturn
EndFunc   ;==>_RustDesk_GetID

Func _RustDesk_SetPassword($pPassword)
	_Debug(">> _Set_RustDesk_Password(" & $pPassword & ")")
	Local $lReturn = _Run($g_RustDesk_InstalledPath, "--password " & $pPassword, $g_RustDesk_InstalledPathDIR)
	If Not $lReturn = "Done!" Then
		Return SetError(8, 0, False)
	Else
		Return True
	EndIf
EndFunc   ;==>_RustDesk_SetPassword
#EndRegion ##### RUSTDESK UDFs #####

#Region ##### MISC UDFs #####

Func _Run($pExe, $pArguments = Default, $pPath = Default, $pHide = True)
	_Debug(">> _Run(" & $pExe & "," & $pArguments & "," & $pPath & "," & String($pHide) & ")")
	If $pHide = True Then
		$pHide = @SW_HIDE
	Else
		$pHide = @SW_SHOW
	EndIf

	Local $hProcess, $sOutput = "", $sLine
	If $pArguments = Default Then
		$hProcess = Run($pExe, $pPath, $pHide)
		If @error Then Return SetError(6, 0, False)
	Else
		$hProcess = Run($pExe & " " & $pArguments, $pPath, $pHide, $STDERR_CHILD & $STDOUT_CHILD)
		If @error Then Return SetError(7, 0, False)
	EndIf

	While 1
		$sLine = StdoutRead($hProcess)
		If @error Then ExitLoop ; Exit loop if there's no more data
		$sOutput &= $sLine
	WEnd

	ProcessClose($hProcess)
	_Debug("-> _Run(): $sOutput = " & $sOutput & "")
	Return $sOutput
EndFunc   ;==>_Run

Func _Download($pURL, $pPath, $pText = "Data...")
	_Debug(">> _Download(" & $pURL & "," & $pPath & ")")
	Local $lTempFile, $lBytesReceived, $lPercentage, $lFileSize, $lINet_HNDL
	ProgressOn("Download", "Downloading " & $pText, "0%")
	$lINet_HNDL = InetGet($pURL, $pPath, $INET_FORCERELOAD, $INET_DOWNLOADBACKGROUND)
	If @error Then Return SetError(1, 0, False)
	$lFileSize = InetGetSize($pURL)
	If @error Then SetError(2, 0, False)
	While Not InetGetInfo($lINet_HNDL, 2)
		Sleep(500)
		$lBytesReceived = InetGetInfo($lINet_HNDL, 0)
		If @error Then SetError(3, 0, False)
		$lPercentage = Int($lBytesReceived / $lFileSize * 100)
		ProgressSet($lPercentage, $lPercentage & "%")
	WEnd
	ProgressOff()
EndFunc   ;==>_Download

Func _DebugInit()
	If FileExists($g_RustDesk_DebugPath) Then
		FileDelete($g_RustDesk_DebugPath)
	Else
		_Debug("----- DEBUG LOG BEGIN -----")
	EndIf
EndFunc   ;==>_DebugInit

Func _Debug($pData)
	ConsoleWrite($pData & @CRLF)
	If $g_RustDesk_DebugToFile Then FileWriteLine($g_RustDesk_DebugPath, $pData)
EndFunc   ;==>_Debug

Func _ServiceExists($pName)
	Local $lReturn = _Run("sc.exe", "query " & $pName, @SystemDir)
	If StringInStr($lReturn, "does not exist") Then
		Return False
	Else
		Return True
	EndIf
EndFunc   ;==>_ServiceExists

Func _ServiceStart($pName)
	Local $lReturn = _Run("sc.exe", "start " & $pName, @SystemDir)
	If StringInStr($lReturn, "START_PENDING") Then
		Return SetExtended(1, True)
	ElseIf StringInStr($lReturn, "does not exist") Then
		Return SetError(1, 0, False)
	ElseIf StringInStr($lReturn, "already running") Then
		Return SetExtended(2, True)
	Else
		Return SetError(2, 0, False)
	EndIf
EndFunc   ;==>_ServiceStart


Func _ServiceStatus($pName)
	Local $lReturn = _Run("sc.exe", "query " & $pName, @SystemDir)
	If StringInStr($lReturn, "RUNNING") Then
		Return "RUNNING"
	ElseIf StringInStr($lReturn, "STOPPED") Then
		Return "STOPPED"
	ElseIf StringInStr($lReturn, "PAUSED") Then
		Return "PAUSED"
	ElseIf StringInStr($lReturn, "STARTING") Then
		Return "STARTING"
	ElseIf StringInStr($lReturn, "STOPPING") Then
		Return "STOPPING"
	Else
		SetError(1, 0, False)
	EndIf
EndFunc   ;==>_ServiceStatus
#EndRegion ##### MISC UDFs #####