insight/src/Agent/Insight.Agent/Network/Handlers/ConsoleHandler.cs

using Insight.Domain.Interfaces;
using Insight.Domain.Messages;
using Insight.Domain.Messages.Agent;
using System.Management.Automation;
using System.Management.Automation.Runspaces;

namespace Insight.Agent.Network.Handlers;

public class ConsoleHandler : IMessageHandler<AgentSession>
{
    public async ValueTask HandleAsync<TMessage>(AgentSession sender, TMessage message, CancellationToken cancellationToken) where TMessage : IMessage
    {
        if (message is Proxy<ConsoleQueryRequest> consoleQueryRequest)
        {
            await OnConsoleQueryRequestAsync(sender, consoleQueryRequest, cancellationToken);
        }
    }

    private async ValueTask OnConsoleQueryRequestAsync(AgentSession sender, Proxy<ConsoleQueryRequest> consoleQueryRequest, CancellationToken cancellationToken)
    {
        var result = await QueryScriptAsync(consoleQueryRequest.Message.Query);

        await sender.SendAsync(new Proxy<ConsoleQuery>
        {
            RequestId = consoleQueryRequest.RequestId,
            HostId = consoleQueryRequest.HostId,
            Message = new ConsoleQuery
            {
                Data = result.Data,
                Errors = result.Errors,
                HadErrors = result.HadErrors
            }
        }, cancellationToken);
    }

    private static async Task<QueryResult> QueryScriptAsync(string query)
    {
        var result = new QueryResult();
        var errors = new List<string>();

        try
        {
            using var runspace = RunspaceFactory.CreateRunspace();
            runspace.Open();
            runspace.SessionStateProxy.LanguageMode = PSLanguageMode.FullLanguage;

            using var ps = PowerShell.Create(runspace);
            ps.AddScript("Set-ExecutionPolicy unrestricted -Scope Process");
            ps.AddScript(query);
            ps.AddCommand("ConvertTo-Json"); // -Depth 10

            result.Query = query;

            var queryResult = await ps.InvokeAsync();

            if (ps.HadErrors)
            {
                result.HadErrors = true;
                errors.AddRange(ps.Streams.Error.Select(e => e.ToString()));
            }
            else
            {
                result.Data = queryResult[0].ToString();

                //if (string.IsNullOrWhiteSpace(jsonString)) return result;

                //if (jsonString.TrimStart().StartsWith("[")) // It's an array
                //{
                //    result.IsArray = true;

                //    var deserialized = JsonSerializer.Deserialize<List<Dictionary<string, object?>>>(jsonString, new JsonSerializerOptions { Encoder = JavaScriptEncoder.UnsafeRelaxedJsonEscaping });
                //    if (deserialized is null) return result;

                //    result.Data.AddRange(deserialized);

                //    //Console.WriteLine("Deserialized to List<Dictionary<string, object>>");
                //}
                //else
                //{
                //    if (jsonString.TrimStart().StartsWith("{") is false) // It's an object
                //    {
                //        result.IsString = true;
                //        result.Data.Add(new Dictionary<string, object?> { { query, jsonString.Trim('"') } });
                //    }
                //    else
                //    {
                //        var deserialized = JsonSerializer.Deserialize<Dictionary<string, object?>>(jsonString, new JsonSerializerOptions { Encoder = JavaScriptEncoder.UnsafeRelaxedJsonEscaping });
                //        if (deserialized is null) return result;

                //        result.Data.Add(deserialized);

                //        //Console.WriteLine("Deserialized to Dictionary<string, object>");
                //    }
                //}
            }
        }
        catch (Exception ex)
        {
            result.HadErrors = true;
            errors.Add(ex.Message);
        }

        result.Errors = string.Join("\n", errors);
        return result;
    }
}

public class QueryResult
{
    public bool HadErrors { get; set; }
    public string? Query { get; set; }
    public string? Data { get; set; }
    public string? Errors { get; set; }
}