diff --git a/Borealis.ps1 b/Borealis.ps1 index 79080f7..350582f 100644 --- a/Borealis.ps1 +++ b/Borealis.ps1 @@ -340,20 +340,29 @@ function Ensure-AgentTasks { if (-not (Test-Path $regScript)) { Write-Host "Register helper script not found: $regScript" -ForegroundColor Red; return } if (-not (Test-Path $wdSource)) { Write-Host "Watchdog script not found: $wdSource" -ForegroundColor Red; return } - $psi = New-Object System.Diagnostics.ProcessStartInfo - $psi.FileName = 'powershell.exe' - $psi.Verb = 'runas' - $psi.UseShellExecute = $true - $psi.ArgumentList = @( - '-NoProfile','-ExecutionPolicy','Bypass', - '-File', $regScript, - '-SupName', $supName, - '-PythonExe', $py, - '-SupScript', $supScript, - '-WdName', $wdName, - '-WdSource', $wdSource - ) - try { $proc = [System.Diagnostics.Process]::Start($psi); $proc.WaitForExit() } catch { + # Launch registrar elevated using -EncodedCommand to avoid quoting/binding issues + $qSupName = $supName -replace "'","''" + $qPy = $py -replace "'","''" + $qSupScript= $supScript-replace "'","''" + $qWdName = $wdName -replace "'","''" + $qWdSource = $wdSource -replace "'","''" + $qRegScript= $regScript-replace "'","''" + $inline = @" +`$p = @{ + SupName = '$qSupName' + PythonExe = '$qPy' + SupScript = '$qSupScript' + WdName = '$qWdName' + WdSource = '$qWdSource' +} +& '$qRegScript' @p +"@ + $bytes = [System.Text.Encoding]::Unicode.GetBytes($inline) + $encoded = [Convert]::ToBase64String($bytes) + $argList = @('-NoProfile','-ExecutionPolicy','Bypass','-EncodedCommand', $encoded) + try { + Start-Process -FilePath 'powershell.exe' -ArgumentList ($argList -join ' ') -Verb RunAs -Wait | Out-Null + } catch { Write-Host "Failed to elevate for task registration." -ForegroundColor Red } } diff --git a/Data/Agent/Scripts/register_agent_tasks.ps1 b/Data/Agent/Scripts/register_agent_tasks.ps1 index 6852746..c1184ea 100644 --- a/Data/Agent/Scripts/register_agent_tasks.ps1 +++ b/Data/Agent/Scripts/register_agent_tasks.ps1 @@ -1,4 +1,4 @@ -1param( +param( [Parameter(Mandatory=$true)] [string]$SupName, [Parameter(Mandatory=$true)] [string]$PythonExe, [Parameter(Mandatory=$true)] [string]$SupScript, @@ -40,4 +40,3 @@ try { Write-Error $_ exit 1 } - diff --git a/Data/Agent/agent_deployment.py b/Data/Agent/agent_deployment.py index 15d9e8e..e6ead35 100644 --- a/Data/Agent/agent_deployment.py +++ b/Data/Agent/agent_deployment.py @@ -108,14 +108,14 @@ def ensure_user_logon_task(paths): pyw = paths.get("venv_pythonw") or paths["venv_python"] cmd = f'"{pyw}" -W ignore::SyntaxWarning "{paths["agent_script"]}"' # Try create non-elevated - q = run(["schtasks.exe", "/Query", "/TN", task_name]) + q = run(["schtasks.exe", "/Query", "/TN", task_name], capture=True) if q.returncode == 0: - d = run(["schtasks.exe", "/Delete", "/TN", task_name, "/F"]) + d = run(["schtasks.exe", "/Delete", "/TN", task_name, "/F"], capture=True) if d.returncode != 0: pass - c = run(["schtasks.exe", "/Create", "/SC", "ONLOGON", "/TN", task_name, "/TR", cmd, "/F", "/RL", "LIMITED"]) + c = run(["schtasks.exe", "/Create", "/SC", "ONLOGON", "/TN", task_name, "/TR", cmd, "/F", "/RL", "LIMITED"], capture=True) if c.returncode == 0: - run(["schtasks.exe", "/Run", "/TN", task_name]) + run(["schtasks.exe", "/Run", "/TN", task_name], capture=True) return True # Elevated fallback using ScheduledTasks cmdlets for better reliability ps = f""" @@ -160,4 +160,3 @@ def main(argv): if __name__ == "__main__": sys.exit(main(sys.argv)) - diff --git a/Data/Server/server.py b/Data/Server/server.py index 0c05751..fd872ff 100644 --- a/Data/Server/server.py +++ b/Data/Server/server.py @@ -1278,10 +1278,22 @@ def on_agent_heartbeat(data): hostname = data.get("hostname") if hostname: + # Avoid duplicate entries per-hostname. Prefer non-script agents over script helpers. + try: + is_current_script = isinstance(agent_id, str) and agent_id.lower().endswith('-script') + except Exception: + is_current_script = False for aid, info in list(registered_agents.items()): - if aid != agent_id and info.get("hostname") == hostname: - registered_agents.pop(aid, None) - agent_configurations.pop(aid, None) + if aid == agent_id: + continue + if info.get("hostname") == hostname: + if info.get('is_script_agent') and not is_current_script: + # Replace script helper with full agent record + registered_agents.pop(aid, None) + agent_configurations.pop(aid, None) + else: + # Keep existing non-script agent; do not evict it for script heartbeats + pass rec = registered_agents.setdefault(agent_id, {}) rec["agent_id"] = agent_id