Last active
May 4, 2023 03:27
-
-
Save SeeminglyScience/1b72f0fde3840b275dd23b4104776c5d to your computer and use it in GitHub Desktop.
Proof of concept for "psedit" working outside of PSES.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
function Enter-PSSessionWithEdit { | |
[CmdletBinding()] | |
param( | |
[Parameter(Mandatory)] | |
[ValidateNotNullOrEmpty()] | |
[string] $ComputerName | |
) | |
end { | |
$enterEventName = 'RemoteSessionEditor.Enter' | |
if (-not $Host.Runspace.Events.GetEventSubscribers($enterEventName)) { | |
& { | |
$enterEventName = $enterEventName | |
# Called on the destination machine when the local editor closes. | |
$OnRemoteEditorExit = { | |
param([string] $path, [byte[]] $bytes) | |
end { | |
$stream = $null | |
try { | |
$stream = [System.IO.FileStream]::new( | |
$path, | |
[System.IO.FileMode]::OpenOrCreate, | |
[System.IO.FileAccess]::Write, | |
[System.IO.FileShare]::Read) | |
# Use async methods to avoid blocking the pipeline thread. | |
$task = $stream.WriteAsync($bytes, 0, $bytes.Length) | |
while (-not $task.AsyncWaitHandle.WaitOne(200)) { } | |
$null = $task.GetAwaiter().GetResult() | |
$task = $stream.FlushAsync() | |
while (-not $task.AsyncWaitHandle.WaitOne(200)) { } | |
$null = $task.GetAwaiter().GetResult() | |
} finally { | |
if ($null -ne $stream) { | |
$stream.Dispose() | |
} | |
} | |
} | |
} | |
# Called on the local machine when an edit session is requested. This will | |
# create a temporary file that represents the remote file, open the desired | |
# editor, then send changes back to the remote machine. | |
$OnRemoteEditorEnter = { | |
param([string] $path, [byte[]] $bytes) | |
end { | |
$tempFile = [System.IO.Path]::ChangeExtension( | |
[System.IO.Path]::GetRandomFileName(), | |
[System.IO.Path]::GetExtension($path)) | |
$tempFile = [System.IO.Path]::Combine( | |
[System.IO.Path]::GetTempPath(), | |
$tempFile) | |
$stream = $null | |
try { | |
$stream = [System.IO.FileStream]::new( | |
$tempFile, | |
[System.IO.FileMode]::CreateNew, | |
[System.IO.FileAccess]::Write, | |
[System.IO.FileShare]'ReadWrite, Delete') | |
# Use async methods to avoid blocking the pipeline thread. | |
$task = $stream.WriteAsync($bytes, 0, $bytes.Length) | |
while (-not $task.AsyncWaitHandle.WaitOne(200)) { } | |
$null = $task.GetAwaiter().GetResult() | |
$task = $stream.FlushAsync() | |
while (-not $task.AsyncWaitHandle.WaitOne(200)) { } | |
$null = $task.GetAwaiter().GetResult() | |
} finally { | |
if ($null -ne $stream) { | |
$stream.Dispose() | |
} | |
} | |
# Switch the current runspace back to the local machine so the editor can | |
# attach to the terminal. | |
$runspace = $Host.Runspace | |
$Host.PopRunspace() | |
$ps = $null | |
$nestedPs = $null | |
try { | |
$editor = $PSEdit | |
if ([string]::IsNullOrEmpty($editor)) { | |
foreach ($editorName in 'vi', 'vim', 'nano', 'notepad') { | |
if ($command = Get-Command $editorName -ErrorAction Ignore) { | |
$editor = $command.Path | |
break | |
} | |
} | |
} | |
if ([string]::IsNullOrEmpty($editor)) { | |
# Since this is being processed in an event subscriber, neither throw | |
# nor Write-Error will display any messages. Since this is only for | |
# interactive use it should be fine. | |
$Host.UI.WriteErrorLine( | |
'Cannot determine desired editor. Please populate "$global:PSEdit" with the command to use.') | |
return | |
} | |
# The command is constructed this way so that the native command processor | |
# knows to consider it to be "standalone". Without this, stdout is not | |
# displayed. | |
$nestedPs = [powershell]:: | |
Create('CurrentRunspace'). | |
AddScript('& $args[0] $args[1]', $false). | |
AddArgument($editor). | |
AddArgument($tempFile). | |
AddCommand('Out-Default') | |
$nestedPs.Invoke() | |
# This should be a lot smarter. At the very least it shouldn't send | |
# an event back if there was no changes. | |
$newBytes = [System.IO.File]::ReadAllBytes($tempFile) | |
$ps = [powershell]::Create() | |
$ps.Runspace = $runspace | |
$ps.AddScript($OnRemoteEditorExit). | |
AddArgument($path). | |
AddArgument($newBytes). | |
Invoke() | |
} finally { | |
if ($null -ne $ps) { | |
$ps.Dispose() | |
} | |
if ($null -ne $nestedPs) { | |
$nestedPs.Dispose() | |
} | |
Remove-Item $tempFile -ErrorAction Stop | |
$Host.PushRunspace($runspace) | |
} | |
} | |
} | |
# Create the event subscriber on the local machine that will handle opening | |
# an editor. | |
$sub = $Host.Runspace.Events.SubscribeEvent( | |
<# source: #> $null, | |
<# eventName: #> $enterEventName, | |
<# sourceIdentifier: #> $enterEventName, | |
<# data: #> $null, | |
<# action: #> $OnRemoteEditorEnter, | |
<# supportEvent: #> $true, | |
<# forwardEvent: #> $false) | |
$sub.Action.Module.SessionState.PSVariable.Set( | |
'OnRemoteEditorExit', | |
$OnRemoteEditorExit) | |
} | |
} | |
$session = New-PSSession $ComputerName -ErrorAction Stop | |
$ps = $null | |
try { | |
$ps = [powershell]::Create() | |
$ps.Runspace = $session.Runspace | |
$initializationScript = { | |
param([string] $enterEventName) | |
end { | |
$global:__enterEventName = $enterEventName | |
# Create a forward event on the remote machine that will send generated events | |
# to the runspace of the local machine. | |
$null = $Host.Runspace.Events.SubscribeEvent( | |
<# source: #> $null, | |
<# eventName: #> $enterEventName, | |
<# sourceIdentifier: #> $enterEventName, | |
<# data: #> $null, | |
<# action: #> $null, | |
<# supportEvent: #> $true, | |
<# forwardEvent: #> $true) | |
function psedit { | |
[CmdletBinding()] | |
param([string] $path) | |
end { | |
$fullPath = $PSCmdlet.SessionState.Path.GetUnresolvedProviderPathFromPSPath( | |
$path) | |
if (-not (Test-Path $fullPath)) { | |
$null = $Host.Runspace.Events.GenerateEvent( | |
<# sourceIdentifier: #> $global:__enterEventName, | |
<# sender: #> $null, | |
<# args: #> @($fullPath, [byte[]]::new(0)), | |
<# extraData: #> $null) | |
return | |
} | |
$fullPath = (Get-Item $path -ErrorAction Stop).FullName | |
$bytes = [System.IO.File]::ReadAllBytes($fullPath) | |
$null = $Host.Runspace.Events.GenerateEvent( | |
<# sourceIdentifier: #> $global:__enterEventName, | |
<# sender: #> $null, | |
<# args: #> @($fullPath, $bytes), | |
<# extraData: #> $null) | |
} | |
} | |
} | |
} | |
$ps.AddScript($initializationScript, <# useLocalScope: #> $false). | |
AddArgument($enterEventName). | |
Invoke() | |
Enter-PSSession $session | |
} finally { | |
if ($null -ne $ps) { | |
$ps.Dispose() | |
} | |
} | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment