using System.Management.Automation; using System.Collections.Concurrent; using System.Management.Automation.Language; using System.Management.Automation.Runspaces; using MaksIT.UScheduler.Shared; using MaksIT.UScheduler.Shared.Helpers; using MaksIT.UScheduler.Shared.Extensions; namespace MaksIT.UScheduler.Services; /// /// Service responsible for executing and managing PowerShell scripts. /// Provides parallel script execution using a RunspacePool and supports /// signature verification for signed scripts. /// public sealed class PSScriptService : IPSScriptService { private readonly ILogger _logger; private readonly ILoggerFactory _loggerFactory; private readonly ConcurrentDictionary _runningScripts = new(); private readonly RunspacePool _runspacePool; private bool _disposed; /// /// Initializes a new instance of the class. /// Creates a RunspacePool for parallel script execution. /// /// The logger instance for this service. /// The logger factory for creating script-specific loggers. public PSScriptService( ILogger logger, ILoggerFactory loggerFactory ) { _logger = logger; _loggerFactory = loggerFactory; // Create a RunspacePool to allow parallel script execution _runspacePool = RunspaceFactory.CreateRunspacePool(1, Environment.ProcessorCount); _runspacePool.Open(); _logger.LogInformation($"RunspacePool opened with max {Environment.ProcessorCount} runspaces"); } /// /// Executes a PowerShell script asynchronously with optional signature verification. /// Automatically passes Automated and CurrentDateTimeUtc parameters to the script. /// /// The path to the PowerShell script to execute. /// If true, validates the script's Authenticode signature before execution. /// Cancellation token to signal when script execution should stop. /// When true, the script honors its schedule. A manual launch passes false. /// A task representing the asynchronous operation. public async Task RunScriptAsync(string scriptPath, bool signed, CancellationToken stoppingToken, bool automated = true) { // Resolve relative paths against application base directory var resolvedPath = PathHelper.ResolvePath(scriptPath); _logger.LogInformation($"Preparing to run script {resolvedPath}"); if (GetRunningScriptTasks().Contains(resolvedPath)) { _logger.LogInformation($"PowerShell script {resolvedPath} is already running"); return; } if (!File.Exists(resolvedPath)) { _logger.LogError($"Script file {resolvedPath} does not exist"); if (!automated) LogManualOutcome(resolvedPath, $"Script file {resolvedPath} does not exist"); return; } if (!EnsureDependenciesUnblocked(resolvedPath)) { _logger.LogError($"Script or dependencies for {resolvedPath} could not be unblocked. Aborting execution."); if (!automated) LogManualOutcome(resolvedPath, $"Script or dependencies for {resolvedPath} could not be unblocked."); return; } var ps = PowerShell.Create(); ps.RunspacePool = _runspacePool; if (!_runningScripts.TryAdd(resolvedPath, ps)) { _logger.LogWarning($"Script {resolvedPath} was already added by another thread"); ps.Dispose(); return; } ILogger? scriptLogger = null; try { scriptLogger = _loggerFactory.CreateFolderLogger(resolvedPath); var scriptPolicy = signed ? "AllSigned" : "Unrestricted"; ps.AddScript($"Set-ExecutionPolicy -Scope Process -ExecutionPolicy {scriptPolicy}"); await Task.Factory.FromAsync(ps.BeginInvoke(), ps.EndInvoke); if (signed) { ps.Commands.Clear(); ps.AddScript($"Get-AuthenticodeSignature \"{resolvedPath}\""); var signatureResults = await Task.Factory.FromAsync(ps.BeginInvoke(), ps.EndInvoke); if (signatureResults.Count == 0 || ((Signature)signatureResults[0].BaseObject).Status != SignatureStatus.Valid) { _logger.LogWarning($"Script {resolvedPath} signature is invalid. Correct and restart the service."); if (!automated) scriptLogger.LogWarning($"[PS Output] [Warning] Script {resolvedPath} signature is invalid."); return; } } stoppingToken.ThrowIfCancellationRequested(); _logger.LogInformation($"Invoking: {resolvedPath}"); ps.Commands.Clear(); var myCommand = new Command(resolvedPath); var currentDateTimeUtcString = DateTime.UtcNow.ToString("o"); myCommand.Parameters.Add(new CommandParameter("Automated", automated)); myCommand.Parameters.Add(new CommandParameter("CurrentDateTimeUtc", currentDateTimeUtcString)); ps.Commands.Commands.Add(myCommand); _logger.LogInformation($"Added parameters: Automated={automated}, CurrentDateTimeUtc={currentDateTimeUtcString}"); await InvokeAndLogAsync(ps, scriptLogger); } catch (OperationCanceledException) { _logger.LogInformation($"Stopping script {resolvedPath} due to cancellation request"); if (!automated) scriptLogger?.LogInformation("[PS Output] [Info] Script canceled"); } catch (Exception ex) { _logger.LogError($"Error running script {resolvedPath}: {ex.Message}"); if (!automated) scriptLogger?.LogError($"[PS Error] {ex.Message}"); } finally { if (!automated) scriptLogger?.LogInformation(ScriptRunMarkers.FinishedMessage); TerminateScript(resolvedPath); _logger.LogInformation($"Script {resolvedPath} completed and removed from running scripts"); } } /// /// Writes one failure into the script log and marks the manual run finished so the progress window can close its watch. /// private void LogManualOutcome(string resolvedPath, string error) { var scriptLogger = _loggerFactory.CreateFolderLogger(resolvedPath); scriptLogger.LogError($"[PS Error] {error}"); scriptLogger.LogInformation(ScriptRunMarkers.FinishedMessage); } /// /// Runs the current command and writes each output, error, and host line as it arrives. /// private static async Task InvokeAndLogAsync(PowerShell ps, ILogger scriptLogger) { ps.Streams.ClearStreams(); var output = new PSDataCollection(); void OnOutput(object? _, DataAddedEventArgs e) { if (e.Index < 0 || e.Index >= output.Count) return; var text = output[e.Index]?.BaseObject?.ToString(); if (!string.IsNullOrEmpty(text)) scriptLogger.LogInformation($"[PS Output] {text}"); } void OnError(object? _, DataAddedEventArgs e) { if (e.Index < 0 || e.Index >= ps.Streams.Error.Count) return; scriptLogger.LogError($"[PS Error] {ps.Streams.Error[e.Index]}"); } void OnInformation(object? _, DataAddedEventArgs e) { if (e.Index < 0 || e.Index >= ps.Streams.Information.Count) return; var text = ps.Streams.Information[e.Index].MessageData?.ToString(); if (!string.IsNullOrEmpty(text)) scriptLogger.LogInformation($"[PS Output] {text}"); } output.DataAdded += OnOutput; ps.Streams.Error.DataAdded += OnError; ps.Streams.Information.DataAdded += OnInformation; try { var pending = ps.BeginInvoke(null, output); await Task.Factory.FromAsync(pending, ar => ps.EndInvoke(ar)); } finally { output.DataAdded -= OnOutput; ps.Streams.Error.DataAdded -= OnError; ps.Streams.Information.DataAdded -= OnInformation; output.Dispose(); } } /// /// Gets a list of script paths that are currently being executed. /// /// A list of script paths currently running. public List GetRunningScriptTasks() { _logger.LogInformation($"Retrieving running script tasks. Current count: {_runningScripts.Count}"); return _runningScripts.Keys.ToList(); } /// /// Terminates a running PowerShell script by its path. /// Stops the script execution and disposes of its PowerShell instance. /// /// The path of the script to terminate. public void TerminateScript(string scriptPath) { _logger.LogInformation($"Attempting to terminate script {scriptPath}"); if (_runningScripts.TryRemove(scriptPath, out var ps)) { ps.Stop(); ps.Dispose(); _logger.LogInformation($"Script {scriptPath} terminated"); } else { _logger.LogWarning($"Failed to terminate script {scriptPath}. Script not found."); } } /// /// Terminates all currently running PowerShell scripts. /// public void TerminateAllScripts() { _logger.LogInformation("Terminating all running scripts"); foreach (var scriptPath in _runningScripts.Keys.ToList()) { TerminateScript(scriptPath); } } /// /// Releases all resources used by the . /// Terminates all running scripts and closes the RunspacePool. /// public void Dispose() { if (_disposed) return; _disposed = true; TerminateAllScripts(); _runspacePool?.Close(); _runspacePool?.Dispose(); _logger.LogInformation("RunspacePool disposed"); } /// /// Recursively scans a PowerShell script for module and dot-sourced dependencies and unblocks them. /// /// The entry script path. /// True if all scripts and dependencies were unblocked; false otherwise. private bool EnsureDependenciesUnblocked(string scriptPath) { var visited = new HashSet(StringComparer.OrdinalIgnoreCase); var queue = new Queue(); queue.Enqueue(scriptPath); bool allUnblocked = true; while (queue.Count > 0) { var current = queue.Dequeue(); if (!visited.Add(current)) continue; if (!TryUnblockScript(current)) { _logger.LogError($"Failed to unblock dependency: {current}"); allUnblocked = false; } // Scan for dependencies try { if (!File.Exists(current)) continue; var currentDir = Path.GetDirectoryName(current); var ast = Parser.ParseFile(current, out var tokens, out var errors); // Handle 'using module' statements (UsingStatementAst) foreach (var usingAst in ast.FindAll(a => a is UsingStatementAst, true)) { var usingStmt = (UsingStatementAst)usingAst; if (usingStmt.UsingStatementKind == UsingStatementKind.Module && usingStmt.Name != null) { var depName = usingStmt.Name.Value; var depPath = ResolveModulePath(depName, currentDir); if (!string.IsNullOrEmpty(depPath)) queue.Enqueue(depPath); } } // Handle Import-Module commands foreach (var cmdAst in ast.FindAll(a => a is CommandAst, true)) { var cmd = (CommandAst)cmdAst; var name = cmd.GetCommandName(); if (string.Equals(name, "Import-Module", StringComparison.OrdinalIgnoreCase)) { foreach (var arg in cmd.CommandElements.Skip(1)) { var depName = arg.ToString().Trim('"', '\'', ' '); // Skip parameters like -Name, -Force, etc. if (depName.StartsWith("-", StringComparison.Ordinal)) continue; var depPath = ResolveModulePath(depName, currentDir); if (!string.IsNullOrEmpty(depPath)) queue.Enqueue(depPath); } } } // Handle dot-sourcing: . ./file.ps1 foreach (var cmdAst in ast.FindAll(a => a is CommandAst, true)) { var cmd = (CommandAst)cmdAst; if (cmd.InvocationOperator == TokenKind.Dot) { var arg = cmd.CommandElements.FirstOrDefault(); if (arg != null && !string.IsNullOrEmpty(currentDir)) { var depName = arg.ToString().Trim('"', '\'', ' '); var depPath = Path.Combine(currentDir, depName); if (File.Exists(depPath)) queue.Enqueue(depPath); } } } } catch (Exception ex) { _logger.LogWarning($"Dependency scan failed for {current}: {ex.Message}"); } } return allUnblocked; } /// /// Attempts to resolve a module path from a module name or path. /// /// Module name or path. /// Base directory for relative paths. /// Resolved module file path or null. private string? ResolveModulePath(string moduleName, string? baseDir) { // If it's a path, resolve relative to baseDir if (moduleName.EndsWith(".psm1", StringComparison.OrdinalIgnoreCase) || moduleName.EndsWith(".psd1", StringComparison.OrdinalIgnoreCase)) { if (Path.IsPathRooted(moduleName)) { if (File.Exists(moduleName)) return moduleName; } else if (!string.IsNullOrEmpty(baseDir)) { var path = Path.Combine(baseDir, moduleName); if (File.Exists(path)) return path; } } // Try to find module in same directory if (!string.IsNullOrEmpty(baseDir)) { var psm1 = Path.Combine(baseDir, moduleName + ".psm1"); if (File.Exists(psm1)) return psm1; var psd1 = Path.Combine(baseDir, moduleName + ".psd1"); if (File.Exists(psd1)) return psd1; // Try subfolder with module name (common module structure) var subPsm1 = Path.Combine(baseDir, moduleName, moduleName + ".psm1"); if (File.Exists(subPsm1)) return subPsm1; var subPsd1 = Path.Combine(baseDir, moduleName, moduleName + ".psd1"); if (File.Exists(subPsd1)) return subPsd1; } return null; } /// /// Attempts to unblock a downloaded script by removing the Zone.Identifier alternate data stream. /// This is equivalent to right-clicking a file and selecting "Unblock" in Windows. /// /// The path to the script to unblock. /// True if the script was successfully unblocked or was not blocked; false if unblocking failed. private bool TryUnblockScript(string scriptPath) { try { var zoneIdentifier = scriptPath + ":Zone.Identifier"; if (File.Exists(zoneIdentifier)) { File.Delete(zoneIdentifier); _logger.LogInformation($"Unblocked script {scriptPath} by removing Zone.Identifier."); } return true; } catch (Exception ex) { _logger.LogWarning($"Failed to unblock script {scriptPath}: {ex.Message}"); return false; } } }