using System.Collections.Concurrent; using System.Management.Automation; using System.Management.Automation.Language; using System.Management.Automation.Runspaces; 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. /// A task representing the asynchronous operation. public async Task RunScriptAsync(string scriptPath, bool signed, CancellationToken stoppingToken) { // 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"); return; } if (!EnsureDependenciesUnblocked(resolvedPath)) { _logger.LogError($"Script or dependencies for {resolvedPath} could not be unblocked. Aborting execution."); 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; } try { // Set execution policy 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."); 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", true)); myCommand.Parameters.Add(new CommandParameter("CurrentDateTimeUtc", currentDateTimeUtcString)); ps.Commands.Commands.Add(myCommand); _logger.LogInformation($"Added parameters: Automated=true, CurrentDateTimeUtc={currentDateTimeUtcString}"); // Execute asynchronously var outputResults = await Task.Factory.FromAsync(ps.BeginInvoke(), ps.EndInvoke); #region Script Output Logging var scriptLogger = _loggerFactory.CreateFolderLogger(resolvedPath); // Log standard output if (outputResults != null && outputResults.Count > 0) { foreach (var outputItem in outputResults) { scriptLogger.LogInformation($"[PS Output] {outputItem}"); } } // Log errors if (ps.Streams.Error.Count > 0) { foreach (var errorItem in ps.Streams.Error) { scriptLogger.LogError($"[PS Error] {errorItem}"); } } #endregion } catch (OperationCanceledException) { _logger.LogInformation($"Stopping script {resolvedPath} due to cancellation request"); } catch (Exception ex) { _logger.LogError($"Error running script {resolvedPath}: {ex.Message}"); } finally { TerminateScript(resolvedPath); _logger.LogInformation($"Script {resolvedPath} completed and removed from running scripts"); } } /// /// 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; } } }