begin persistence/attachment overhaul
This commit is contained in:
Nathan
195
remote_worker_controller.ps1
Normal file
195
remote_worker_controller.ps1
Normal file
@@ -0,0 +1,195 @@
|
||||
param(
|
||||
[Parameter(Mandatory = $true)]
|
||||
[string]$WorkerName,
|
||||
|
||||
[Parameter(Mandatory = $true)]
|
||||
[string]$WorkerType,
|
||||
|
||||
[Parameter(Mandatory = $true)]
|
||||
[string]$PayloadBase64,
|
||||
|
||||
[string]$DataRoot = "$env:ProgramData\UnifiedWorkers",
|
||||
|
||||
[int]$MaxRestarts = 5,
|
||||
|
||||
[int]$RestartDelaySeconds = 10
|
||||
)
|
||||
|
||||
$ErrorActionPreference = 'Stop'
|
||||
|
||||
# region Path setup
|
||||
$workerRoot = Join-Path -Path $DataRoot -ChildPath $WorkerType
|
||||
$instanceRoot = Join-Path -Path $workerRoot -ChildPath $WorkerName
|
||||
New-Item -ItemType Directory -Path $instanceRoot -Force | Out-Null
|
||||
|
||||
$logsRoot = Join-Path -Path $instanceRoot -ChildPath 'logs'
|
||||
$stateRoot = Join-Path -Path $instanceRoot -ChildPath 'state'
|
||||
New-Item -ItemType Directory -Path $logsRoot -Force | Out-Null
|
||||
New-Item -ItemType Directory -Path $stateRoot -Force | Out-Null
|
||||
|
||||
$logPath = Join-Path -Path $logsRoot -ChildPath 'worker.log'
|
||||
$metaPath = Join-Path -Path $stateRoot -ChildPath 'worker-info.json'
|
||||
$commandPath = Join-Path -Path $stateRoot -ChildPath 'commands.txt'
|
||||
$payloadPath = Join-Path -Path $stateRoot -ChildPath "payload-$([Guid]::NewGuid().ToString()).ps1"
|
||||
# endregion
|
||||
|
||||
# region Helpers
|
||||
function Write-ControllerLog {
|
||||
param([string]$Message)
|
||||
$timestamp = (Get-Date).ToString('u')
|
||||
Add-Content -Path $logPath -Value "[CTRL $timestamp] $Message" -Encoding UTF8
|
||||
}
|
||||
|
||||
function Write-Metadata {
|
||||
param(
|
||||
[string]$Status,
|
||||
[nullable[int]]$WorkerPid = $null,
|
||||
[nullable[int]]$ControllerPid = $PID,
|
||||
[int]$Restarts = 0,
|
||||
[nullable[int]]$LastExitCode = $null
|
||||
)
|
||||
|
||||
$payload = [pscustomobject]@{
|
||||
WorkerName = $WorkerName
|
||||
WorkerType = $WorkerType
|
||||
Status = $Status
|
||||
ControllerPid = $ControllerPid
|
||||
WorkerPid = $WorkerPid
|
||||
Restarts = $Restarts
|
||||
LastExitCode = $LastExitCode
|
||||
LogPath = $logPath
|
||||
CommandPath = $commandPath
|
||||
PayloadPath = $payloadPath
|
||||
UpdatedAtUtc = (Get-Date).ToUniversalTime()
|
||||
}
|
||||
|
||||
$payload | ConvertTo-Json -Depth 5 | Set-Content -Path $metaPath -Encoding UTF8
|
||||
}
|
||||
|
||||
function Get-PendingCommands {
|
||||
if (-not (Test-Path $commandPath)) {
|
||||
return @()
|
||||
}
|
||||
|
||||
try {
|
||||
$lines = Get-Content -Path $commandPath -ErrorAction Stop
|
||||
Remove-Item -Path $commandPath -Force -ErrorAction SilentlyContinue
|
||||
return $lines | Where-Object { -not [string]::IsNullOrWhiteSpace($_) }
|
||||
}
|
||||
catch {
|
||||
return @()
|
||||
}
|
||||
}
|
||||
# endregion
|
||||
|
||||
try {
|
||||
# Write payload script to disk
|
||||
$payloadBytes = [Convert]::FromBase64String($PayloadBase64)
|
||||
[IO.File]::WriteAllBytes($payloadPath, $payloadBytes)
|
||||
Write-ControllerLog "Payload written to $payloadPath"
|
||||
}
|
||||
catch {
|
||||
Write-Error "Unable to write payload: $($_.Exception.Message)"
|
||||
exit 1
|
||||
}
|
||||
|
||||
$restartCount = 0
|
||||
$controllerPid = $PID
|
||||
|
||||
while ($restartCount -le $MaxRestarts) {
|
||||
try {
|
||||
# Initialize worker process
|
||||
$psi = [System.Diagnostics.ProcessStartInfo]::new()
|
||||
$psi.FileName = (Get-Command pwsh -ErrorAction SilentlyContinue)?.Source
|
||||
if (-not $psi.FileName) {
|
||||
$psi.FileName = (Get-Command powershell -ErrorAction Stop).Source
|
||||
}
|
||||
|
||||
$psi.Arguments = "-NoLogo -NoProfile -ExecutionPolicy Bypass -File `"$payloadPath`""
|
||||
$psi.UseShellExecute = $false
|
||||
$psi.RedirectStandardInput = $true
|
||||
$psi.RedirectStandardOutput = $true
|
||||
$psi.RedirectStandardError = $true
|
||||
$psi.CreateNoWindow = $true
|
||||
|
||||
$workerProcess = New-Object System.Diagnostics.Process
|
||||
$workerProcess.StartInfo = $psi
|
||||
|
||||
if (-not $workerProcess.Start()) {
|
||||
throw "Failed to start worker process."
|
||||
}
|
||||
|
||||
Write-ControllerLog "Worker process started with PID $($workerProcess.Id)"
|
||||
Write-Metadata -Status 'running' -WorkerPid $workerProcess.Id -ControllerPid $controllerPid -Restarts $restartCount
|
||||
|
||||
$logWriter = [System.IO.StreamWriter]::new($logPath, $true, [System.Text.Encoding]::UTF8)
|
||||
$logWriter.AutoFlush = $true
|
||||
|
||||
$outputHandler = [System.Diagnostics.DataReceivedEventHandler]{
|
||||
param($s, $e)
|
||||
if ($e.Data) {
|
||||
$ts = (Get-Date).ToString('u')
|
||||
$logWriter.WriteLine("[OUT $ts] $($e.Data)")
|
||||
}
|
||||
}
|
||||
$errorHandler = [System.Diagnostics.DataReceivedEventHandler]{
|
||||
param($s, $e)
|
||||
if ($e.Data) {
|
||||
$ts = (Get-Date).ToString('u')
|
||||
$logWriter.WriteLine("[ERR $ts] $($e.Data)")
|
||||
}
|
||||
}
|
||||
|
||||
$workerProcess.add_OutputDataReceived($outputHandler)
|
||||
$workerProcess.add_ErrorDataReceived($errorHandler)
|
||||
$workerProcess.BeginOutputReadLine()
|
||||
$workerProcess.BeginErrorReadLine()
|
||||
|
||||
while (-not $workerProcess.HasExited) {
|
||||
$commands = Get-PendingCommands
|
||||
foreach ($command in $commands) {
|
||||
$trimmed = $command.Trim()
|
||||
if (-not $trimmed) { continue }
|
||||
|
||||
Write-ControllerLog "Received command '$trimmed'"
|
||||
try {
|
||||
$workerProcess.StandardInput.WriteLine($trimmed)
|
||||
$workerProcess.StandardInput.Flush()
|
||||
}
|
||||
catch {
|
||||
Write-ControllerLog "Failed to forward command '$trimmed': $($_.Exception.Message)"
|
||||
}
|
||||
|
||||
if ($trimmed -ieq 'quit') {
|
||||
Write-ControllerLog "Quit command issued. Waiting for worker to exit."
|
||||
}
|
||||
}
|
||||
|
||||
Start-Sleep -Milliseconds 500
|
||||
}
|
||||
|
||||
$exitCode = $workerProcess.ExitCode
|
||||
Write-ControllerLog "Worker exited with code $exitCode"
|
||||
Write-Metadata -Status 'stopped' -WorkerPid $null -ControllerPid $controllerPid -Restarts $restartCount -LastExitCode $exitCode
|
||||
|
||||
if ($exitCode -eq 0) {
|
||||
break
|
||||
}
|
||||
}
|
||||
catch {
|
||||
Write-ControllerLog "Controller error: $($_.Exception.Message)"
|
||||
}
|
||||
|
||||
$restartCount++
|
||||
if ($restartCount -gt $MaxRestarts) {
|
||||
Write-ControllerLog "Maximum restart attempts reached. Controller stopping."
|
||||
break
|
||||
}
|
||||
|
||||
Write-ControllerLog "Restarting worker in $RestartDelaySeconds seconds (attempt $restartCount of $MaxRestarts)."
|
||||
Start-Sleep -Seconds $RestartDelaySeconds
|
||||
}
|
||||
|
||||
Write-Metadata -Status 'inactive' -WorkerPid $null -ControllerPid $controllerPid -Restarts $restartCount
|
||||
Write-ControllerLog "Controller exiting."
|
||||
|
||||
Reference in New Issue
Block a user