winmr-api/Services/Scanners/ProcessScanner.cs

/* This software is licensed by the MIT License, see LICENSE file */
/*                           Copyright © 2024-2025 Gregory Lirent */

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)>> _wmiTask;

        async private Task<Dictionary<int, (string? Name, string? CommandLine, int ParentPID)>> GetWmiDataAsync()
        {
            var wmi = new Dictionary<int, (string? Name, string? CommandLine, int ParentPID)>();
            await Task.Run(() =>
            {
                try
                {
                    using (var searcher = new ManagementObjectSearcher("SELECT ProcessId, Name, CommandLine, ParentProcessId FROM Win32_Process"))
                    using (var processes = searcher.Get())
                    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, Dictionary<int, (string? Name, string? CommandLine, int ParentPID)> wmiData)
        {
            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 async internal Task<bool> ScanAsync(Dictionary<long, ProcessInfo> data)
        {
            var wmiData = await _wmiTask;
            foreach (var process in await Task.Run(Process.GetProcesses))
            using (process)
            {
                ProcessInfo? info;

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

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

                if (Populate(info, process, wmiData))
                {
                    data.Add(info.PID, info);
                }
                else
                {
                    Provider.Logger.LogWarning($"Skipping process {process.Id} due to incomplete information during population.");
                }
            }

            return true;
        }

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