mirror of
https://github.com/bunny-lab-io/Borealis.git
synced 2025-09-11 01:48:42 -06:00
Continued Work on Remote Script Execution Code
This commit is contained in:
176
Borealis.ps1
176
Borealis.ps1
@@ -19,6 +19,8 @@
|
||||
param(
|
||||
[switch]$Server,
|
||||
[switch]$Agent,
|
||||
[ValidateSet('install','repair','remove','launch','')]
|
||||
[string]$AgentAction = '',
|
||||
[switch]$Vite,
|
||||
[switch]$Flask,
|
||||
[switch]$Quick
|
||||
@@ -27,6 +29,7 @@ param(
|
||||
# Preselect menu choices from CLI args (optional)
|
||||
$choice = $null
|
||||
$modeChoice = $null
|
||||
$agentSubChoice = $null
|
||||
|
||||
if ($Server -and $Agent) {
|
||||
Write-Host "Cannot use -Server and -Agent together." -ForegroundColor Red
|
||||
@@ -42,6 +45,13 @@ if ($Server) {
|
||||
$choice = '1'
|
||||
} elseif ($Agent) {
|
||||
$choice = '2'
|
||||
switch ($AgentAction) {
|
||||
'install' { $agentSubChoice = '1' }
|
||||
'repair' { $agentSubChoice = '2' }
|
||||
'remove' { $agentSubChoice = '3' }
|
||||
'launch' { $agentSubChoice = '4' }
|
||||
default { }
|
||||
}
|
||||
}
|
||||
|
||||
if ($Server) {
|
||||
@@ -82,6 +92,125 @@ $symbols = @{
|
||||
Info = [char]0x2139
|
||||
}
|
||||
|
||||
# Ensure log directories
|
||||
function Ensure-AgentLogDir {
|
||||
$logRoot = Join-Path $scriptDir 'Logs'
|
||||
$agentLogDir = Join-Path $logRoot 'Agent'
|
||||
if (-not (Test-Path $agentLogDir)) { New-Item -ItemType Directory -Path $agentLogDir -Force | Out-Null }
|
||||
return $agentLogDir
|
||||
}
|
||||
|
||||
function Write-AgentLog {
|
||||
param(
|
||||
[string]$FileName,
|
||||
[string]$Message
|
||||
)
|
||||
$dir = Ensure-AgentLogDir
|
||||
$path = Join-Path $dir $FileName
|
||||
$ts = Get-Date -Format s
|
||||
"[$ts] $Message" | Out-File -FilePath $path -Append -Encoding UTF8
|
||||
}
|
||||
|
||||
# Forcefully remove legacy and current Borealis services and tasks
|
||||
function Remove-BorealisServicesAndTasks {
|
||||
param([string]$LogName)
|
||||
$svcNames = @('BorealisAgent','BorealisScriptService','BorealisScriptAgent')
|
||||
foreach ($n in $svcNames) {
|
||||
Write-AgentLog -FileName $LogName -Message "Attempting to stop service: $n"
|
||||
try { sc.exe stop $n 2>$null | Out-Null } catch {}
|
||||
Start-Sleep -Milliseconds 300
|
||||
Write-AgentLog -FileName $LogName -Message "Attempting to delete service: $n"
|
||||
try { sc.exe delete $n 2>$null | Out-Null } catch {}
|
||||
}
|
||||
# Remove scheduled task if it exists
|
||||
$taskName = 'Borealis Agent'
|
||||
Write-AgentLog -FileName $LogName -Message "Attempting to delete scheduled task: $taskName"
|
||||
try { schtasks.exe /Delete /TN "$taskName" /F 2>$null | Out-Null } catch {}
|
||||
}
|
||||
|
||||
# Repair routine: cleans services, ensures venv files, reinstalls and starts BorealisAgent
|
||||
function Repair-BorealisAgent {
|
||||
$logName = 'Repair.log'
|
||||
Write-AgentLog -FileName $logName -Message "=== Repair start ==="
|
||||
$agentDir = Join-Path $scriptDir 'Agent'
|
||||
$venvPython = Join-Path $agentDir 'Scripts\python.exe'
|
||||
$deployScript = Join-Path $agentDir 'Borealis\agent_deployment.py'
|
||||
|
||||
# Aggressive cleanup first
|
||||
Remove-BorealisServicesAndTasks -LogName $logName
|
||||
Start-Sleep -Seconds 1
|
||||
|
||||
# Ensure venv and files exist by reusing the install block
|
||||
Write-AgentLog -FileName $logName -Message "Ensuring agent venv and files"
|
||||
$venvFolder = 'Agent'
|
||||
$agentSourcePath = 'Data\Agent\borealis-agent.py'
|
||||
$agentRequirements = 'Data\Agent\agent-requirements.txt'
|
||||
$agentDestinationFolder = "$venvFolder\Borealis"
|
||||
$venvPythonPath = Join-Path $scriptDir $venvFolder | Join-Path -ChildPath 'Scripts\python.exe'
|
||||
|
||||
if (-not (Test-Path "$venvFolder\Scripts\Activate")) {
|
||||
$pythonForVenv = $pythonExe
|
||||
if (-not (Test-Path $pythonForVenv)) {
|
||||
$pyCmd = Get-Command py -ErrorAction SilentlyContinue
|
||||
$pythonCmd = Get-Command python -ErrorAction SilentlyContinue
|
||||
if ($pyCmd) { $pythonForVenv = $pyCmd.Source }
|
||||
elseif ($pythonCmd) { $pythonForVenv = $pythonCmd.Source }
|
||||
else { throw "Python not found. Please run Server setup (option 1) first." }
|
||||
}
|
||||
Write-AgentLog -FileName $logName -Message "Creating venv using: $pythonForVenv"
|
||||
& $pythonForVenv -m venv $venvFolder | Out-Null
|
||||
}
|
||||
|
||||
if (Test-Path $agentSourcePath) {
|
||||
Write-AgentLog -FileName $logName -Message "Refreshing Agent/Borealis files"
|
||||
Remove-Item $agentDestinationFolder -Recurse -Force -ErrorAction SilentlyContinue
|
||||
New-Item -Path $agentDestinationFolder -ItemType Directory -Force | Out-Null
|
||||
Copy-Item 'Data\Agent\borealis-agent.py' $agentDestinationFolder -Recurse
|
||||
Copy-Item 'Data\Agent\agent_info.py' $agentDestinationFolder -Recurse
|
||||
Copy-Item 'Data\Agent\agent_roles.py' $agentDestinationFolder -Recurse
|
||||
Copy-Item 'Data\Agent\Python_API_Endpoints' $agentDestinationFolder -Recurse
|
||||
Copy-Item 'Data\Agent\windows_script_service.py' $agentDestinationFolder -Force
|
||||
Copy-Item 'Data\Agent\agent_deployment.py' $agentDestinationFolder -Force
|
||||
Copy-Item 'Data\Agent\tray_launcher.py' $agentDestinationFolder -Force
|
||||
if (Test-Path 'Data\Agent\Borealis.ico') { Copy-Item 'Data\Agent\Borealis.ico' $agentDestinationFolder -Force }
|
||||
}
|
||||
|
||||
. "$venvFolder\Scripts\Activate"
|
||||
if (Test-Path $agentRequirements) {
|
||||
Write-AgentLog -FileName $logName -Message "Installing agent requirements"
|
||||
& $venvPythonPath -m pip install --disable-pip-version-check -q -r $agentRequirements | Out-Null
|
||||
}
|
||||
|
||||
# Install service via deployment helper (sets PythonHome/PythonPath)
|
||||
Write-AgentLog -FileName $logName -Message "Running agent_deployment.py ensure-all"
|
||||
$deployScript = Join-Path $agentDestinationFolder 'agent_deployment.py'
|
||||
& $venvPythonPath -W ignore::SyntaxWarning $deployScript ensure-all *>&1 | Tee-Object -FilePath (Join-Path (Ensure-AgentLogDir) $logName) -Append | Out-Null
|
||||
|
||||
# Start the service explicitly
|
||||
Write-AgentLog -FileName $logName -Message "Starting service BorealisAgent"
|
||||
try { sc.exe start BorealisAgent 2>$null | Out-Null } catch {}
|
||||
Start-Sleep -Seconds 2
|
||||
Write-AgentLog -FileName $logName -Message "Query service state"
|
||||
sc.exe query BorealisAgent *>> (Join-Path (Ensure-AgentLogDir) $logName)
|
||||
Write-AgentLog -FileName $logName -Message "=== Repair end ==="
|
||||
}
|
||||
|
||||
function Remove-BorealisAgent {
|
||||
$logName = 'Removal.log'
|
||||
Write-AgentLog -FileName $logName -Message "=== Removal start ==="
|
||||
Remove-BorealisServicesAndTasks -LogName $logName
|
||||
# Kill stray helpers
|
||||
Write-AgentLog -FileName $logName -Message "Terminating stray helper processes"
|
||||
Get-Process python,pythonw -ErrorAction SilentlyContinue | Where-Object { $_.Path -like (Join-Path $scriptDir 'Agent\*') } | ForEach-Object {
|
||||
try { $_ | Stop-Process -Force } catch {}
|
||||
}
|
||||
# Remove venv folder
|
||||
$venvFolder = Join-Path $scriptDir 'Agent'
|
||||
Write-AgentLog -FileName $logName -Message "Removing folder: $venvFolder"
|
||||
try { Remove-Item $venvFolder -Recurse -Force -ErrorAction SilentlyContinue } catch {}
|
||||
Write-AgentLog -FileName $logName -Message "=== Removal end ==="
|
||||
}
|
||||
|
||||
function Write-ProgressStep {
|
||||
param (
|
||||
[string]$Message,
|
||||
@@ -422,7 +551,7 @@ function Ensure-BorealisScriptAgent-Service {
|
||||
[string]$ServiceScript
|
||||
)
|
||||
if (-not (Test-IsWindows)) { return }
|
||||
$serviceName = 'BorealisScriptService'
|
||||
$serviceName = 'BorealisAgent'
|
||||
$svc = Get-CimInstance -ClassName Win32_Service -Filter "Name='$serviceName'" -ErrorAction SilentlyContinue
|
||||
$needsInstall = $false
|
||||
if (-not $svc) { $needsInstall = $true }
|
||||
@@ -778,12 +907,42 @@ switch ($choice) {
|
||||
|
||||
"2" {
|
||||
$host.UI.RawUI.WindowTitle = "Borealis Agent"
|
||||
# Agent Deployment (Client / Data Collector)
|
||||
Write-Host " "
|
||||
Write-Host "Agent Menu:" -ForegroundColor Cyan
|
||||
Write-Host " 1) Install/Update Agent"
|
||||
Write-Host " 2) Repair Borealis Agent"
|
||||
Write-Host " 3) Remove Agent"
|
||||
Write-Host " 4) Launch UserSession Helper (current session)"
|
||||
Write-Host " 5) Back"
|
||||
if (-not $agentSubChoice) { $agentSubChoice = Read-Host "Select an option" }
|
||||
|
||||
switch ($agentSubChoice) {
|
||||
'2' {
|
||||
Repair-BorealisAgent
|
||||
break
|
||||
}
|
||||
'3' {
|
||||
Remove-BorealisAgent
|
||||
break
|
||||
}
|
||||
'4' {
|
||||
# Manually launch helper for the current session (optional)
|
||||
$venvPythonw = Join-Path $scriptDir 'Agent\Scripts\pythonw.exe'
|
||||
$helper = Join-Path $scriptDir 'Agent\Borealis\borealis-agent.py'
|
||||
if (-not (Test-Path $venvPythonw)) { Write-Host "pythonw.exe not found under Agent\Scripts" -ForegroundColor Yellow }
|
||||
if (-not (Test-Path $helper)) { Write-Host "Helper not found under Agent\Borealis" -ForegroundColor Yellow }
|
||||
if ((Test-Path $venvPythonw) -and (Test-Path $helper)) {
|
||||
Start-Process -FilePath $venvPythonw -ArgumentList @('-W','ignore::SyntaxWarning', $helper) -WorkingDirectory (Split-Path $helper -Parent)
|
||||
Write-Host "Launched user-session helper." -ForegroundColor Green
|
||||
}
|
||||
break
|
||||
}
|
||||
'5' { break }
|
||||
Default {
|
||||
# 1) Install/Update Agent (original behavior)
|
||||
Write-Host "Ensuring Agent Dependencies Exist..." -ForegroundColor DarkCyan
|
||||
Install_Shared_Dependencies
|
||||
Install_Agent_Dependencies
|
||||
# Confirm Python presence and update PATH
|
||||
if (-not (Test-Path $pythonExe)) {
|
||||
Write-Host "`r$($symbols.Fail) Bundled Python not found at '$pythonExe'." -ForegroundColor Red
|
||||
exit 1
|
||||
@@ -814,13 +973,8 @@ switch ($choice) {
|
||||
& $pythonForVenv -m venv $venvFolder
|
||||
}
|
||||
if (Test-Path $agentSourcePath) {
|
||||
# Remove Existing "Agent/Borealis" folder.
|
||||
Remove-Item $agentDestinationFolder -Recurse -Force -ErrorAction SilentlyContinue
|
||||
|
||||
# Create New "Agent/Borealis" folder.
|
||||
New-Item -Path $agentDestinationFolder -ItemType Directory -Force | Out-Null
|
||||
|
||||
# Agent Files and Modules
|
||||
Copy-Item "Data\Agent\borealis-agent.py" $agentDestinationFolder -Recurse
|
||||
Copy-Item "Data\Agent\agent_info.py" $agentDestinationFolder -Recurse
|
||||
Copy-Item "Data\Agent\agent_roles.py" $agentDestinationFolder -Recurse
|
||||
@@ -839,7 +993,7 @@ switch ($choice) {
|
||||
}
|
||||
}
|
||||
|
||||
Write-Host "`nConfiguring Borealis Agent (service + logon task)..." -ForegroundColor Blue
|
||||
Write-Host "`nConfiguring Borealis Agent (service)..." -ForegroundColor Blue
|
||||
Write-Host "===================================================================================="
|
||||
$deployScript = Join-Path $agentDestinationFolder 'agent_deployment.py'
|
||||
& $venvPython -W ignore::SyntaxWarning $deployScript ensure-all
|
||||
@@ -848,7 +1002,9 @@ switch ($choice) {
|
||||
Write-Host " - See logs under: $(Join-Path $scriptDir 'Logs')" -ForegroundColor Red
|
||||
exit 1
|
||||
} else {
|
||||
Write-Host "Agent setup complete. Service + logon task ensured." -ForegroundColor DarkGreen
|
||||
Write-Host "Agent setup complete. Service ensured." -ForegroundColor DarkGreen
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
@@ -159,7 +159,7 @@ def current_service_path(service_name):
|
||||
|
||||
|
||||
def ensure_script_service(paths):
|
||||
service_name = "BorealisScriptService"
|
||||
service_name = "BorealisAgent"
|
||||
log_name = "Borealis_ScriptService_Install.log"
|
||||
ensure_dirs(paths)
|
||||
log_write(paths, log_name, "[INFO] Ensuring script execution service...")
|
||||
@@ -191,14 +191,16 @@ $post = "{postinstall}"
|
||||
$pyhome = "{py_home}"
|
||||
try {{
|
||||
try {{ New-Item -ItemType Directory -Force -Path (Split-Path $log -Parent) | Out-Null }} catch {{}}
|
||||
# Remove legacy service name if present
|
||||
# Remove legacy service names if present
|
||||
try {{ sc.exe stop BorealisScriptAgent 2>$null | Out-Null }} catch {{}}
|
||||
try {{ sc.exe delete BorealisScriptAgent 2>$null | Out-Null }} catch {{}}
|
||||
try {{ sc.exe stop BorealisScriptService 2>$null | Out-Null }} catch {{}}
|
||||
try {{ sc.exe delete BorealisScriptService 2>$null | Out-Null }} catch {{}}
|
||||
if (Test-Path $post) {{ & $venv $post -install *>> "$log" }} else {{ & $venv -m pywin32_postinstall -install *>> "$log" }}
|
||||
try {{ & $venv $srv remove *>> "$log" }} catch {{}}
|
||||
& $venv $srv --startup auto install *>> "$log"
|
||||
# Ensure registry points to correct module and PY path
|
||||
reg add "HKLM\SYSTEM\CurrentControlSet\Services\{service_name}\PythonClass" /ve /t REG_SZ /d "windows_script_service.BorealisScriptAgentService" /f | Out-Null
|
||||
reg add "HKLM\SYSTEM\CurrentControlSet\Services\{service_name}\PythonClass" /ve /t REG_SZ /d "windows_script_service.BorealisAgentService" /f | Out-Null
|
||||
reg add "HKLM\SYSTEM\CurrentControlSet\Services\{service_name}\PythonPath" /ve /t REG_SZ /d "{paths['borealis_dir']}" /f | Out-Null
|
||||
reg add "HKLM\SYSTEM\CurrentControlSet\Services\{service_name}\PythonHome" /ve /t REG_SZ /d "$pyhome" /f | Out-Null
|
||||
sc.exe config {service_name} obj= LocalSystem start= auto | Out-File -FilePath "$log" -Append -Encoding UTF8
|
||||
@@ -226,14 +228,42 @@ try {{
|
||||
# Remove legacy service if it exists
|
||||
run(["sc.exe", "stop", "BorealisScriptAgent"]) # ignore rc
|
||||
run(["sc.exe", "delete", "BorealisScriptAgent"]) # ignore rc
|
||||
run(["sc.exe", "stop", "BorealisScriptService"]) # ignore rc
|
||||
run(["sc.exe", "delete", "BorealisScriptService"]) # ignore rc
|
||||
if need_install:
|
||||
run([paths["venv_python"], paths["service_script"], "remove"]) # ignore rc
|
||||
r1 = run([paths["venv_python"], paths["service_script"], "--startup", "auto", "install"], capture=True)
|
||||
log_write(paths, log_name, f"[INFO] install rc={r1.returncode} out={r1.stdout}\nerr={r1.stderr}")
|
||||
# fix registry for module import and path
|
||||
run(["reg", "add", fr"HKLM\\SYSTEM\\CurrentControlSet\\Services\\{service_name}\\PythonClass", "/ve", "/t", "REG_SZ", "/d", "windows_script_service.BorealisScriptAgentService", "/f"]) # noqa
|
||||
run(["reg", "add", fr"HKLM\\SYSTEM\\CurrentControlSet\\Services\\{service_name}\\PythonPath", "/ve", "/t", "REG_SZ", "/d", paths["borealis_dir"], "/f"]) # noqa
|
||||
run(["reg", "add", fr"HKLM\\SYSTEM\\CurrentControlSet\\Services\\{service_name}\\PythonHome", "/ve", "/t", "REG_SZ", "/d", paths["venv_root"], "/f"]) # noqa
|
||||
# fix registry for module import and runtime resolution
|
||||
# PythonHome: base interpreter home (from pyvenv.cfg 'home') so pythonservice can load pythonXY.dll
|
||||
# PythonPath: add Borealis dir and venv site-packages including pywin32 dirs
|
||||
try:
|
||||
cfg = os.path.join(paths["venv_root"], "pyvenv.cfg")
|
||||
base_home = None
|
||||
if os.path.isfile(cfg):
|
||||
with open(cfg, "r", encoding="utf-8", errors="ignore") as f:
|
||||
for line in f:
|
||||
if line.strip().lower().startswith("home ="):
|
||||
base_home = line.split("=",1)[1].strip()
|
||||
break
|
||||
if not base_home:
|
||||
# fallback to parent of venv Scripts
|
||||
base_home = os.path.dirname(os.path.dirname(paths["venv_python"]))
|
||||
except Exception:
|
||||
base_home = os.path.dirname(os.path.dirname(paths["venv_python"]))
|
||||
|
||||
site = os.path.join(paths["venv_root"], "Lib", "site-packages")
|
||||
pypath = ";".join([
|
||||
paths["borealis_dir"],
|
||||
site,
|
||||
os.path.join(site, "win32"),
|
||||
os.path.join(site, "win32", "lib"),
|
||||
os.path.join(site, "pywin32_system32"),
|
||||
])
|
||||
|
||||
run(["reg", "add", fr"HKLM\\SYSTEM\\CurrentControlSet\\Services\\{service_name}\\PythonClass", "/ve", "/t", "REG_SZ", "/d", "windows_script_service.BorealisAgentService", "/f"]) # noqa
|
||||
run(["reg", "add", fr"HKLM\\SYSTEM\\CurrentControlSet\\Services\\{service_name}\\PythonPath", "/ve", "/t", "REG_SZ", "/d", pypath, "/f"]) # noqa
|
||||
run(["reg", "add", fr"HKLM\\SYSTEM\\CurrentControlSet\\Services\\{service_name}\\PythonHome", "/ve", "/t", "REG_SZ", "/d", base_home, "/f"]) # noqa
|
||||
run(["sc.exe", "config", service_name, "obj=", "LocalSystem"]) # ensure LocalSystem
|
||||
run(["sc.exe", "start", service_name])
|
||||
# quick validate
|
||||
@@ -332,8 +362,8 @@ def ensure_all():
|
||||
paths = project_paths()
|
||||
ensure_dirs(paths)
|
||||
ok_svc = ensure_script_service(paths)
|
||||
ok_task = ensure_user_logon_task(paths)
|
||||
return 0 if (ok_svc and ok_task) else 1
|
||||
# Service now launches per-session helper; scheduled task is not required.
|
||||
return 0 if ok_svc else 1
|
||||
|
||||
|
||||
def main(argv):
|
||||
@@ -350,7 +380,7 @@ def main(argv):
|
||||
if cmd == "service-install":
|
||||
return 0 if ensure_script_service(paths) else 1
|
||||
if cmd == "service-remove":
|
||||
name = "BorealisScriptService"
|
||||
name = "BorealisAgent"
|
||||
if not is_admin():
|
||||
ps = f"try {{ sc.exe stop {name} }} catch {{}}; try {{ sc.exe delete {name} }} catch {{}}"
|
||||
return run_elevated_powershell(paths, ps, "Borealis_ScriptService_Remove.log")
|
||||
|
@@ -996,17 +996,16 @@ async def on_quick_job_run(payload):
|
||||
script_type = (payload.get('script_type') or '').lower()
|
||||
run_mode = (payload.get('run_mode') or 'current_user').lower()
|
||||
content = payload.get('script_content') or ''
|
||||
# Only handle non-SYSTEM runs here; SYSTEM runs are handled by the LocalSystem service agent
|
||||
if run_mode == 'system':
|
||||
# Optionally, could emit a status indicating delegation; for now, just ignore to avoid overlap
|
||||
return
|
||||
if script_type != 'powershell':
|
||||
await sio.emit('quick_job_result', { 'job_id': job_id, 'status': 'Failed', 'stdout': '', 'stderr': f"Unsupported type: {script_type}" })
|
||||
return
|
||||
path = _write_temp_script(content, '.ps1')
|
||||
rc = 0; out = ''; err = ''
|
||||
if run_mode == 'system':
|
||||
if not _is_admin_windows():
|
||||
rc, out, err = -1, '', 'Agent is not elevated. SYSTEM execution requires running the agent as Administrator or service.'
|
||||
else:
|
||||
rc, out, err = await _run_powershell_as_system(path)
|
||||
elif run_mode == 'admin':
|
||||
if run_mode == 'admin':
|
||||
# Admin credentialed runs are disabled in current design
|
||||
rc, out, err = -1, '', 'Admin credentialed runs are disabled; use SYSTEM (service) or Current User.'
|
||||
else:
|
||||
|
@@ -8,6 +8,8 @@ import subprocess
|
||||
import tempfile
|
||||
|
||||
import socketio
|
||||
import platform
|
||||
import time
|
||||
|
||||
|
||||
def get_project_root():
|
||||
@@ -72,6 +74,10 @@ async def main():
|
||||
target = (payload.get('target_hostname') or '').strip().lower()
|
||||
if target and target != hostname.lower():
|
||||
return
|
||||
run_mode = (payload.get('run_mode') or 'current_user').lower()
|
||||
# Only the SYSTEM service handles system-mode jobs; ignore others
|
||||
if run_mode != 'system':
|
||||
return
|
||||
job_id = payload.get('job_id')
|
||||
script_type = (payload.get('script_type') or '').lower()
|
||||
content = payload.get('script_content') or ''
|
||||
@@ -106,11 +112,33 @@ async def main():
|
||||
async def disconnect():
|
||||
print("[ScriptAgent] Disconnected")
|
||||
|
||||
async def heartbeat_loop():
|
||||
# Minimal heartbeat so device appears online even without a user helper
|
||||
while True:
|
||||
try:
|
||||
await sio.emit("agent_heartbeat", {
|
||||
"agent_id": f"{hostname}-script",
|
||||
"hostname": hostname,
|
||||
"agent_operating_system": f"{platform.system()} {platform.release()} (Service)",
|
||||
"last_seen": int(time.time())
|
||||
})
|
||||
except Exception:
|
||||
pass
|
||||
await asyncio.sleep(30)
|
||||
|
||||
url = get_server_url()
|
||||
while True:
|
||||
try:
|
||||
await sio.connect(url, transports=['websocket'])
|
||||
# Heartbeat while connected
|
||||
hb = asyncio.create_task(heartbeat_loop())
|
||||
try:
|
||||
await sio.wait()
|
||||
finally:
|
||||
try:
|
||||
hb.cancel()
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as e:
|
||||
print(f"[ScriptAgent] reconnect in 5s: {e}")
|
||||
await asyncio.sleep(5)
|
||||
|
@@ -37,7 +37,7 @@ class TrayApp(QtWidgets.QSystemTrayIcon):
|
||||
self.action_show_console = self.menu.addAction('Switch to Foreground Mode')
|
||||
self.action_hide_console = self.menu.addAction('Switch to Background Mode')
|
||||
self.action_restart = self.menu.addAction('Restart Agent')
|
||||
self.action_restart_service = self.menu.addAction('Restart Script Execution Service')
|
||||
self.action_restart_service = self.menu.addAction('Restart Borealis Agent Service')
|
||||
self.menu.addSeparator()
|
||||
self.action_quit = self.menu.addAction('Quit Agent and Tray')
|
||||
|
||||
@@ -103,7 +103,7 @@ class TrayApp(QtWidgets.QSystemTrayIcon):
|
||||
|
||||
def restart_script_service(self):
|
||||
# Try direct stop/start; if fails (likely due to permissions), attempt elevation via PowerShell
|
||||
service_name = 'BorealisScriptService'
|
||||
service_name = 'BorealisAgent'
|
||||
try:
|
||||
# Stop service
|
||||
subprocess.run(["sc.exe", "stop", service_name], check=False, capture_output=True)
|
||||
|
@@ -6,17 +6,31 @@ import subprocess
|
||||
import os
|
||||
import sys
|
||||
import datetime
|
||||
import threading
|
||||
|
||||
# Session/process helpers for per-user helper launch
|
||||
try:
|
||||
import win32ts
|
||||
import win32con
|
||||
import win32process
|
||||
import win32security
|
||||
import win32profile
|
||||
import win32api
|
||||
except Exception:
|
||||
win32ts = None
|
||||
|
||||
|
||||
class BorealisScriptAgentService(win32serviceutil.ServiceFramework):
|
||||
_svc_name_ = "BorealisScriptService"
|
||||
_svc_display_name_ = "Borealis Script Execution Service"
|
||||
_svc_description_ = "Executes automation scripts (PowerShell, etc.) as LocalSystem and bridges to Borealis Server."
|
||||
class BorealisAgentService(win32serviceutil.ServiceFramework):
|
||||
_svc_name_ = "BorealisAgent"
|
||||
_svc_display_name_ = "Borealis Agent"
|
||||
_svc_description_ = "Background agent for data collection and remote script execution."
|
||||
|
||||
def __init__(self, args):
|
||||
win32serviceutil.ServiceFramework.__init__(self, args)
|
||||
self.hWaitStop = win32event.CreateEvent(None, 0, 0, None)
|
||||
self.proc = None
|
||||
self.user_helpers = {}
|
||||
self.helpers_thread = None
|
||||
try:
|
||||
self._log("Service initialized")
|
||||
except Exception:
|
||||
@@ -36,6 +50,17 @@ class BorealisScriptAgentService(win32serviceutil.ServiceFramework):
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
# Stop user helpers
|
||||
try:
|
||||
for sid, h in list(self.user_helpers.items()):
|
||||
try:
|
||||
hp = h.get('hProcess')
|
||||
if hp:
|
||||
win32api.TerminateProcess(hp, 0)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
win32event.SetEvent(self.hWaitStop)
|
||||
|
||||
def SvcDoRun(self):
|
||||
@@ -59,43 +84,191 @@ class BorealisScriptAgentService(win32serviceutil.ServiceFramework):
|
||||
def main(self):
|
||||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
agent_py = os.path.join(script_dir, 'script_agent.py')
|
||||
python = sys.executable
|
||||
|
||||
def _venv_python():
|
||||
try:
|
||||
exe_dir = os.path.dirname(sys.executable)
|
||||
py = os.path.join(exe_dir, 'python.exe')
|
||||
if os.path.isfile(py):
|
||||
return py
|
||||
except Exception:
|
||||
pass
|
||||
return sys.executable
|
||||
|
||||
def _start_script_agent():
|
||||
python = _venv_python()
|
||||
try:
|
||||
self._log(f"Launching script_agent via {python}")
|
||||
self.proc = subprocess.Popen(
|
||||
return subprocess.Popen(
|
||||
[python, '-W', 'ignore::SyntaxWarning', agent_py],
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
creationflags=0x08000000 if os.name == 'nt' else 0
|
||||
)
|
||||
self._log("script_agent started")
|
||||
except Exception:
|
||||
self.proc = None
|
||||
except Exception as e:
|
||||
try:
|
||||
self._log("Failed to start script_agent")
|
||||
self._log(f"Failed to start script_agent: {e}")
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
# Launch the system script agent (and keep it alive)
|
||||
self.proc = _start_script_agent()
|
||||
|
||||
# Start per-user helper manager in a background thread
|
||||
if win32ts is not None:
|
||||
try:
|
||||
self.helpers_thread = threading.Thread(target=self._manage_user_helpers_loop, daemon=True)
|
||||
self.helpers_thread.start()
|
||||
except Exception:
|
||||
try:
|
||||
self._log("Failed to start user helper manager thread")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Wait until stop or child exits
|
||||
# Monitor stop event and child process; restart child if it exits
|
||||
while True:
|
||||
rc = win32event.WaitForSingleObject(self.hWaitStop, 1000)
|
||||
rc = win32event.WaitForSingleObject(self.hWaitStop, 2000)
|
||||
if rc == win32event.WAIT_OBJECT_0:
|
||||
break
|
||||
if self.proc and self.proc.poll() is not None:
|
||||
break
|
||||
try:
|
||||
if not self.proc or (self.proc and self.proc.poll() is not None):
|
||||
# child exited; attempt restart
|
||||
self._log("script_agent exited; attempting restart")
|
||||
self.proc = _start_script_agent()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _log(self, msg: str):
|
||||
try:
|
||||
root = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, os.pardir))
|
||||
logs = os.path.join(root, 'Logs')
|
||||
os.makedirs(logs, exist_ok=True)
|
||||
p = os.path.join(logs, 'ScriptService_Runtime.log')
|
||||
p = os.path.join(logs, 'AgentService_Runtime.log')
|
||||
ts = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
|
||||
with open(p, 'a', encoding='utf-8') as f:
|
||||
f.write(f"{ts} {msg}\n")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ---------------------- Per-Session User Helper Management ----------------------
|
||||
def _enable_privileges(self):
|
||||
try:
|
||||
hProc = win32api.GetCurrentProcess()
|
||||
hTok = win32security.OpenProcessToken(
|
||||
hProc,
|
||||
win32con.TOKEN_ADJUST_PRIVILEGES | win32con.TOKEN_QUERY,
|
||||
)
|
||||
for name in [
|
||||
win32security.SE_ASSIGNPRIMARYTOKEN_NAME,
|
||||
win32security.SE_INCREASE_QUOTA_NAME,
|
||||
win32security.SE_TCB_NAME,
|
||||
win32security.SE_BACKUP_NAME,
|
||||
win32security.SE_RESTORE_NAME,
|
||||
]:
|
||||
try:
|
||||
luid = win32security.LookupPrivilegeValue(None, name)
|
||||
win32security.AdjustTokenPrivileges(
|
||||
hTok, False, [(luid, win32con.SE_PRIVILEGE_ENABLED)]
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _active_session_ids(self):
|
||||
ids = []
|
||||
try:
|
||||
sessions = win32ts.WTSEnumerateSessions(None, 1, 0)
|
||||
for s in sessions:
|
||||
# tuple: (SessionId, WinStationName, State)
|
||||
sess_id, _, state = s
|
||||
if state == win32ts.WTSActive:
|
||||
ids.append(sess_id)
|
||||
except Exception:
|
||||
pass
|
||||
return ids
|
||||
|
||||
def _launch_helper_in_session(self, session_id):
|
||||
try:
|
||||
self._enable_privileges()
|
||||
# User token for session
|
||||
hUser = win32ts.WTSQueryUserToken(session_id)
|
||||
# Duplicate to primary
|
||||
primary = win32security.DuplicateTokenEx(
|
||||
hUser,
|
||||
win32con.MAXIMUM_ALLOWED,
|
||||
win32security.SECURITY_ATTRIBUTES(),
|
||||
win32security.SecurityImpersonation,
|
||||
win32con.TOKEN_PRIMARY,
|
||||
)
|
||||
env = win32profile.CreateEnvironmentBlock(primary, True)
|
||||
startup = win32process.STARTUPINFO()
|
||||
startup.lpDesktop = "winsta0\\default"
|
||||
|
||||
# Compute pythonw + helper script
|
||||
venv_dir = os.path.dirname(sys.executable) # e.g., Agent
|
||||
pyw = os.path.join(venv_dir, 'pythonw.exe')
|
||||
agent_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
helper = os.path.join(agent_dir, 'borealis-agent.py')
|
||||
cmd = f'"{pyw}" -W ignore::SyntaxWarning "{helper}"'
|
||||
|
||||
flags = getattr(win32con, 'CREATE_NEW_PROCESS_GROUP', 0)
|
||||
flags |= getattr(win32con, 'CREATE_UNICODE_ENVIRONMENT', 0x00000400)
|
||||
proc_tuple = win32process.CreateProcessAsUser(
|
||||
primary,
|
||||
None,
|
||||
cmd,
|
||||
None,
|
||||
None,
|
||||
False,
|
||||
flags,
|
||||
env,
|
||||
agent_dir,
|
||||
startup,
|
||||
)
|
||||
# proc_tuple: (hProcess, hThread, dwProcessId, dwThreadId)
|
||||
self.user_helpers[session_id] = {
|
||||
'hProcess': proc_tuple[0],
|
||||
'hThread': proc_tuple[1],
|
||||
'pid': proc_tuple[2],
|
||||
'started': datetime.datetime.now(),
|
||||
}
|
||||
self._log(f"Started user helper in session {session_id}")
|
||||
return True
|
||||
except Exception as e:
|
||||
try:
|
||||
self._log(f"Failed to start helper in session {session_id}: {e}")
|
||||
except Exception:
|
||||
pass
|
||||
return False
|
||||
|
||||
def _manage_user_helpers_loop(self):
|
||||
# Periodically ensure one user helper per active session
|
||||
while True:
|
||||
rc = win32event.WaitForSingleObject(self.hWaitStop, 2000)
|
||||
if rc == win32event.WAIT_OBJECT_0:
|
||||
break
|
||||
try:
|
||||
active = set(self._active_session_ids())
|
||||
# Start missing
|
||||
for sid in active:
|
||||
if sid not in self.user_helpers:
|
||||
self._launch_helper_in_session(sid)
|
||||
# Cleanup ended
|
||||
for sid in list(self.user_helpers.keys()):
|
||||
if sid not in active:
|
||||
info = self.user_helpers.pop(sid, None)
|
||||
try:
|
||||
hp = info and info.get('hProcess')
|
||||
if hp:
|
||||
win32api.TerminateProcess(hp, 0)
|
||||
except Exception:
|
||||
pass
|
||||
self._log(f"Cleaned helper for session {sid}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
win32serviceutil.HandleCommandLine(BorealisScriptAgentService)
|
||||
win32serviceutil.HandleCommandLine(BorealisAgentService)
|
||||
|
Reference in New Issue
Block a user