attempting fix
This commit is contained in:
Nathan
@@ -50,9 +50,91 @@ $workers = @(
|
||||
}
|
||||
)
|
||||
|
||||
$global:UnifiedWorkerDataRoot = 'C:\ProgramData\UnifiedWorkers'
|
||||
$script:ControllerScriptBase64 = $null
|
||||
$script:AttachHelperScriptBase64 = $null
|
||||
$script:WorkerBasePathCache = @{}
|
||||
|
||||
function Build-SshArgsFromParts {
|
||||
param(
|
||||
[pscustomobject]$Parts,
|
||||
[switch]$Interactive
|
||||
)
|
||||
|
||||
$args = @('-o','ServerAliveInterval=60','-o','ServerAliveCountMax=30')
|
||||
if ($Interactive -and $Parts.RequestPty) {
|
||||
$args += '-t'
|
||||
}
|
||||
elseif (-not $Interactive) {
|
||||
$args += '-T'
|
||||
}
|
||||
|
||||
$args += $Parts.Options
|
||||
|
||||
if ($Parts.Port) {
|
||||
$args += '-p'
|
||||
$args += $Parts.Port
|
||||
}
|
||||
|
||||
$args += $Parts.Host
|
||||
return $args
|
||||
}
|
||||
|
||||
function Build-ScpArgsFromParts {
|
||||
param(
|
||||
[pscustomobject]$Parts
|
||||
)
|
||||
|
||||
$args = @('-o','ServerAliveInterval=60','-o','ServerAliveCountMax=30')
|
||||
$args += $Parts.Options
|
||||
if ($Parts.Port) {
|
||||
$args += '-P'
|
||||
$args += $Parts.Port
|
||||
}
|
||||
return $args
|
||||
}
|
||||
|
||||
function Get-SshArgs {
|
||||
param(
|
||||
[object]$Worker,
|
||||
[switch]$Interactive
|
||||
)
|
||||
|
||||
$parts = Get-WorkerConnectionParts -RawArgs $Worker.SSHArgs -DefaultHost $Worker.Name
|
||||
return Build-SshArgsFromParts -Parts $parts -Interactive:$Interactive
|
||||
}
|
||||
|
||||
function Get-WorkerBasePath {
|
||||
param(
|
||||
[object]$Worker,
|
||||
[pscustomobject]$ConnectionParts = $null
|
||||
)
|
||||
|
||||
if ($script:WorkerBasePathCache.ContainsKey($Worker.Name)) {
|
||||
return $script:WorkerBasePathCache[$Worker.Name]
|
||||
}
|
||||
|
||||
if (-not $ConnectionParts) {
|
||||
$ConnectionParts = Get-WorkerConnectionParts -RawArgs $Worker.SSHArgs -DefaultHost $Worker.Name
|
||||
}
|
||||
|
||||
$sshArgs = Build-SshArgsFromParts -Parts $ConnectionParts -Interactive:$false
|
||||
$scriptBlock = "`$ProgressPreference='SilentlyContinue'; [Environment]::GetFolderPath('LocalApplicationData')"
|
||||
$encoded = [Convert]::ToBase64String([Text.Encoding]::Unicode.GetBytes($scriptBlock))
|
||||
$remoteCmd = "powershell -NoLogo -NoProfile -ExecutionPolicy Bypass -EncodedCommand $encoded"
|
||||
$output = & ssh @sshArgs $remoteCmd
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
throw "Unable to determine LocalAppData on $($Worker.Name)."
|
||||
}
|
||||
|
||||
$base = ($output | Where-Object { -not [string]::IsNullOrWhiteSpace($_) } | Select-Object -Last 1).Trim()
|
||||
if (-not $base) {
|
||||
throw "Unable to read LocalAppData path on $($Worker.Name)."
|
||||
}
|
||||
|
||||
$final = Join-Path $base 'UnifiedWorkers'
|
||||
$script:WorkerBasePathCache[$Worker.Name] = $final
|
||||
return $final
|
||||
}
|
||||
|
||||
function Get-FileBase64 {
|
||||
param([string]$Path)
|
||||
@@ -88,63 +170,139 @@ function Get-WorkerSshArgs {
|
||||
return $RawArgs -split '\s+' | Where-Object { $_.Trim().Length -gt 0 }
|
||||
}
|
||||
|
||||
function Get-WorkerConnectionParts {
|
||||
param(
|
||||
[string]$RawArgs,
|
||||
[string]$DefaultHost
|
||||
)
|
||||
|
||||
$tokens = Get-WorkerSshArgs -RawArgs $RawArgs
|
||||
$options = New-Object System.Collections.Generic.List[string]
|
||||
$targetHost = $null
|
||||
$port = $null
|
||||
$requestPty = $false
|
||||
$optionsWithArgs = @('-i','-o','-c','-D','-E','-F','-I','-J','-L','-l','-m','-O','-Q','-R','-S','-W','-w')
|
||||
|
||||
for ($i = 0; $i -lt $tokens.Count; $i++) {
|
||||
$token = $tokens[$i]
|
||||
if ($token -eq '-t' -or $token -eq '-tt') {
|
||||
$requestPty = $true
|
||||
continue
|
||||
}
|
||||
|
||||
if ($token -eq '-p' -and ($i + 1) -lt $tokens.Count) {
|
||||
$port = $tokens[$i + 1]
|
||||
$i++
|
||||
continue
|
||||
}
|
||||
|
||||
if ($token.StartsWith('-')) {
|
||||
$options.Add($token)
|
||||
if ($optionsWithArgs -contains $token -and ($i + 1) -lt $tokens.Count) {
|
||||
$options.Add($tokens[$i + 1])
|
||||
$i++
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if (-not $targetHost) {
|
||||
$targetHost = $token
|
||||
continue
|
||||
}
|
||||
|
||||
$options.Add($token)
|
||||
}
|
||||
|
||||
if (-not $targetHost) {
|
||||
$targetHost = $DefaultHost
|
||||
}
|
||||
|
||||
return [pscustomobject]@{
|
||||
Host = $targetHost
|
||||
Options = $options.ToArray()
|
||||
Port = $port
|
||||
RequestPty = $requestPty
|
||||
}
|
||||
}
|
||||
|
||||
function Invoke-RemotePowerShell {
|
||||
param(
|
||||
[Parameter(Mandatory = $true)][object]$Worker,
|
||||
[Parameter(Mandatory = $true)][string]$Script
|
||||
[Parameter(Mandatory = $true)][string]$Script,
|
||||
[switch]$Interactive
|
||||
)
|
||||
|
||||
$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"
|
||||
$parts = Get-WorkerConnectionParts -RawArgs $Worker.SSHArgs -DefaultHost $Worker.Name
|
||||
if (-not $parts.Host) {
|
||||
throw "Unable to determine SSH host for $($Worker.Name)"
|
||||
}
|
||||
|
||||
& ssh @sshArgs
|
||||
return $LASTEXITCODE
|
||||
$sshBaseArgs = Build-SshArgsFromParts -Parts $parts -Interactive:$Interactive
|
||||
$remoteBasePath = Get-WorkerBasePath -Worker $Worker -ConnectionParts $parts
|
||||
$localTemp = [System.IO.Path]::GetTempFileName() + '.ps1'
|
||||
Set-Content -Path $localTemp -Value $Script -Encoding UTF8
|
||||
|
||||
$remoteTmpDir = Join-Path $remoteBasePath 'tmp'
|
||||
$remoteScriptWin = Join-Path $remoteTmpDir ("script-{0}.ps1" -f ([guid]::NewGuid().ToString()))
|
||||
$remoteScriptScp = $remoteScriptWin -replace '\\','/'
|
||||
$remoteTarget = "{0}:{1}" -f $parts.Host, ('"'+$remoteScriptScp+'"')
|
||||
|
||||
$ensureScript = "New-Item -ItemType Directory -Path '$remoteTmpDir' -Force | Out-Null"
|
||||
$ensureCmd = "powershell -NoLogo -NoProfile -ExecutionPolicy Bypass -EncodedCommand " + [Convert]::ToBase64String([Text.Encoding]::Unicode.GetBytes($ensureScript))
|
||||
& ssh @sshBaseArgs $ensureCmd
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
Remove-Item $localTemp -ErrorAction SilentlyContinue
|
||||
return $LASTEXITCODE
|
||||
}
|
||||
|
||||
$scpArgs = Build-ScpArgsFromParts -Parts $parts
|
||||
$scpArgs += $localTemp
|
||||
$scpArgs += $remoteTarget
|
||||
|
||||
& scp @scpArgs
|
||||
$scpExit = $LASTEXITCODE
|
||||
Remove-Item $localTemp -ErrorAction SilentlyContinue
|
||||
if ($scpExit -ne 0) {
|
||||
return $scpExit
|
||||
}
|
||||
|
||||
$execCmd = "powershell -NoLogo -NoProfile -ExecutionPolicy Bypass -File `"$remoteScriptWin`""
|
||||
& ssh @sshBaseArgs $execCmd
|
||||
$execExit = $LASTEXITCODE
|
||||
|
||||
$cleanupScript = "Remove-Item -LiteralPath '$remoteScriptWin' -ErrorAction SilentlyContinue"
|
||||
$cleanupCmd = "powershell -NoLogo -NoProfile -ExecutionPolicy Bypass -EncodedCommand " + [Convert]::ToBase64String([Text.Encoding]::Unicode.GetBytes($cleanupScript))
|
||||
& ssh @sshBaseArgs $cleanupCmd | Out-Null
|
||||
|
||||
return $execExit
|
||||
}
|
||||
|
||||
function Ensure-ControllerDeployed {
|
||||
param([object]$Worker)
|
||||
$controllerBase64 = Get-ControllerScriptBase64
|
||||
$script = @'
|
||||
`$dataRoot = '{0}'
|
||||
$script = @"
|
||||
`$ProgressPreference = 'SilentlyContinue'
|
||||
`$dataRoot = Join-Path ([Environment]::GetFolderPath('LocalApplicationData')) 'UnifiedWorkers'
|
||||
New-Item -ItemType Directory -Path `$dataRoot -Force | Out-Null
|
||||
`$controllerPath = Join-Path `$dataRoot 'controller.ps1'
|
||||
[IO.File]::WriteAllBytes(`$controllerPath, [Convert]::FromBase64String('{1}'))
|
||||
Write-Host "Controller ready at `$controllerPath"
|
||||
'@ -f $global:UnifiedWorkerDataRoot, $controllerBase64
|
||||
[IO.File]::WriteAllBytes(`$controllerPath, [Convert]::FromBase64String('$controllerBase64'))
|
||||
"@
|
||||
Invoke-RemotePowerShell -Worker $Worker -Script $script | Out-Null
|
||||
}
|
||||
|
||||
function Ensure-AttachHelperDeployed {
|
||||
param([object]$Worker)
|
||||
$helperBase64 = Get-AttachHelperScriptBase64
|
||||
$script = @'
|
||||
`$dataRoot = '{0}'
|
||||
$script = @"
|
||||
`$ProgressPreference = 'SilentlyContinue'
|
||||
`$dataRoot = Join-Path ([Environment]::GetFolderPath('LocalApplicationData')) 'UnifiedWorkers'
|
||||
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
|
||||
[IO.File]::WriteAllBytes(`$attachPath, [Convert]::FromBase64String('$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
|
||||
return [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,
|
||||
@@ -152,15 +310,24 @@ function Get-EnsureWorkerScript {
|
||||
[string]$PayloadBase64
|
||||
)
|
||||
|
||||
$paths = Get-RemoteWorkerPaths -WorkerType $WorkerType -WorkerName $WorkerName
|
||||
$controllerPath = Join-Path $global:UnifiedWorkerDataRoot 'controller.ps1'
|
||||
$payload = @{
|
||||
WorkerName = $WorkerName
|
||||
WorkerType = $WorkerType
|
||||
PayloadBase64 = $PayloadBase64
|
||||
} | ConvertTo-Json -Compress
|
||||
|
||||
return @'
|
||||
`$controllerPath = '{0}'
|
||||
`$metaPath = '{1}'
|
||||
`$workerName = '{2}'
|
||||
`$workerType = '{3}'
|
||||
`$payloadBase64 = '{4}'
|
||||
$payloadBase64 = ConvertTo-Base64Unicode -Content $payload
|
||||
|
||||
return @"
|
||||
`$ProgressPreference = 'SilentlyContinue'
|
||||
`$params = ConvertFrom-Json ([Text.Encoding]::Unicode.GetString([Convert]::FromBase64String('$payloadBase64')))
|
||||
`$workerName = `$params.WorkerName
|
||||
`$workerType = `$params.WorkerType
|
||||
`$payloadBase64 = `$params.PayloadBase64
|
||||
`$dataRoot = Join-Path ([Environment]::GetFolderPath('LocalApplicationData')) 'UnifiedWorkers'
|
||||
`$instanceRoot = Join-Path (Join-Path `$dataRoot `$workerType) `$workerName
|
||||
`$metaPath = Join-Path `$instanceRoot 'state\worker-info.json'
|
||||
`$controllerPath = Join-Path `$dataRoot 'controller.ps1'
|
||||
|
||||
if (-not (Test-Path `$controllerPath)) {
|
||||
throw "Controller missing at `$controllerPath"
|
||||
@@ -182,8 +349,11 @@ if (Test-Path `$metaPath) {
|
||||
}
|
||||
|
||||
if (`$shouldStart) {
|
||||
`$psExe = (Get-Command pwsh -ErrorAction SilentlyContinue)?.Source
|
||||
if (-not `$psExe) {
|
||||
`$pwsh = Get-Command pwsh -ErrorAction SilentlyContinue
|
||||
if (`$pwsh) {
|
||||
`$psExe = `$pwsh.Source
|
||||
}
|
||||
else {
|
||||
`$psExe = (Get-Command powershell -ErrorAction Stop).Source
|
||||
}
|
||||
|
||||
@@ -198,7 +368,7 @@ if (`$shouldStart) {
|
||||
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 {
|
||||
@@ -236,17 +406,14 @@ function Invoke-WorkerAttach {
|
||||
$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
|
||||
$parts = Get-WorkerConnectionParts -RawArgs $Worker.SSHArgs -DefaultHost $Worker.Name
|
||||
$sshArgs = Build-SshArgsFromParts -Parts $parts -Interactive:(!$CommandOnly)
|
||||
$remoteBasePath = Get-WorkerBasePath -Worker $Worker -ConnectionParts $parts
|
||||
$remoteHelper = Join-Path $remoteBasePath 'attach-helper.ps1'
|
||||
$quotedArgs = ($paramsBlock | ForEach-Object { '"' + ($_ -replace '"','""') + '"' }) -join ' '
|
||||
$remoteCmd = "powershell -NoLogo -NoProfile -ExecutionPolicy Bypass -File `"$remoteHelper`" $quotedArgs"
|
||||
|
||||
Invoke-RemotePowerShell -Worker $Worker -Script $script | Out-Null
|
||||
& ssh @sshArgs $remoteCmd
|
||||
}
|
||||
|
||||
# FUNCTIONS
|
||||
|
||||
Reference in New Issue
Block a user