using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Logging; using Remotely.Shared.Dtos; using Remotely.Shared.Models; using System; using System.Collections.Concurrent; using System.Diagnostics; using System.Linq; using System.Management.Automation; using System.Threading.Tasks; namespace Remotely.Agent.Services; public interface IPsCoreShell { string? SenderConnectionId { get; set; } CommandCompletion GetCompletions(string inputText, int currentIndex, bool? forward); Task WriteInput(string input); } public class PsCoreShell : IPsCoreShell { private static readonly ConcurrentDictionary _sessions = new(); private readonly IConfigService _configService; private readonly ConnectionInfo _connectionInfo; private readonly ILogger _logger; private readonly PowerShell _powershell; private CommandCompletion? _lastCompletion; private string? _lastInputText; public PsCoreShell( IConfigService configService, ILogger logger) { _configService = configService; _connectionInfo = _configService.GetConnectionInfo(); _logger = logger; _powershell = PowerShell.Create(); _powershell.AddScript($@"$VerbosePreference = ""Continue""; $DebugPreference = ""Continue""; $InformationPreference = ""Continue""; $WarningPreference = ""Continue""; $env:DeviceId = ""{_connectionInfo.DeviceID}""; $env:ServerUrl = ""{_connectionInfo.Host}"""); _powershell.Invoke(); } public string? SenderConnectionId { get; set; } // TODO: Turn into cache and factory. public static IPsCoreShell GetCurrent(string senderConnectionId) { if (_sessions.TryGetValue(senderConnectionId, out var session)) { return session; } else { session = Program.Services.GetRequiredService(); session.SenderConnectionId = senderConnectionId; _sessions.AddOrUpdate(senderConnectionId, session, (id, b) => session); return session; } } public CommandCompletion GetCompletions(string inputText, int currentIndex, bool? forward) { if (_lastCompletion is null || inputText != _lastInputText) { _lastInputText = inputText; _lastCompletion = CommandCompletion.CompleteInput(inputText, currentIndex, new(), _powershell); } if (forward.HasValue) { _lastCompletion.GetNextResult(forward.Value); } return _lastCompletion; } public async Task WriteInput(string input) { var deviceId = _configService.GetConnectionInfo().DeviceID; var sw = Stopwatch.StartNew(); try { _powershell.Streams.ClearStreams(); _powershell.Commands.Clear(); _powershell.AddScript(input); var results = _powershell.Invoke(); using var ps = PowerShell.Create(); ps.AddScript("$args[0] | Out-String"); ps.AddArgument(results); var result = await ps.InvokeAsync(); var hostOutput = result.Count > 0 ? $"{result[0].BaseObject}" : string.Empty; var verboseOut = _powershell.Streams.Verbose.ReadAll().Select(x => x.Message); var debugOut = _powershell.Streams.Debug.ReadAll().Select(x => x.Message); var errorOut = _powershell.Streams.Error.ReadAll().Select(x => x.Exception.ToString() + Environment.NewLine + x.ScriptStackTrace); var infoOut = _powershell.Streams.Information.Select(x => x.MessageData.ToString()); var warningOut = _powershell.Streams.Warning.Select(x => x.Message); var standardOut = hostOutput.Split(Environment.NewLine) .Concat(infoOut) .Concat(debugOut) .Concat(verboseOut) .Select(x => $"{x}"); var errorAndWarningOut = errorOut.Concat(warningOut).ToArray(); return new ScriptResultDto() { DeviceID = _configService.GetConnectionInfo().DeviceID, SenderConnectionID = SenderConnectionId, ScriptInput = input, Shell = Shared.Enums.ScriptingShell.PSCore, StandardOutput = standardOut.ToArray(), ErrorOutput = errorAndWarningOut, RunTime = sw.Elapsed, HadErrors = _powershell.HadErrors || errorAndWarningOut.Any() }; } catch (Exception ex) { _logger.LogError(ex, "Error while writing input to PSCore."); return new ScriptResultDto() { DeviceID = deviceId, SenderConnectionID = SenderConnectionId, ScriptInput = input, Shell = Shared.Enums.ScriptingShell.PSCore, StandardOutput = Array.Empty(), ErrorOutput = new[] { "Error while writing input." }, RunTime = sw.Elapsed, HadErrors = true }; } } }