begin persistence/attachment overhaul
This commit is contained in:
Nathan
@@ -32,6 +32,196 @@ $workers = @(
|
||||
@{ ID = 6; Name = "i9-13ks"; SSHArgs = "-t -p 22146 i9-13ks"; Enabled = $true }
|
||||
)
|
||||
|
||||
$global:UnifiedWorkerDataRoot = 'C:\ProgramData\UnifiedWorkers'
|
||||
$script:ControllerScriptBase64 = $null
|
||||
$script:AttachHelperScriptBase64 = $null
|
||||
|
||||
function Get-FileBase64 {
|
||||
param([string]$Path)
|
||||
[Convert]::ToBase64String([IO.File]::ReadAllBytes($Path))
|
||||
}
|
||||
|
||||
function Get-ControllerScriptBase64 {
|
||||
if (-not $script:ControllerScriptBase64) {
|
||||
$controllerPath = Join-Path $PSScriptRoot 'remote_worker_controller.ps1'
|
||||
$script:ControllerScriptBase64 = Get-FileBase64 -Path $controllerPath
|
||||
}
|
||||
return $script:ControllerScriptBase64
|
||||
}
|
||||
|
||||
function Get-AttachHelperScriptBase64 {
|
||||
if (-not $script:AttachHelperScriptBase64) {
|
||||
$helperPath = Join-Path $PSScriptRoot 'remote_worker_attach.ps1'
|
||||
$script:AttachHelperScriptBase64 = Get-FileBase64 -Path $helperPath
|
||||
}
|
||||
return $script:AttachHelperScriptBase64
|
||||
}
|
||||
|
||||
function ConvertTo-Base64Unicode {
|
||||
param([string]$Content)
|
||||
[Convert]::ToBase64String([Text.Encoding]::Unicode.GetBytes($Content))
|
||||
}
|
||||
|
||||
function Get-WorkerSshArgs {
|
||||
param([string]$RawArgs)
|
||||
if ([string]::IsNullOrWhiteSpace($RawArgs)) {
|
||||
return @()
|
||||
}
|
||||
return $RawArgs -split '\s+' | Where-Object { $_.Trim().Length -gt 0 }
|
||||
}
|
||||
|
||||
function Invoke-RemotePowerShell {
|
||||
param(
|
||||
[object]$Worker,
|
||||
[string]$Script
|
||||
)
|
||||
|
||||
$bytes = [Text.Encoding]::Unicode.GetBytes($Script)
|
||||
$encoded = [Convert]::ToBase64String($bytes)
|
||||
$sshArgs = @('-o','ServerAliveInterval=60','-o','ServerAliveCountMax=30')
|
||||
$sshArgs += (Get-WorkerSshArgs -RawArgs $Worker.SSHArgs)
|
||||
$sshArgs += "powershell -NoLogo -NoProfile -ExecutionPolicy Bypass -EncodedCommand $encoded"
|
||||
|
||||
& ssh @sshArgs
|
||||
return $LASTEXITCODE
|
||||
}
|
||||
|
||||
function Ensure-ControllerDeployed {
|
||||
param([object]$Worker)
|
||||
$controllerBase64 = Get-ControllerScriptBase64
|
||||
$script = @'
|
||||
`$dataRoot = '{0}'
|
||||
New-Item -ItemType Directory -Path `$dataRoot -Force | Out-Null
|
||||
`$controllerPath = Join-Path `$dataRoot 'controller.ps1'
|
||||
[IO.File]::WriteAllBytes(`$controllerPath, [Convert]::FromBase64String('{1}'))
|
||||
'@ -f $global:UnifiedWorkerDataRoot, $controllerBase64
|
||||
Invoke-RemotePowerShell -Worker $Worker -Script $script | Out-Null
|
||||
}
|
||||
|
||||
function Ensure-AttachHelperDeployed {
|
||||
param([object]$Worker)
|
||||
$helperBase64 = Get-AttachHelperScriptBase64
|
||||
$script = @'
|
||||
`$dataRoot = '{0}'
|
||||
New-Item -ItemType Directory -Path `$dataRoot -Force | Out-Null
|
||||
`$attachPath = Join-Path `$dataRoot 'attach-helper.ps1'
|
||||
[IO.File]::WriteAllBytes(`$attachPath, [Convert]::FromBase64String('{1}'))
|
||||
'@ -f $global:UnifiedWorkerDataRoot, $helperBase64
|
||||
Invoke-RemotePowerShell -Worker $Worker -Script $script | Out-Null
|
||||
}
|
||||
|
||||
function Get-RemoteWorkerPaths {
|
||||
param([string]$WorkerType, [string]$WorkerName)
|
||||
$base = Join-Path $global:UnifiedWorkerDataRoot $WorkerType
|
||||
$instance = Join-Path $base $WorkerName
|
||||
[pscustomobject]@{
|
||||
InstanceRoot = $instance
|
||||
MetadataPath = Join-Path $instance 'state\worker-info.json'
|
||||
CommandPath = Join-Path $instance 'state\commands.txt'
|
||||
LogPath = Join-Path $instance 'logs\worker.log'
|
||||
}
|
||||
}
|
||||
|
||||
function Get-EnsureWorkerScript {
|
||||
param(
|
||||
[string]$WorkerName,
|
||||
[string]$WorkerType,
|
||||
[string]$PayloadBase64
|
||||
)
|
||||
|
||||
$paths = Get-RemoteWorkerPaths -WorkerType $WorkerType -WorkerName $WorkerName
|
||||
$controllerPath = Join-Path $global:UnifiedWorkerDataRoot 'controller.ps1'
|
||||
|
||||
@'
|
||||
`$controllerPath = '{0}'
|
||||
`$metaPath = '{1}'
|
||||
`$workerName = '{2}'
|
||||
`$workerType = '{3}'
|
||||
`$payloadBase64 = '{4}'
|
||||
|
||||
if (-not (Test-Path `$controllerPath)) {
|
||||
throw "Controller missing at `$controllerPath"
|
||||
}
|
||||
|
||||
`$shouldStart = `$true
|
||||
if (Test-Path `$metaPath) {
|
||||
try {
|
||||
`$meta = Get-Content `$metaPath -Raw | ConvertFrom-Json
|
||||
if (`$meta.Status -eq 'running' -and `$meta.WorkerPid) {
|
||||
if (Get-Process -Id `$meta.WorkerPid -ErrorAction SilentlyContinue) {
|
||||
Write-Host "Worker `$workerName already running (PID `$($meta.WorkerPid))."
|
||||
`$shouldStart = `$false
|
||||
}
|
||||
}
|
||||
} catch {
|
||||
Write-Host "Failed to read metadata. Controller will restart worker." -ForegroundColor Yellow
|
||||
}
|
||||
}
|
||||
|
||||
if (`$shouldStart) {
|
||||
`$psExe = (Get-Command pwsh -ErrorAction SilentlyContinue)?.Source
|
||||
if (-not `$psExe) {
|
||||
`$psExe = (Get-Command powershell -ErrorAction Stop).Source
|
||||
}
|
||||
|
||||
`$controllerArgs = @(
|
||||
'-NoLogo','-NoProfile','-ExecutionPolicy','Bypass',
|
||||
'-File',"`$controllerPath",
|
||||
'-WorkerName',"`$workerName",
|
||||
'-WorkerType',"`$workerType",
|
||||
'-PayloadBase64',"`$payloadBase64"
|
||||
)
|
||||
|
||||
Start-Process -FilePath `$psExe -ArgumentList `$controllerArgs -WindowStyle Hidden | Out-Null
|
||||
Write-Host "Worker `$workerName started under controller." -ForegroundColor Green
|
||||
}
|
||||
'@ -f $controllerPath, $paths.MetadataPath, $WorkerName, $WorkerType, $PayloadBase64
|
||||
}
|
||||
|
||||
function Ensure-PersistentWorker {
|
||||
param(
|
||||
[object]$Worker,
|
||||
[string]$WorkerType,
|
||||
[string]$PayloadScript
|
||||
)
|
||||
|
||||
Ensure-ControllerDeployed -Worker $Worker
|
||||
$payloadBase64 = ConvertTo-Base64Unicode -Content $PayloadScript
|
||||
$ensureScript = Get-EnsureWorkerScript -WorkerName $Worker.Name -WorkerType $WorkerType -PayloadBase64 $payloadBase64
|
||||
Invoke-RemotePowerShell -Worker $Worker -Script $ensureScript | Out-Null
|
||||
}
|
||||
|
||||
function Invoke-WorkerAttach {
|
||||
param(
|
||||
[object]$Worker,
|
||||
[string]$WorkerType,
|
||||
[switch]$CommandOnly,
|
||||
[string]$Command
|
||||
)
|
||||
|
||||
Ensure-AttachHelperDeployed -Worker $Worker
|
||||
$paramsBlock = @("-WorkerName","$($Worker.Name)","-WorkerType","$WorkerType")
|
||||
if ($CommandOnly) {
|
||||
$paramsBlock += "-CommandOnly"
|
||||
}
|
||||
if ($Command) {
|
||||
$paramsBlock += "-Command"
|
||||
$paramsBlock += $Command
|
||||
}
|
||||
|
||||
$paramsLiteral = ($paramsBlock | ForEach-Object { "'$_'" }) -join ', '
|
||||
$script = @'
|
||||
`$helperPath = Join-Path '{0}' 'attach-helper.ps1'
|
||||
if (-not (Test-Path `$helperPath)) {
|
||||
throw "Attach helper missing at `$helperPath"
|
||||
}
|
||||
`$arguments = @({1})
|
||||
& `$helperPath @arguments
|
||||
'@ -f $global:UnifiedWorkerDataRoot, $paramsLiteral
|
||||
|
||||
Invoke-RemotePowerShell -Worker $Worker -Script $script | Out-Null
|
||||
}
|
||||
|
||||
function Get-RemoteSheepItCommand {
|
||||
param(
|
||||
[string]$RenderKey,
|
||||
@@ -135,6 +325,40 @@ try {
|
||||
throw
|
||||
}
|
||||
}
|
||||
|
||||
function Ensure-SheepItWorkerController {
|
||||
param([object]$Worker)
|
||||
Initialize-SheepItCredentials
|
||||
$payloadScript = Get-RemoteSheepItCommand -RenderKey $script:SheepItRenderKey -Username $script:SheepItUsername
|
||||
Ensure-PersistentWorker -Worker $Worker -WorkerType 'sheepit' -PayloadScript $payloadScript
|
||||
}
|
||||
|
||||
function Start-SheepItWorker {
|
||||
param([object]$Worker)
|
||||
Ensure-SheepItWorkerController -Worker $Worker
|
||||
Invoke-WorkerAttach -Worker $Worker -WorkerType 'sheepit'
|
||||
}
|
||||
|
||||
function Start-AllSheepIt {
|
||||
Initialize-SheepItCredentials
|
||||
foreach ($worker in $workers | Where-Object { $_.Enabled }) {
|
||||
Ensure-SheepItWorkerController -Worker $worker
|
||||
}
|
||||
Write-Host "All enabled SheepIt workers ensured running under controllers." -ForegroundColor Green
|
||||
Write-Host "Use the attach option to monitor any worker." -ForegroundColor Cyan
|
||||
Read-Host "Press Enter to continue" | Out-Null
|
||||
}
|
||||
|
||||
function Send-SheepItCommandAll {
|
||||
param([string]$CommandText)
|
||||
|
||||
foreach ($worker in $workers | Where-Object { $_.Enabled }) {
|
||||
Write-Host "[$($worker.Name)] Sending '$CommandText'..." -ForegroundColor Gray
|
||||
Invoke-WorkerAttach -Worker $worker -WorkerType 'sheepit' -CommandOnly -Command $CommandText
|
||||
}
|
||||
Write-Host "Command '$CommandText' dispatched to all enabled workers." -ForegroundColor Green
|
||||
Read-Host "Press Enter to continue" | Out-Null
|
||||
}
|
||||
catch {
|
||||
Write-Host ('Error: {0}' -f `$_.Exception.Message) -ForegroundColor Red
|
||||
Write-Host ('Stack trace: {0}' -f `$_.ScriptStackTrace) -ForegroundColor DarkRed
|
||||
@@ -142,127 +366,7 @@ catch {
|
||||
"@
|
||||
}
|
||||
|
||||
function New-SheepItSessionScript {
|
||||
param(
|
||||
[object]$Worker,
|
||||
[string]$RemoteCommand
|
||||
)
|
||||
|
||||
$rawArgs = @($Worker.SSHArgs -split '\s+' | Where-Object { $_ -and $_.Trim().Length -gt 0 })
|
||||
$sshArgsJson = if ($rawArgs) { ($rawArgs | ConvertTo-Json -Compress) } else { '[]' }
|
||||
$targetHost = ($rawArgs | Where-Object { $_ -notmatch '^-{1,2}' } | Select-Object -Last 1)
|
||||
$hostLiteral = if ($targetHost) { "'" + ($targetHost -replace "'", "''") + "'" } else { '$null' }
|
||||
$portValue = $null
|
||||
for ($i = 0; $i -lt $rawArgs.Count; $i++) {
|
||||
if ($rawArgs[$i] -eq '-p' -and $i -lt ($rawArgs.Count - 1)) {
|
||||
$portValue = $rawArgs[$i + 1]
|
||||
break
|
||||
}
|
||||
}
|
||||
$portLiteral = if ($portValue) { "'" + ($portValue -replace "'", "''") + "'" } else { '$null' }
|
||||
$remoteScriptBase64 = [Convert]::ToBase64String([System.Text.Encoding]::Unicode.GetBytes($RemoteCommand))
|
||||
|
||||
$scriptContent = @"
|
||||
Write-Host 'Connecting to $($Worker.Name)...' -ForegroundColor Cyan
|
||||
`$sshArgs = ConvertFrom-Json '$sshArgsJson'
|
||||
`$targetHost = $hostLiteral
|
||||
`$port = $portLiteral
|
||||
`$scriptBase64 = '$remoteScriptBase64'
|
||||
|
||||
if (-not `$targetHost) {
|
||||
Write-Host "Unable to determine SSH host for $($Worker.Name)." -ForegroundColor Red
|
||||
return
|
||||
}
|
||||
|
||||
`$localTemp = [System.IO.Path]::GetTempFileName() + '.ps1'
|
||||
`$remoteScript = [System.Text.Encoding]::Unicode.GetString([Convert]::FromBase64String(`$scriptBase64))
|
||||
Set-Content -Path `$localTemp -Value `$remoteScript -Encoding UTF8
|
||||
|
||||
`$remoteFile = "sheepit-$([System.Guid]::NewGuid().ToString()).ps1"
|
||||
|
||||
`$scpArgs = @('-o','ServerAliveInterval=60','-o','ServerAliveCountMax=30')
|
||||
if (`$port) {
|
||||
`$scpArgs += '-P'
|
||||
`$scpArgs += `$port
|
||||
}
|
||||
`$scpArgs += `$localTemp
|
||||
`$scpArgs += ("${targetHost}:`$remoteFile")
|
||||
|
||||
Write-Host "Transferring SheepIt script to `$targetHost..." -ForegroundColor Gray
|
||||
& scp @scpArgs
|
||||
if (`$LASTEXITCODE -ne 0) {
|
||||
Write-Host "Failed to copy script to `$targetHost." -ForegroundColor Red
|
||||
Remove-Item `$localTemp -ErrorAction SilentlyContinue
|
||||
return
|
||||
}
|
||||
|
||||
`$sshBase = @('-o','ServerAliveInterval=60','-o','ServerAliveCountMax=30') + `$sshArgs
|
||||
`$execArgs = `$sshBase + @('powershell','-NoLogo','-NoProfile','-ExecutionPolicy','Bypass','-File',"`$remoteFile")
|
||||
|
||||
Write-Host "Starting SheepIt client on `$targetHost..." -ForegroundColor Cyan
|
||||
& ssh @execArgs
|
||||
|
||||
Write-Host "Cleaning up remote script..." -ForegroundColor Gray
|
||||
`$cleanupArgs = `$sshBase + @('powershell','-NoLogo','-NoProfile','-ExecutionPolicy','Bypass','-Command',"`$ErrorActionPreference='SilentlyContinue'; Remove-Item -LiteralPath ``"$remoteFile``" -ErrorAction SilentlyContinue")
|
||||
& ssh @cleanupArgs | Out-Null
|
||||
|
||||
Remove-Item `$localTemp -ErrorAction SilentlyContinue
|
||||
Write-Host "`nSSH session ended." -ForegroundColor Yellow
|
||||
Read-Host "Press Enter to close"
|
||||
"@
|
||||
|
||||
return $scriptContent
|
||||
}
|
||||
|
||||
function Start-SheepItTab {
|
||||
param(
|
||||
[string]$Title,
|
||||
[string]$Content
|
||||
)
|
||||
|
||||
$tempScript = [System.IO.Path]::GetTempFileName() + '.ps1'
|
||||
Set-Content -Path $tempScript -Value $Content -Encoding UTF8
|
||||
|
||||
if (Get-Command wt.exe -ErrorAction SilentlyContinue) {
|
||||
Start-Process wt.exe -ArgumentList "-w 0 new-tab --title `"$Title`" powershell -NoLogo -NoProfile -ExecutionPolicy Bypass -File `"$tempScript`""
|
||||
} else {
|
||||
Start-Process powershell -ArgumentList "-NoLogo -NoProfile -ExecutionPolicy Bypass -File `"$tempScript`""
|
||||
}
|
||||
}
|
||||
|
||||
function Start-SheepItWorker {
|
||||
param([object]$Worker)
|
||||
|
||||
if (-not $Worker.Enabled) {
|
||||
Write-Host "$($Worker.Name) is not enabled for SheepIt." -ForegroundColor Yellow
|
||||
return
|
||||
}
|
||||
|
||||
Initialize-SheepItCredentials
|
||||
|
||||
$remoteCommand = Get-RemoteSheepItCommand -RenderKey $script:SheepItRenderKey -Username $script:SheepItUsername
|
||||
$sessionScript = New-SheepItSessionScript -Worker $Worker -RemoteCommand $remoteCommand
|
||||
$title = "$($Worker.Name) - SheepIt"
|
||||
|
||||
Start-SheepItTab -Title $title -Content $sessionScript
|
||||
Write-Host "Opened SheepIt session for $($Worker.Name) in a new terminal tab." -ForegroundColor Green
|
||||
}
|
||||
|
||||
function Start-AllSheepIt {
|
||||
Initialize-SheepItCredentials
|
||||
|
||||
$targets = $workers | Where-Object { $_.Enabled }
|
||||
|
||||
if (-not $targets) {
|
||||
Write-Host "No systems are ready for SheepIt." -ForegroundColor Yellow
|
||||
return
|
||||
}
|
||||
|
||||
foreach ($worker in $targets) {
|
||||
Start-SheepItWorker -Worker $worker
|
||||
Start-Sleep -Milliseconds 200
|
||||
}
|
||||
}
|
||||
|
||||
function Select-SheepItWorker {
|
||||
while ($true) {
|
||||
@@ -308,20 +412,22 @@ Initialize-SheepItCredentials
|
||||
while ($true) {
|
||||
Show-Header
|
||||
Write-Host "Main Menu:" -ForegroundColor Magenta
|
||||
Write-Host "1. Launch SheepIt on a single system" -ForegroundColor Yellow
|
||||
Write-Host "2. Launch SheepIt on all ready systems" -ForegroundColor Yellow
|
||||
Write-Host "3. Exit" -ForegroundColor Yellow
|
||||
Write-Host "1. Launch/Attach SheepIt on a single system" -ForegroundColor Yellow
|
||||
Write-Host "2. Ensure all ready systems are running" -ForegroundColor Yellow
|
||||
Write-Host "3. Pause all workers" -ForegroundColor Yellow
|
||||
Write-Host "4. Resume all workers" -ForegroundColor Yellow
|
||||
Write-Host "5. Quit all workers" -ForegroundColor Yellow
|
||||
Write-Host "6. Exit" -ForegroundColor Yellow
|
||||
|
||||
$choice = Read-Host "Select option (1-3)"
|
||||
|
||||
switch ($choice) {
|
||||
'1' { Select-SheepItWorker }
|
||||
'2' {
|
||||
Start-AllSheepIt
|
||||
Write-Host
|
||||
Read-Host "Press Enter to return to the main menu" | Out-Null
|
||||
}
|
||||
'3' { break }
|
||||
'2' { Start-AllSheepIt }
|
||||
'3' { Send-SheepItCommandAll -CommandText 'pause' }
|
||||
'4' { Send-SheepItCommandAll -CommandText 'resume' }
|
||||
'5' { Send-SheepItCommandAll -CommandText 'quit' }
|
||||
'6' { break }
|
||||
default {
|
||||
Write-Host "Invalid selection." -ForegroundColor Red
|
||||
Start-Sleep -Seconds 1
|
||||
|
||||
Reference in New Issue
Block a user