process_scanner.py

"""
WeInjectDLL - Process Scanner
Multi-mode process enumeration and filtering.
Made by Hardik - https://github.com/ewwhardik
"""

import sys
import time
import struct
import ctypes
import threading
from typing import List, Optional, Callable
from dataclasses import dataclass, field

IS_WINDOWS = sys.platform == "win32"

try:
    import psutil
    HAS_PSUTIL = True
except ImportError:
    HAS_PSUTIL = False

if IS_WINDOWS:
    kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
    psapi    = ctypes.WinDLL("psapi",    use_last_error=True)


@dataclass
class ProcessInfo:
    pid:          int
    name:         str
    path:         str         = ""
    arch:         str         = "?"   # "x64" | "x86" | "?"
    cpu_percent:  float       = 0.0
    mem_mb:       float       = 0.0
    username:     str         = ""
    status:       str         = ""
    modules:      List[str]   = field(default_factory=list)
    is_64bit:     bool        = True
    is_protected: bool        = False   # anti-cheat / protected process
    window_title: str         = ""


class ScanMode:
    ALL           = "All Processes"
    WINDOWED      = "Windowed Only"
    X64           = "64-bit Only"
    X86           = "32-bit Only"
    GAMES         = "Game Processes"
    NAME_FILTER   = "Name Filter"


# Common game-related executable substrings
GAME_HINTS = {
    "game", "steam", "epic", "launcher", "unity", "unreal", "source",
    "client", "engine", "play", "run", "cs2", "valorant", "fortnite",
    "minecraft", "roblox", "overwatch", "apex", "cod", "warzone",
}


def _get_window_title(pid: int) -> str:
    """Return the window title of the main window of `pid`, if any."""
    if not IS_WINDOWS:
        return ""
    result = ctypes.create_unicode_buffer(256)
    found  = ctypes.c_int(0)

    EnumWindowsProc = ctypes.WINFUNCTYPE(ctypes.c_bool,
                                          ctypes.wintypes.HWND,
                                          ctypes.wintypes.LPARAM)

    def callback(hwnd, lParam):
        out_pid = ctypes.c_ulong(0)
        ctypes.windll.user32.GetWindowThreadProcessId(hwnd, ctypes.byref(out_pid))
        if out_pid.value == pid:
            ctypes.windll.user32.GetWindowTextW(hwnd, result, 256)
            if result.value:
                found.value = 1
                return False
        return True

    try:
        ctypes.windll.user32.EnumWindows(EnumWindowsProc(callback), 0)
    except Exception:
        pass
    return result.value


def _is_process_protected(pid: int) -> bool:
    """Detect if a process is a protected process (PPL / anti-cheat)."""
    if not IS_WINDOWS:
        return False
    PROCESS_QUERY_LIMITED = 0x1000
    handle = kernel32.OpenProcess(PROCESS_QUERY_LIMITED, False, pid)
    if not handle:
        return True   # Cannot open → likely protected
    kernel32.CloseHandle(handle)
    return False


def _detect_arch(pid: int) -> str:
    if not IS_WINDOWS:
        return "x64"
    PROCESS_QUERY_LIMITED = 0x1000
    handle = kernel32.OpenProcess(PROCESS_QUERY_LIMITED, False, pid)
    if not handle:
        return "?"
    try:
        wow64 = ctypes.c_int(0)
        kernel32.IsWow64Process(handle, ctypes.byref(wow64))
        return "x86" if wow64.value else "x64"
    finally:
        kernel32.CloseHandle(handle)


class ProcessScanner:
    """
    Enumerate running processes with rich metadata.
    Uses psutil where available; falls back to WinAPI TH32 snapshot.
    """

    def __init__(self):
        self._cache: List[ProcessInfo] = []
        self._lock  = threading.Lock()

    # ── Public API ──────────────────────────────────────────────────────────

    def scan(
        self,
        mode:        str = ScanMode.ALL,
        name_filter: str = "",
        progress_cb: Optional[Callable] = None,
    ) -> List[ProcessInfo]:
        """Return a fresh list of ProcessInfo objects."""
        results: List[ProcessInfo] = []

        if HAS_PSUTIL:
            results = self._scan_psutil(progress_cb)
        else:
            results = self._scan_winapi(progress_cb)

        results = self._apply_filter(results, mode, name_filter)
        with self._lock:
            self._cache = results
        return results

    def get_cached(self) -> List[ProcessInfo]:
        with self._lock:
            return list(self._cache)

    def refresh_modules(self, info: ProcessInfo) -> List[str]:
        """Populate `info.modules` in-place and return the list."""
        mods = self._get_modules(info.pid)
        info.modules = mods
        return mods

    # ── Scan backends ────────────────────────────────────────────────────────

    def _scan_psutil(self, progress_cb) -> List[ProcessInfo]:
        procs = []
        all_p = list(psutil.process_iter(
            ["pid", "name", "exe", "username", "status",
             "cpu_percent", "memory_info"]
        ))
        total = max(len(all_p), 1)

        for i, p in enumerate(all_p):
            if progress_cb and i % 20 == 0:
                progress_cb(f"Scanning… ({i}/{total})", int(i / total * 100))
            try:
                info = self._psutil_to_info(p)
                procs.append(info)
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                continue

        if progress_cb:
            progress_cb("Scan complete.", 100)
        return procs

    @staticmethod
    def _psutil_to_info(p) -> ProcessInfo:
        try:
            mem_mb = p.info["memory_info"].rss / (1024 * 1024) if p.info.get("memory_info") else 0
        except Exception:
            mem_mb = 0
        try:
            arch = _detect_arch(p.info["pid"])
        except Exception:
            arch = "?"
        return ProcessInfo(
            pid         = p.info["pid"],
            name        = p.info["name"] or "",
            path        = p.info.get("exe") or "",
            arch        = arch,
            cpu_percent = p.info.get("cpu_percent") or 0.0,
            mem_mb      = mem_mb,
            username    = p.info.get("username") or "",
            status      = p.info.get("status") or "",
            is_64bit    = (arch == "x64"),
        )

    def _scan_winapi(self, progress_cb) -> List[ProcessInfo]:
        """Fallback: use CreateToolhelp32Snapshot."""
        if not IS_WINDOWS:
            return self._demo_processes()

        TH32CS_SNAPPROCESS = 0x00000002

        class PROCESSENTRY32W(ctypes.Structure):
            _fields_ = [
                ("dwSize",              ctypes.c_ulong),
                ("cntUsage",            ctypes.c_ulong),
                ("th32ProcessID",       ctypes.c_ulong),
                ("th32DefaultHeapID",   ctypes.POINTER(ctypes.c_ulong)),
                ("th32ModuleID",        ctypes.c_ulong),
                ("cntThreads",          ctypes.c_ulong),
                ("th32ParentProcessID", ctypes.c_ulong),
                ("pcPriClassBase",      ctypes.c_long),
                ("dwFlags",             ctypes.c_ulong),
                ("szExeFile",           ctypes.c_wchar * 260),
            ]

        snap = kernel32.CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0)
        procs = []
        pe = PROCESSENTRY32W()
        pe.dwSize = ctypes.sizeof(PROCESSENTRY32W)

        if kernel32.Process32FirstW(snap, ctypes.byref(pe)):
            while True:
                arch = _detect_arch(pe.th32ProcessID)
                procs.append(ProcessInfo(
                    pid     = pe.th32ProcessID,
                    name    = pe.szExeFile,
                    arch    = arch,
                    is_64bit= (arch == "x64"),
                ))
                pe.dwSize = ctypes.sizeof(PROCESSENTRY32W)
                if not kernel32.Process32NextW(snap, ctypes.byref(pe)):
                    break
        kernel32.CloseHandle(snap)
        return procs

    @staticmethod
    def _demo_processes() -> List[ProcessInfo]:
        """Return mock data for non-Windows demo mode."""
        demo = [
            ProcessInfo(1,    "System",           "/",           "x64", 0.1, 4.0,   "SYSTEM",  "running"),
            ProcessInfo(4,    "smss.exe",          "C:/Windows",  "x64", 0.0, 1.2,   "SYSTEM",  "running"),
            ProcessInfo(580,  "csrss.exe",         "C:/Windows",  "x64", 0.2, 3.5,   "SYSTEM",  "running"),
            ProcessInfo(900,  "winlogon.exe",      "C:/Windows",  "x64", 0.0, 5.1,   "SYSTEM",  "running"),
            ProcessInfo(1234, "explorer.exe",      "C:/Windows",  "x64", 1.2, 64.0,  "User",    "running"),
            ProcessInfo(2048, "notepad.exe",       "C:/Windows",  "x64", 0.0, 8.0,   "User",    "running"),
            ProcessInfo(3000, "chrome.exe",        "C:/Program Files","x64",5.4,320.0,"User",   "running"),
            ProcessInfo(3100, "steam.exe",         "C:/Program Files","x64",2.1,200.0,"User",   "running"),
            ProcessInfo(3200, "cs2.exe",           "C:/SteamApps","x64",35.0,2048.0, "User",   "running"),
            ProcessInfo(3300, "RobloxPlayerBeta.exe","C:/Users",  "x64",10.0,512.0,  "User",   "running"),
            ProcessInfo(3400, "Discord.exe",       "C:/Users",    "x64", 2.3,180.0,  "User",   "running"),
            ProcessInfo(4000, "svchost.exe",       "C:/Windows",  "x64", 0.5, 25.0,  "NETWORK SERVICE","running"),
            ProcessInfo(4100, "taskmgr.exe",       "C:/Windows",  "x64", 0.8, 15.0,  "User",   "running"),
            ProcessInfo(4200, "code.exe",          "C:/Program Files","x64",4.2,210.0,"User",  "running"),
            ProcessInfo(4300, "python.exe",        "C:/Python",   "x64", 3.1, 45.0,  "User",   "running"),
        ]
        return demo

    # ── Module enumeration ────────────────────────────────────────────────────

    def _get_modules(self, pid: int) -> List[str]:
        if not IS_WINDOWS:
            return ["kernel32.dll", "ntdll.dll", "user32.dll", "test.dll"]
        PROCESS_VM_READ   = 0x0010
        PROCESS_QUERY_INFO= 0x0400
        handle = kernel32.OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_INFO, False, pid)
        if not handle:
            return []
        try:
            hMods    = (ctypes.wintypes.HMODULE * 1024)()
            cb_needed= ctypes.c_ulong()
            psapi.EnumProcessModulesEx(handle, hMods, ctypes.sizeof(hMods),
                                       ctypes.byref(cb_needed), 0x03)
            count  = cb_needed.value // ctypes.sizeof(ctypes.wintypes.HMODULE)
            mods   = []
            name_b = ctypes.create_unicode_buffer(260)
            for i in range(count):
                psapi.GetModuleFileNameExW(handle, hMods[i], name_b, 260)
                mods.append(name_b.value)
            return mods
        finally:
            kernel32.CloseHandle(handle)

    # ── Filter logic ─────────────────────────────────────────────────────────

    @staticmethod
    def _apply_filter(procs: List[ProcessInfo], mode: str, name_filter: str) -> List[ProcessInfo]:
        filt = procs

        if mode == ScanMode.X64:
            filt = [p for p in filt if p.arch == "x64"]
        elif mode == ScanMode.X86:
            filt = [p for p in filt if p.arch == "x86"]
        elif mode == ScanMode.GAMES:
            filt = [p for p in filt
                    if any(h in p.name.lower() for h in GAME_HINTS)]
        elif mode == ScanMode.WINDOWED:
            filt = [p for p in filt if p.window_title]

        if name_filter:
            q = name_filter.lower()
            filt = [p for p in filt
                    if q in p.name.lower() or
                       q in str(p.pid) or
                       q in p.path.lower()]

        return sorted(filt, key=lambda p: p.name.lower())