winmr-api/Services/Scanners/ProcessScanner.cs

// File: Services/Scanners/ProcessScanner.cs

using System.Diagnostics;
using System.Management;
using WebmrAPI.Models;
using WebmrAPI.Utils;

namespace WebmrAPI.Services.Scanners
{
    public class ProcessScanner : AbstractCpuScanner<ProcessInfo>
    {
        override public ScanTarget Target { get => ScanTarget.Processes; }

        private Task<Dictionary<int, (string? Name, string? CommandLine, int ParentPID)>> _wmi;

        private Dictionary<int, (string? Name, string? CommandLine, int ParentPID)> WmiData
        {
            get => _wmi.Result;
        }

        async private Task<Dictionary<int, (string? Name, string? CommandLine, int ParentPID)>> GetWmiDataAsync()
        {
            var wmi = new Dictionary<int, (string? Name, string? CommandLine, int ParentPID)>();

            try
            {
                using (var searcher = new ManagementObjectSearcher("SELECT ProcessId, Name, CommandLine, ParentProcessId FROM Win32_Process"))
                using (var processes = await Task.Run(() => { return searcher.Get(); }))
                await Task.Run(() =>
                {
                    foreach (var obj in processes)
                    {
                        int pid = Convert.ToInt32(obj["ProcessId"]);
                        int ppid = Convert.ToInt32(obj["ParentProcessId"]);
                        string? name = obj["Name"]?.ToString();
                        string? cmd = obj["CommandLine"]?.ToString();

                        wmi[pid] = (name, cmd, ppid);
                    }
                });
            }
            catch (Exception ex)
            {
                Provider.Logger.LogError(ex, "Failed to retrieve WMI process data for all processes.");
            }

            return wmi;
        }

        private bool Populate(ProcessInfo info, Process process)
        {
            info.TotalProcessorTime = process.TotalProcessorTime;
            info.MemorySize = (ulong)process.WorkingSet64;
            info.ThreadCount = process.Threads.Count;

            try
            {
                if (info.MemoryAddress == 0 && process.MainModule != null)
                {
                    info.FileName = process.MainModule.FileName;
                    info.MemoryAddress = process.MainModule.BaseAddress.ToInt64();
                }
            }
            catch (Exception ex)
            {
                Provider.Logger.LogWarning($"Process {process.ProcessName} (PID: {process.Id}), an error occurred while getting process address: {ex.Message}");
            }

            if (info.MemoryAddress != 0)
            {
                try
                {
                    if (WmiData.TryGetValue(process.Id, out var entry))
                    {
                        info.Name = entry.Name;
                        info.CommandLine = entry.CommandLine;
                        info.ParentPID = entry.ParentPID;
                    }
                    else
                    {
                        info.Name = process.ProcessName;
                        Provider.Logger.LogDebug($"WMI data not found for PID {process.Id}. Falling back to Process.ProcessName.");
                    }

                    if (info.StartTime == null)
                    {
                        try { info.StartTime = process.StartTime; }
                        catch (Exception ex)
                        {
                            Provider.Logger.LogDebug($"Could not get StartTime for PID {process.Id}: {ex.Message}");
                        }
                    }
                    process.Refresh();
                    info.Status = process.Responding ? ProcessInfo.ProcessStatus.Running : ProcessInfo.ProcessStatus.NotResponding;
                }
                catch (Exception ex) 
                { 
                    info.Status = ProcessInfo.ProcessStatus.Undefined;
                    Provider.Logger.LogDebug($"Could not get Status for PID {process.Id}: {ex.Message}");
                }
                return true;
            }
            return false;
        }

        override internal bool Scan(out Dictionary<long, ProcessInfo>? data)
        {
            data = new();

            foreach (var process in Process.GetProcesses())
            using (process)
            {
                ProcessInfo? info;

                if (process.Id == 0 || process.Id == 4)
                {
                    process.Dispose();
                    continue;
                }

                if (GetFromCacheOrNew(process.Id, out info))
                {
                    info.CpuUsage = CalcCpuUsage(info.ProcessorTime, Container.Elapsed.TotalMilliseconds);
                }
                else
                {
                    info.PID = process.Id;
                }

                try
                {
                    if (Populate(info, process))
                    {
                        data.Add(info.PID, info);
                    }
                } catch (Exception) { }
            }

            return true;
        }

        public ProcessScanner(IScanProvider scanner, LazyConcurrentContainer<ProcessInfo> container)
            : base(scanner, container)
        {
            _wmi = GetWmiDataAsync();
        }
    }
}