Team Ai
Datasetpublic

MegaBites-AI/Windows-powershell

sourceHugging Facemitupdated 6mo agoView on Hugging Face
0likes372downloads
RunspaceRef.cs432 linesDownload Raw Back to client
1// Copyright (c) Microsoft Corporation.2// Licensed under the MIT License.3 4using System.Collections.ObjectModel;5using System.Management.Automation.Internal;6using System.Management.Automation.Runspaces;7using System.Management.Automation.Runspaces.Internal;8using System.Management.Automation.Security;9 10using Dbg = System.Management.Automation.Diagnostics;11 12namespace System.Management.Automation.Remoting13{14    /// <summary>15    /// Provides a reference to a runspace that can be used to temporarily16    /// push a remote runspace on top of a local runspace. This is17    /// primary used by Start-PSSession. The purpose of this class is to hide18    /// the CreatePipeline method and force it to be used as defined in this19    /// class.20    /// </summary>21    internal class RunspaceRef22    {23        /// <summary>24        /// Runspace ref.25        /// </summary>26        private readonly ObjectRef<Runspace> _runspaceRef;27        private bool _stopInvoke;28        private readonly object _localSyncObject;29        private static readonly RobustConnectionProgress s_RCProgress = new RobustConnectionProgress();30 31        /// <summary>32        /// Constructor for RunspaceRef.33        /// </summary>34        internal RunspaceRef(Runspace runspace)35        {36            Dbg.Assert(runspace != null, "Expected runspace != null");37            _runspaceRef = new ObjectRef<Runspace>(runspace);38            _stopInvoke = false;39            _localSyncObject = new object();40        }41 42        /// <summary>43        /// Revert.44        /// </summary>45        internal void Revert()46        {47            _runspaceRef.Revert();48 49            lock (_localSyncObject)50            {51                _stopInvoke = true;52            }53        }54 55        /// <summary>56        /// Runspace.57        /// </summary>58        internal Runspace Runspace59        {60            get61            {62                return _runspaceRef.Value;63            }64        }65 66        internal Runspace OldRunspace67        {68            get { return _runspaceRef.OldValue; }69        }70 71        /// <summary>72        /// Is runspace overridden.73        /// </summary>74        internal bool IsRunspaceOverridden75        {76            get77            {78                return _runspaceRef.IsOverridden;79            }80        }81 82        /// <summary>83        /// Parse ps command using script block.84        /// </summary>85        private PSCommand ParsePsCommandUsingScriptBlock(string line, bool? useLocalScope)86        {87            try88            {89                // Extract execution context from local runspace.90                Runspace localRunspace = _runspaceRef.OldValue;91                ExecutionContext context = localRunspace.ExecutionContext;92 93                // This is trusted input as long as we're in FullLanguage mode94                // and if we are not in a loopback configuration mode, in which case we always force remote script commands95                // to be parsed and evaluated on the remote session (not in the current local session).96                RemoteRunspace remoteRunspace = _runspaceRef.Value as RemoteRunspace;97                bool isConfiguredLoopback = remoteRunspace != null && remoteRunspace.IsConfiguredLoopBack;98 99                bool inFullLanguage = context.LanguageMode == PSLanguageMode.FullLanguage;100                if (context.LanguageMode == PSLanguageMode.ConstrainedLanguage101                    && SystemPolicy.GetSystemLockdownPolicy() == SystemEnforcementMode.Audit)102                {103                    // In audit mode, report but don't enforce.104                    inFullLanguage = true;105                    SystemPolicy.LogWDACAuditMessage(106                        context: context,107                        title: RemotingErrorIdStrings.WDACGetPowerShellLogTitle,108                        message: RemotingErrorIdStrings.WDACGetPowerShellLogMessage,109                        fqid: "GetPowerShellMayFail",110                        dropIntoDebugger: true);111                }112 113                bool isTrustedInput = !isConfiguredLoopback && inFullLanguage;114 115                // Create PowerShell from ScriptBlock.116                ScriptBlock scriptBlock = ScriptBlock.Create(context, line);117                PowerShell powerShell = scriptBlock.GetPowerShell(context, isTrustedInput, useLocalScope, null);118                return powerShell.Commands;119            }120            catch (ScriptBlockToPowerShellNotSupportedException)121            {122            }123            catch (RuntimeException)124            {125            }126 127            // If parsing failed return null.128            return null;129        }130 131        /// <summary>132        /// Create ps command.133        /// </summary>134        internal PSCommand CreatePsCommand(string line, bool isScript, bool? useNewScope)135        {136            // Fall-back to traditional approach if runspace is not pushed.137            if (!this.IsRunspaceOverridden)138            {139                return CreatePsCommandNotOverridden(line, isScript, useNewScope);140            }141 142            // Try to parse commands as script-block.143            PSCommand psCommand = ParsePsCommandUsingScriptBlock(line, useNewScope);144 145            // If that didn't work fall back to traditional approach.146            if (psCommand == null)147            {148                return CreatePsCommandNotOverridden(line, isScript, useNewScope);149            }150 151            // Otherwise return the psCommandCollection we got.152            return psCommand;153        }154 155        /// <summary>156        /// Creates the PSCommand when the runspace is not overridden.157        /// </summary>158        private static PSCommand CreatePsCommandNotOverridden(string line, bool isScript, bool? useNewScope)159        {160            PSCommand command = new PSCommand();161 162            if (isScript)163            {164                if (useNewScope.HasValue)165                {166                    command.AddScript(line, useNewScope.Value);167                }168                else169                {170                    command.AddScript(line);171                }172            }173            else174            {175                if (useNewScope.HasValue)176                {177                    command.AddCommand(line, useNewScope.Value);178                }179                else180                {181                    command.AddCommand(line);182                }183            }184 185            return command;186        }187 188        /// <summary>189        /// Create pipeline.190        /// </summary>191        internal Pipeline CreatePipeline(string line, bool addToHistory, bool useNestedPipelines)192        {193            // This method allows input commands to work against no-language runspaces. If a runspace194            // is pushed, it tries to parse the line using a ScriptBlock object. If a runspace is not195            // pushed, or if the parsing fails, in these cases it reverts to calling CreatePipeline196            // using the unparsed line.197 198            Pipeline pipeline = null;199 200            // In Start-PSSession scenario try to create a pipeline by parsing the line as a script block.201            if (this.IsRunspaceOverridden)202            {203                // Win8: exit should work to escape from the restrictive session204                if ((_runspaceRef.Value is RemoteRunspace) &&205                    (!string.IsNullOrEmpty(line) && string.Equals(line.Trim(), "exit", StringComparison.OrdinalIgnoreCase)))206                {207                    line = "Exit-PSSession";208                }209 210                PSCommand psCommand = ParsePsCommandUsingScriptBlock(line, null);211                if (psCommand != null)212                {213                    pipeline = useNestedPipelines ?214                        _runspaceRef.Value.CreateNestedPipeline(psCommand.Commands[0].CommandText, addToHistory) :215                        _runspaceRef.Value.CreatePipeline(psCommand.Commands[0].CommandText, addToHistory);216 217                    pipeline.Commands.Clear();218 219                    foreach (Command command in psCommand.Commands)220                    {221                        pipeline.Commands.Add(command);222                    }223                }224            }225 226            // If that didn't work out fall-back to the traditional approach.227            pipeline ??= useNestedPipelines ?228                _runspaceRef.Value.CreateNestedPipeline(line, addToHistory) :229                _runspaceRef.Value.CreatePipeline(line, addToHistory);230 231            // Add robust connection callback if this is a pushed runspace.232            RemotePipeline remotePipeline = pipeline as RemotePipeline;233            if (this.IsRunspaceOverridden && remotePipeline != null)234            {235                PowerShell shell = remotePipeline.PowerShell;236                if (shell.RemotePowerShell != null)237                {238                    shell.RemotePowerShell.RCConnectionNotification += HandleRCConnectionNotification;239                }240 241                // Add callback to write robust connection errors from stream.242                shell.ErrorBuffer.DataAdded += (sender, eventArgs) =>243                {244                    RemoteRunspace remoteRunspace = _runspaceRef.Value as RemoteRunspace;245                    PSDataCollection<ErrorRecord> erBuffer = sender as PSDataCollection<ErrorRecord>;246                    if (remoteRunspace != null && erBuffer != null &&247                        remoteRunspace.RunspacePool.RemoteRunspacePoolInternal.Host != null)248                    {249                        Collection<ErrorRecord> erRecords = erBuffer.ReadAll();250                        foreach (var er in erRecords)251                        {252                            remoteRunspace.RunspacePool.RemoteRunspacePoolInternal.Host.UI.WriteErrorLine(er.ToString());253                        }254                    }255                };256            }257 258            pipeline.SetHistoryString(line);259 260            return pipeline;261        }262 263        /// <summary>264        /// Create pipeline.265        /// </summary>266        internal Pipeline CreatePipeline()267        {268            return _runspaceRef.Value.CreatePipeline();269        }270 271        /// <summary>272        /// Create nested pipeline.273        /// </summary>274        internal Pipeline CreateNestedPipeline()275        {276            return _runspaceRef.Value.CreateNestedPipeline();277        }278 279        /// <summary>280        /// Override.281        /// </summary>282        internal void Override(RemoteRunspace remoteRunspace)283        {284            bool isRunspacePushed = false;285            Override(remoteRunspace, null, out isRunspacePushed);286        }287 288        /// <summary>289        /// Override inside a safe lock.290        /// </summary>291        /// <param name="remoteRunspace">Runspace to override.</param>292        /// <param name="syncObject">Object to use in synchronization.</param>293        /// <param name="isRunspacePushed">Set is runspace pushed.</param>294        internal void Override(RemoteRunspace remoteRunspace, object syncObject, out bool isRunspacePushed)295        {296            lock (_localSyncObject)297            {298                _stopInvoke = false;299            }300 301            try302            {303                if (syncObject != null)304                {305                    lock (syncObject)306                    {307                        _runspaceRef.Override(remoteRunspace);308                        isRunspacePushed = true;309                    }310                }311                else312                {313                    _runspaceRef.Override(remoteRunspace);314                    isRunspacePushed = true;315                }316 317                if ((remoteRunspace.GetCurrentlyRunningPipeline() != null))318                {319                    // Don't execute command if pushed runspace is already running one.320                    return;321                }322 323                using (PowerShell powerShell = PowerShell.Create())324                {325                    powerShell.AddCommand("Get-Command");326                    powerShell.AddParameter("Name", new string[] { "Out-Default", "Exit-PSSession" });327                    powerShell.Runspace = _runspaceRef.Value;328 329                    bool isReleaseCandidateBackcompatibilityMode = _runspaceRef.Value.GetRemoteProtocolVersion() == RemotingConstants.ProtocolVersion_2_0;330                    powerShell.IsGetCommandMetadataSpecialPipeline = !isReleaseCandidateBackcompatibilityMode;331                    int expectedNumberOfResults = isReleaseCandidateBackcompatibilityMode ? 2 : 3;332 333                    powerShell.RemotePowerShell.HostCallReceived += HandleHostCall;334 335                    IAsyncResult asyncResult = powerShell.BeginInvoke();336                    PSDataCollection<PSObject> results = new PSDataCollection<PSObject>();337 338                    while (!_stopInvoke)339                    {340                        asyncResult.AsyncWaitHandle.WaitOne(1000);341 342                        if (asyncResult.IsCompleted)343                        {344                            results = powerShell.EndInvoke(asyncResult);345                            break;346                        }347                    }348 349                    if (powerShell.Streams.Error.Count > 0 || results.Count < expectedNumberOfResults)350                    {351                        throw RemoteHostExceptions.NewRemoteRunspaceDoesNotSupportPushRunspaceException();352                    }353                }354            }355            catch (Exception)356            {357                _runspaceRef.Revert();358                isRunspacePushed = false;359                throw;360            }361        }362 363        /// <summary>364        /// </summary>365        /// <param name="sender"></param>366        /// <param name="eventArgs"></param>367        private void HandleHostCall(object sender, RemoteDataEventArgs<RemoteHostCall> eventArgs)368        {369            ClientRemotePowerShell.ExitHandler(sender, eventArgs);370        }371 372        #region Robust Connection Support373 374        private void HandleRCConnectionNotification(object sender, PSConnectionRetryStatusEventArgs e)375        {376            switch (e.Notification)377            {378                case PSConnectionRetryStatus.NetworkFailureDetected:379                    StartProgressBar(sender.GetHashCode(), e.ComputerName, (e.MaxRetryConnectionTime / 1000));380                    break;381 382                case PSConnectionRetryStatus.AutoDisconnectStarting:383                case PSConnectionRetryStatus.ConnectionRetrySucceeded:384                    StopProgressBar(sender.GetHashCode());385                    break;386 387                case PSConnectionRetryStatus.AutoDisconnectSucceeded:388                case PSConnectionRetryStatus.InternalErrorAbort:389                    WriteRCFailedError();390                    StopProgressBar(sender.GetHashCode());391                    break;392            }393        }394 395        private void WriteRCFailedError()396        {397            RemoteRunspace remoteRunspace = _runspaceRef.Value as RemoteRunspace;398            if (remoteRunspace != null &&399                remoteRunspace.RunspacePool.RemoteRunspacePoolInternal.Host != null)400            {401                remoteRunspace.RunspacePool.RemoteRunspacePoolInternal.Host.UI.WriteErrorLine(402                    StringUtil.Format(RemotingErrorIdStrings.RCAutoDisconnectingError,403                    remoteRunspace.ConnectionInfo.ComputerName));404            }405        }406 407        private void StartProgressBar(408            long sourceId,409            string computerName,410            int totalSeconds)411        {412            RemoteRunspace remoteRunspace = _runspaceRef.Value as RemoteRunspace;413            if (remoteRunspace != null)414            {415                s_RCProgress.StartProgress(416                    sourceId,417                    computerName,418                    totalSeconds,419                    remoteRunspace.RunspacePool.RemoteRunspacePoolInternal.Host);420            }421        }422 423        private static void StopProgressBar(424            long sourceId)425        {426            s_RCProgress.StopProgress(sourceId);427        }428 429        #endregion430    }431}432