343 lines
11 KiB
Python
343 lines
11 KiB
Python
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))
|