55
66using Microsoft . PowerShell . EditorServices . Extensions ;
77using Microsoft . PowerShell . EditorServices . Utility ;
8+ using System ;
89using System . Collections . Generic ;
910using System . Diagnostics ;
1011using System . IO ;
1112using System . Linq ;
1213using System . Management . Automation ;
14+ using System . Management . Automation . Runspaces ;
1315using System . Threading . Tasks ;
1416
1517namespace Microsoft . PowerShell . EditorServices . Session
1618{
1719 /// <summary>
1820 /// Manages files that are accessed from a remote PowerShell session.
19- /// Also manages the registration and handling of the 'psedit' function
20- /// in 'LocalProcess' and 'Remote' runspaces.
21+ /// Also manages the registration and handling of the 'psedit' function.
2122 /// </summary>
2223 public class RemoteFileManager
2324 {
@@ -31,6 +32,51 @@ public class RemoteFileManager
3132 private Dictionary < RunspaceDetails , RemotePathMappings > filesPerRunspace =
3233 new Dictionary < RunspaceDetails , RemotePathMappings > ( ) ;
3334
35+ private const string RemoteSessionOpenFile = "PSESRemoteSessionOpenFile" ;
36+
37+ private const string PSEditFunctionScript = @"
38+ param (
39+ [Parameter(Mandatory=$true)] [String[]] $FileNames
40+ )
41+
42+ foreach ($fileName in $FileNames)
43+ {
44+ dir $fileName | where { ! $_.PSIsContainer } | foreach {
45+ $filePathName = $_.FullName
46+
47+ # Get file contents
48+ $contentBytes = Get-Content -Path $filePathName -Raw -Encoding Byte
49+
50+ # Notify client for file open.
51+ New-Event -SourceIdentifier PSESRemoteSessionOpenFile -EventArguments @($filePathName, $contentBytes) > $null
52+ }
53+ }
54+ " ;
55+
56+ // This script is templated so that the '-Forward' parameter can be added
57+ // to the script when in non-local sessions
58+ private const string CreatePSEditFunctionScript = @"
59+ param (
60+ [string] $PSEditFunction
61+ )
62+
63+ Register-EngineEvent -SourceIdentifier PSESRemoteSessionOpenFile {0}
64+
65+ if ((Test-Path -Path 'function:\global:PSEdit') -eq $false)
66+ {{
67+ Set-Item -Path 'function:\global:PSEdit' -Value $PSEditFunction
68+ }}
69+ " ;
70+
71+ private const string RemovePSEditFunctionScript = @"
72+ if ((Test-Path -Path 'function:\global:PSEdit') -eq $true)
73+ {
74+ Remove-Item -Path 'function:\global:PSEdit' -Force
75+ }
76+
77+ Get-EventSubscriber -SourceIdentifier PSESRemoteSessionOpenFile -EA Ignore | Remove-Event
78+ " ;
79+
3480 #endregion
3581
3682 #region Constructors
@@ -52,7 +98,7 @@ public RemoteFileManager(
5298 Validate . IsNotNull ( nameof ( editorOperations ) , editorOperations ) ;
5399
54100 this . powerShellContext = powerShellContext ;
55- this . powerShellContext . RunspaceChanged += PowerShellContext_RunspaceChanged ;
101+ this . powerShellContext . RunspaceChanged += HandleRunspaceChanged ;
56102
57103 this . editorOperations = editorOperations ;
58104
@@ -65,6 +111,9 @@ public RemoteFileManager(
65111
66112 // Delete existing temporary file cache path if it already exists
67113 this . TryDeleteTemporaryPath ( ) ;
114+
115+ // Register the psedit function in the current runspace
116+ this . RegisterPSEditFunction ( this . powerShellContext . CurrentRunspace ) ;
68117 }
69118
70119 #endregion
@@ -114,16 +163,14 @@ public async Task<string> FetchRemoteFile(
114163
115164 if ( fileContent != null )
116165 {
117- File . WriteAllBytes ( localFilePath , fileContent ) ;
166+ this . StoreRemoteFile ( localFilePath , fileContent , pathMappings ) ;
118167 }
119168 else
120169 {
121170 Logger . Write (
122171 LogLevel . Warning ,
123172 $ "Could not load contents of remote file '{ remoteFilePath } '") ;
124173 }
125-
126- pathMappings . AddOpenedLocalPath ( localFilePath ) ;
127174 }
128175 }
129176 }
@@ -213,6 +260,31 @@ public bool IsUnderRemoteTempPath(string filePath)
213260
214261 #region Private Methods
215262
263+ private string StoreRemoteFile (
264+ string remoteFilePath ,
265+ byte [ ] fileContent ,
266+ RunspaceDetails runspaceDetails )
267+ {
268+ RemotePathMappings pathMappings = this . GetPathMappings ( runspaceDetails ) ;
269+ string localFilePath = pathMappings . GetMappedPath ( remoteFilePath ) ;
270+
271+ this . StoreRemoteFile (
272+ localFilePath ,
273+ fileContent ,
274+ pathMappings ) ;
275+
276+ return localFilePath ;
277+ }
278+
279+ private void StoreRemoteFile (
280+ string localFilePath ,
281+ byte [ ] fileContent ,
282+ RemotePathMappings pathMappings )
283+ {
284+ File . WriteAllBytes ( localFilePath , fileContent ) ;
285+ pathMappings . AddOpenedLocalPath ( localFilePath ) ;
286+ }
287+
216288 private RemotePathMappings GetPathMappings ( RunspaceDetails runspaceDetails )
217289 {
218290 RemotePathMappings remotePathMappings = null ;
@@ -226,11 +298,12 @@ private RemotePathMappings GetPathMappings(RunspaceDetails runspaceDetails)
226298 return remotePathMappings ;
227299 }
228300
229- private async void PowerShellContext_RunspaceChanged ( object sender , RunspaceChangedEventArgs e )
301+ private async void HandleRunspaceChanged ( object sender , RunspaceChangedEventArgs e )
230302 {
303+
231304 if ( e . ChangeAction == RunspaceChangeAction . Enter )
232305 {
233- // TODO: Register psedit function and event handler
306+ this . RegisterPSEditFunction ( e . NewRunspace ) ;
234307 }
235308 else
236309 {
@@ -244,13 +317,116 @@ private async void PowerShellContext_RunspaceChanged(object sender, RunspaceChan
244317 }
245318 }
246319
247- // TODO: Clean up psedit registration
320+ if ( e . PreviousRunspace != null )
321+ {
322+ this . RemovePSEditFunction ( e . PreviousRunspace ) ;
323+ }
248324 }
249325 }
250326
251- #endregion
327+ private void HandlePSEventReceived ( object sender , PSEventArgs args )
328+ {
329+ if ( string . Equals ( RemoteSessionOpenFile , args . SourceIdentifier , StringComparison . CurrentCultureIgnoreCase ) )
330+ {
331+ try
332+ {
333+ if ( args . SourceArgs . Length >= 1 )
334+ {
335+ string localFilePath = string . Empty ;
336+ string remoteFilePath = args . SourceArgs [ 0 ] as string ;
252337
253- #region Private Methods
338+ // Is this a local process runspace? Treat as a local file
339+ if ( this . powerShellContext . CurrentRunspace . Location == RunspaceLocation . Local ||
340+ this . powerShellContext . CurrentRunspace . Location == RunspaceLocation . LocalProcess )
341+ {
342+ localFilePath = remoteFilePath ;
343+ }
344+ else
345+ {
346+ byte [ ] fileContent =
347+ args . SourceArgs . Length == 2
348+ ? ( byte [ ] ) ( ( args . SourceArgs [ 1 ] as PSObject ) . BaseObject )
349+ : new byte [ 0 ] ;
350+
351+ localFilePath =
352+ this . StoreRemoteFile (
353+ remoteFilePath ,
354+ fileContent ,
355+ this . powerShellContext . CurrentRunspace ) ;
356+ }
357+
358+ // Open the file in the editor
359+ this . editorOperations . OpenFile ( localFilePath ) ;
360+ }
361+ }
362+ catch ( NullReferenceException e )
363+ {
364+ Logger . WriteException ( "Could not store null remote file content" , e ) ;
365+ }
366+ }
367+ }
368+
369+ private void RegisterPSEditFunction ( RunspaceDetails runspaceDetails )
370+ {
371+ try
372+ {
373+ runspaceDetails . Runspace . Events . ReceivedEvents . PSEventReceived += HandlePSEventReceived ;
374+
375+ var createScript =
376+ string . Format (
377+ CreatePSEditFunctionScript ,
378+ ( runspaceDetails . Location == RunspaceLocation . Local && ! runspaceDetails . IsAttached )
379+ ? string . Empty : "-Forward" ) ;
380+
381+ PSCommand createCommand = new PSCommand ( ) ;
382+ createCommand
383+ . AddScript ( createScript )
384+ . AddParameter ( "PSEditFunction" , PSEditFunctionScript ) ;
385+
386+ if ( runspaceDetails . IsAttached )
387+ {
388+ this . powerShellContext . ExecuteCommand ( createCommand ) . Wait ( ) ;
389+ }
390+ else
391+ {
392+ using ( var powerShell = System . Management . Automation . PowerShell . Create ( ) )
393+ {
394+ powerShell . Runspace = runspaceDetails . Runspace ;
395+ powerShell . Commands = createCommand ;
396+ powerShell . Invoke ( ) ;
397+ }
398+ }
399+ }
400+ catch ( RemoteException e )
401+ {
402+ Logger . WriteException ( "Could not create psedit function." , e ) ;
403+ }
404+ }
405+
406+ private void RemovePSEditFunction ( RunspaceDetails runspaceDetails )
407+ {
408+ try
409+ {
410+ if ( runspaceDetails . Runspace . Events != null )
411+ {
412+ runspaceDetails . Runspace . Events . ReceivedEvents . PSEventReceived -= HandlePSEventReceived ;
413+ }
414+
415+ if ( runspaceDetails . Runspace . RunspaceStateInfo . State == RunspaceState . Opened )
416+ {
417+ using ( var powerShell = System . Management . Automation . PowerShell . Create ( ) )
418+ {
419+ powerShell . Runspace = runspaceDetails . Runspace ;
420+ powerShell . Commands . AddScript ( RemovePSEditFunctionScript ) ;
421+ powerShell . Invoke ( ) ;
422+ }
423+ }
424+ }
425+ catch ( RemoteException e )
426+ {
427+ Logger . WriteException ( "Could not remove psedit function." , e ) ;
428+ }
429+ }
254430
255431 private void TryDeleteTemporaryPath ( )
256432 {
@@ -265,9 +441,8 @@ private void TryDeleteTemporaryPath()
265441 }
266442 catch ( IOException e )
267443 {
268- Logger . Write (
269- LogLevel . Error ,
270- $ "Could not delete temporary folder for current process: { this . processTempPath } \r \n \r \n { e . ToString ( ) } ") ;
444+ Logger . WriteException (
445+ $ "Could not delete temporary folder for current process: { this . processTempPath } ", e ) ;
271446 }
272447 }
273448
0 commit comments