PYAS_Process.py

from PYAS_Diagnostics import log_exception
import os
import re
import io
import time
import ctypes
import ctypes.wintypes
from PYAS_WinAPI import (
    IO_COUNTERS,
    MIB_TCPROW_OWNER_PID,
    PROCESSENTRY32W,
    PROCESS_BASIC_INFORMATION,
    PYAS_PROCESS_TERMINATE_ACCESS,
    UNICODE_STRING,
)


class ProcessMixin:
    def device_path_to_drive(self, path):
        if not path:
            return ""

        for d in range(65, 91):
            drive = f"{chr(d)}:"
            buf = ctypes.create_unicode_buffer(1024)

            if self.kernel32.QueryDosDeviceW(drive, buf, 1024) and path.startswith(buf.value):
                return path.replace(buf.value, drive, 1)

        return path

    def get_exe_info(self, pid):
        name, file_path = None, None
        h = self.kernel32.OpenProcess(0x1000, False, pid)

        if h:
            try:
                buf = ctypes.create_unicode_buffer(1024)
                size = ctypes.wintypes.DWORD(1024)

                if self.kernel32.QueryFullProcessImageNameW(h, 0, buf, ctypes.byref(size)):
                    file_path = self.norm_path(buf.value)

                    if file_path:
                        name = os.path.basename(file_path)

                elif self.psapi.GetProcessImageFileNameW(h, buf, 1024):
                    file_path = self.norm_path(self.device_path_to_drive(buf.value))

                    if file_path:
                        name = os.path.basename(file_path)

            finally:
                self.kernel32.CloseHandle(h)

        return name, file_path

    def get_process_file(self, h_process):
        buf = ctypes.create_unicode_buffer(1024)

        if self.psapi.GetProcessImageFileNameW(h_process, buf, 1024):
            return self.norm_path(self.device_path_to_drive(buf.value))

        return ""

    def get_process_cmdline(self, h):
        try:
            pbi = PROCESS_BASIC_INFORMATION()
            returned = ctypes.wintypes.ULONG()

            if (
                self.ntdll.NtQueryInformationProcess(
                    h, 0, ctypes.byref(pbi), ctypes.sizeof(pbi), ctypes.byref(returned)
                )
                != 0
            ):
                return ""

            if returned.value != ctypes.sizeof(pbi) or not pbi.PebBaseAddress:
                return ""

            def read_exact(address, destination):
                transferred = ctypes.c_size_t()
                capacity = ctypes.sizeof(destination)
                return (
                    bool(
                        self.kernel32.ReadProcessMemory(
                            h,
                            ctypes.c_void_p(address),
                            ctypes.byref(destination),
                            capacity,
                            ctypes.byref(transferred),
                        )
                    )
                    and transferred.value == capacity
                )

            pointer_size = ctypes.sizeof(ctypes.c_void_p)
            parameters = ctypes.c_void_p()

            if (
                not read_exact(
                    int(pbi.PebBaseAddress) + (0x20 if pointer_size == 8 else 0x10), parameters
                )
                or not parameters.value
            ):
                return ""

            descriptor = UNICODE_STRING()

            if not read_exact(parameters.value + (0x70 if pointer_size == 8 else 0x40), descriptor):
                return ""

            length = int(descriptor.Length)

            if not descriptor.Buffer or not length:
                return ""

            if length % 2 or length > descriptor.MaximumLength:
                raise ValueError("Invalid remote UTF-16 command-line length")

            buffer = (ctypes.c_ubyte * length)()

            if not read_exact(int(descriptor.Buffer), buffer):
                return ""

            return bytes(buffer).decode("utf-16-le", errors="surrogatepass")
        except Exception:
            log_exception("ProcessMixin.get_process_cmdline")
            return ""

    def extract_paths_from_cmdline(self, cmdline):
        if not cmdline:
            return []

        found = []

        for m in re.finditer(r'"([^"]+)"|\'([^\']+)\'', cmdline):
            path = m.group(1) or m.group(2)

            if re.match(r"^[A-Za-z]:\\|^\\\\", path):
                found.append(path)

        if not found:
            m = re.search(
                r'([A-Za-z]:\\[^\*?"<>\|]+\.(?:exe|dll|bat|cmd|vbs|sys|com|pif))',
                cmdline,
                re.IGNORECASE,
            )

            if m:
                found.append(m.group(1))

        if not found:
            argc = ctypes.c_int(0)
            argv = self.shell32.CommandLineToArgvW(cmdline, ctypes.byref(argc))

            if argv:
                args = [argv[i] for i in range(argc.value)]
                self.kernel32.LocalFree(ctypes.cast(argv, ctypes.c_void_p))
                patterns = [
                    r'([A-Za-z]:\\[^"\']+)',
                    r'(\\\\[^"\']+)',
                    r'(\.\\[^"\']+)',
                    r'(\./[^"\']+)',
                    r'([A-Za-z]:/[^"\']+)',
                    r"([^\s]*\\[^\s]+)",
                ]

                for arg in args:
                    for p in patterns:
                        for match in re.finditer(p, arg):
                            found.append(match.group(1).strip('"').strip("'"))

        return list(dict.fromkeys([p.strip('"').strip("'") for p in found]))

    def _enum_processes(self):
        pe = PROCESSENTRY32W()
        pe.dwSize = ctypes.sizeof(PROCESSENTRY32W)
        snapshot = self.kernel32.CreateToolhelp32Snapshot(0x00000002, 0)

        if snapshot in (-1, 0xFFFFFFFF, 0xFFFFFFFFFFFFFFFF):
            return

        try:
            if self.kernel32.Process32FirstW(snapshot, ctypes.byref(pe)):
                while True:
                    yield pe.th32ProcessID, pe.szExeFile

                    if not self.kernel32.Process32NextW(snapshot, ctypes.byref(pe)):
                        break
        finally:
            self.kernel32.CloseHandle(snapshot)

    def get_process_list(self):
        result = []

        for pid, exe_name in self._enum_processes():
            name, file_path = None, None

            if pid > 4:
                name, file_path = self.get_exe_info(pid)

            if not name:
                try:
                    name = exe_name
                except Exception:
                    log_exception("PYAS_Process.ProcessMixin.get_process_list:140")
                    name = "Unknown"

            result.append({"pid": pid, "name": name, "path": file_path or "None"})

        return result

    def get_process_list_pids(self):
        return {pid for pid, _ in self._enum_processes()}

    def kill_process(self, pid, expected_path=None):
        try:
            h = self.kernel32.OpenProcess(PYAS_PROCESS_TERMINATE_ACCESS, False, pid)

            if h:
                try:
                    return self._terminate_process_handle(h, pid=pid, expected_path=expected_path)
                finally:
                    self.kernel32.CloseHandle(h)

        except Exception as e:
            log_exception("PYAS_Process.ProcessMixin.kill_process:159")
            self.write_log(
                "WARN", "kill_process", pid=pid, detail=str(e), operate=True, success=False
            )

        return False

    def get_connections_list(self):
        connections = set()

        try:
            size = ctypes.wintypes.DWORD()

            if self.iphlpapi.GetExtendedTcpTable(None, ctypes.byref(size), True, 2, 5, 0) != 122:
                return connections

            buf = ctypes.create_string_buffer(size.value)

            if self.iphlpapi.GetExtendedTcpTable(buf, ctypes.byref(size), True, 2, 5, 0) != 0:
                return connections

            num_entries = ctypes.cast(buf, ctypes.POINTER(ctypes.wintypes.DWORD)).contents.value

            for i in range(num_entries):
                row = MIB_TCPROW_OWNER_PID.from_address(
                    ctypes.addressof(buf)
                    + ctypes.sizeof(ctypes.wintypes.DWORD)
                    + i * ctypes.sizeof(MIB_TCPROW_OWNER_PID)
                )
                connections.add((row.dwOwningPid, row.dwRemoteAddr, row.dwRemotePort))

        except Exception:
            log_exception("PYAS_Process.ProcessMixin.get_connections_list:179")
            pass

        return connections

    def get_traffic_list(self):
        conns = self.get_connections_list()
        conn_map = {}

        for pid, _, _ in conns:
            conn_map[pid] = conn_map.get(pid, 0) + 1

        current_io = {}
        current_time = time.time()

        with self.lock_io:
            time_diff = current_time - self.last_io_time

            if time_diff <= 0:
                time_diff = 1

            self.last_io_time = current_time
            old_counters = self.last_io_counters.copy()

        result = []
        exist_process = self.get_process_list_pids()

        for pid in exist_process:
            name, file_path = "Unknown", ""

            if pid > 4:
                name, file_path = self.get_exe_info(pid)
            else:
                name = "System"

            down_speed, up_speed = 0, 0
            h = self.kernel32.OpenProcess(0x1000, False, pid)

            if h:
                try:
                    io = IO_COUNTERS()

                    if self.kernel32.GetProcessIoCounters(h, ctypes.byref(io)):
                        current_io[pid] = (io.ReadTransferCount, io.WriteTransferCount)

                        if pid in old_counters:
                            old_read, old_write = old_counters[pid]
                            down_speed = max(0, (io.ReadTransferCount - old_read) / time_diff)
                            up_speed = max(0, (io.WriteTransferCount - old_write) / time_diff)

                finally:
                    self.kernel32.CloseHandle(h)

            count = conn_map.get(pid, 0)

            if count > 0 or down_speed > 0 or up_speed > 0:
                result.append(
                    {
                        "pid": pid,
                        "name": name,
                        "path": file_path or "None",
                        "down": int(down_speed),
                        "up": int(up_speed),
                        "conn": count,
                    }
                )

        with self.lock_io:
            self.last_io_counters = current_io

        return result

    def _terminate_process_handle(self, handle, pid=None, expected_path=None):
        file_path = self.norm_path(self.get_process_file(handle))

        if self.path_equal(file_path, self.file_pyas):
            return False

        if expected_path and not self.path_equal(file_path, expected_path):
            return False

        return bool(self.kernel32.TerminateProcess(handle, 0))