342 lines
11 KiB
Python
342 lines
11 KiB
Python
from __future__ import annotations
|
||
|
||
import ctypes
|
||
import os
|
||
from ctypes import wintypes
|
||
from dataclasses import dataclass
|
||
|
||
|
||
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
|
||
ntdll = ctypes.WinDLL("ntdll", use_last_error=True)
|
||
advapi32 = ctypes.WinDLL("advapi32", use_last_error=True)
|
||
|
||
|
||
PROCESS_DUP_HANDLE = 0x0040
|
||
PROCESS_QUERY_LIMITED_INFORMATION = 0x1000
|
||
DUPLICATE_SAME_ACCESS = 0x00000002
|
||
DUPLICATE_CLOSE_SOURCE = 0x00000001
|
||
|
||
GENERIC_READ = 0x80000000
|
||
FILE_SHARE_READ = 0x00000001
|
||
FILE_SHARE_WRITE = 0x00000002
|
||
FILE_SHARE_DELETE = 0x00000004
|
||
OPEN_EXISTING = 3
|
||
FILE_ATTRIBUTE_NORMAL = 0x00000080
|
||
|
||
|
||
SystemExtendedHandleInformation = 64
|
||
|
||
|
||
class SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX(ctypes.Structure):
|
||
_fields_ = [
|
||
("Object", ctypes.c_void_p),
|
||
("UniqueProcessId", ctypes.c_void_p),
|
||
("HandleValue", ctypes.c_void_p),
|
||
("GrantedAccess", wintypes.ULONG),
|
||
("CreatorBackTraceIndex", wintypes.USHORT),
|
||
("ObjectTypeIndex", wintypes.USHORT),
|
||
("HandleAttributes", wintypes.ULONG),
|
||
("Reserved", wintypes.ULONG),
|
||
]
|
||
|
||
|
||
class SYSTEM_HANDLE_INFORMATION_EX(ctypes.Structure):
|
||
_fields_ = [
|
||
("NumberOfHandles", ctypes.c_void_p),
|
||
("Reserved", ctypes.c_void_p),
|
||
("Handles", SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX * 1),
|
||
]
|
||
|
||
|
||
NtQuerySystemInformation = ntdll.NtQuerySystemInformation
|
||
NtQuerySystemInformation.argtypes = [
|
||
wintypes.ULONG,
|
||
wintypes.LPVOID,
|
||
wintypes.ULONG,
|
||
wintypes.PULONG,
|
||
]
|
||
NtQuerySystemInformation.restype = wintypes.LONG
|
||
|
||
|
||
kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
|
||
kernel32.OpenProcess.restype = wintypes.HANDLE
|
||
|
||
kernel32.CloseHandle.argtypes = [wintypes.HANDLE]
|
||
kernel32.CloseHandle.restype = wintypes.BOOL
|
||
|
||
kernel32.DuplicateHandle.argtypes = [
|
||
wintypes.HANDLE,
|
||
wintypes.HANDLE,
|
||
wintypes.HANDLE,
|
||
ctypes.POINTER(wintypes.HANDLE),
|
||
wintypes.DWORD,
|
||
wintypes.BOOL,
|
||
wintypes.DWORD,
|
||
]
|
||
kernel32.DuplicateHandle.restype = wintypes.BOOL
|
||
|
||
kernel32.GetCurrentProcess.argtypes = []
|
||
kernel32.GetCurrentProcess.restype = wintypes.HANDLE
|
||
|
||
kernel32.GetFinalPathNameByHandleW.argtypes = [
|
||
wintypes.HANDLE,
|
||
wintypes.LPWSTR,
|
||
wintypes.DWORD,
|
||
wintypes.DWORD,
|
||
]
|
||
kernel32.GetFinalPathNameByHandleW.restype = wintypes.DWORD
|
||
|
||
kernel32.GetCurrentProcessId.argtypes = []
|
||
kernel32.GetCurrentProcessId.restype = wintypes.DWORD
|
||
|
||
kernel32.CreateFileW.argtypes = [
|
||
wintypes.LPCWSTR,
|
||
wintypes.DWORD,
|
||
wintypes.DWORD,
|
||
wintypes.LPVOID,
|
||
wintypes.DWORD,
|
||
wintypes.DWORD,
|
||
wintypes.HANDLE,
|
||
]
|
||
kernel32.CreateFileW.restype = wintypes.HANDLE
|
||
|
||
|
||
def _normalize_path(p: str) -> str:
|
||
p = os.path.abspath(p)
|
||
p = os.path.normcase(p)
|
||
return p
|
||
|
||
|
||
def _normalize_final_path(p: str) -> str:
|
||
# GetFinalPathNameByHandleW 常见返回：\\?\C:\path\file 或 \\?\UNC\server\share\path
|
||
if p.startswith("\\\\?\\UNC\\"):
|
||
p = "\\\\" + p[len("\\\\?\\UNC\\") :]
|
||
elif p.startswith("\\\\?\\"):
|
||
p = p[len("\\\\?\\") :]
|
||
return _normalize_path(p)
|
||
|
||
|
||
def _try_enable_debug_privilege() -> None:
|
||
# best-effort：没有也能工作，只是会漏掉部分系统进程
|
||
TOKEN_ADJUST_PRIVILEGES = 0x0020
|
||
TOKEN_QUERY = 0x0008
|
||
SE_PRIVILEGE_ENABLED = 0x0002
|
||
|
||
class LUID(ctypes.Structure):
|
||
_fields_ = [("LowPart", wintypes.DWORD), ("HighPart", wintypes.LONG)]
|
||
|
||
class LUID_AND_ATTRIBUTES(ctypes.Structure):
|
||
_fields_ = [("Luid", LUID), ("Attributes", wintypes.DWORD)]
|
||
|
||
class TOKEN_PRIVILEGES(ctypes.Structure):
|
||
_fields_ = [("PrivilegeCount", wintypes.DWORD), ("Privileges", LUID_AND_ATTRIBUTES * 1)]
|
||
|
||
OpenProcessToken = advapi32.OpenProcessToken
|
||
OpenProcessToken.argtypes = [wintypes.HANDLE, wintypes.DWORD, ctypes.POINTER(wintypes.HANDLE)]
|
||
OpenProcessToken.restype = wintypes.BOOL
|
||
|
||
LookupPrivilegeValueW = advapi32.LookupPrivilegeValueW
|
||
LookupPrivilegeValueW.argtypes = [wintypes.LPCWSTR, wintypes.LPCWSTR, ctypes.POINTER(LUID)]
|
||
LookupPrivilegeValueW.restype = wintypes.BOOL
|
||
|
||
AdjustTokenPrivileges = advapi32.AdjustTokenPrivileges
|
||
AdjustTokenPrivileges.argtypes = [
|
||
wintypes.HANDLE,
|
||
wintypes.BOOL,
|
||
ctypes.POINTER(TOKEN_PRIVILEGES),
|
||
wintypes.DWORD,
|
||
wintypes.PVOID,
|
||
wintypes.PVOID,
|
||
]
|
||
AdjustTokenPrivileges.restype = wintypes.BOOL
|
||
|
||
token = wintypes.HANDLE()
|
||
if not OpenProcessToken(kernel32.GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, ctypes.byref(token)):
|
||
return
|
||
try:
|
||
luid = LUID()
|
||
if not LookupPrivilegeValueW(None, "SeDebugPrivilege", ctypes.byref(luid)):
|
||
return
|
||
tp = TOKEN_PRIVILEGES()
|
||
tp.PrivilegeCount = 1
|
||
tp.Privileges[0].Luid = luid
|
||
tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED
|
||
AdjustTokenPrivileges(token, False, ctypes.byref(tp), 0, None, None)
|
||
finally:
|
||
kernel32.CloseHandle(token)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class FileLock:
|
||
pid: int
|
||
handle: int
|
||
path: str
|
||
|
||
|
||
def _query_system_handles(max_size: int = 1 << 28) -> tuple[ctypes.Array, int] | tuple[None, int]:
|
||
"""返回 (buffer, handle_count)；失败返回 (None, 0)。"""
|
||
size = 1 << 20
|
||
while True:
|
||
buf = ctypes.create_string_buffer(size)
|
||
ret_len = wintypes.ULONG(0)
|
||
status = NtQuerySystemInformation(
|
||
SystemExtendedHandleInformation, buf, size, ctypes.byref(ret_len)
|
||
)
|
||
if status == 0: # STATUS_SUCCESS
|
||
info = ctypes.cast(buf, ctypes.POINTER(SYSTEM_HANDLE_INFORMATION_EX)).contents
|
||
n = int(info.NumberOfHandles)
|
||
return buf, n
|
||
# STATUS_INFO_LENGTH_MISMATCH = 0xC0000004 (signed: -1073741820)
|
||
if status in (-1073741820, 0xC0000004):
|
||
size = max(size * 2, int(ret_len.value) + 0x10000)
|
||
if size > max_size:
|
||
return None, 0
|
||
continue
|
||
return None, 0
|
||
|
||
|
||
def _get_file_object_type_index(buf: ctypes.Array, handle_count: int, sample_path: str) -> int | None:
|
||
"""
|
||
通过“本进程打开一个文件句柄”，再在系统句柄表中找到它对应的 ObjectTypeIndex。
|
||
这样可以只扫描 File 类型句柄，速度会快很多。
|
||
"""
|
||
try:
|
||
h = kernel32.CreateFileW(
|
||
sample_path,
|
||
GENERIC_READ,
|
||
FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE,
|
||
None,
|
||
OPEN_EXISTING,
|
||
FILE_ATTRIBUTE_NORMAL,
|
||
None,
|
||
)
|
||
if not h or h == wintypes.HANDLE(-1).value:
|
||
return None
|
||
except Exception:
|
||
return None
|
||
|
||
try:
|
||
cur_pid = int(kernel32.GetCurrentProcessId())
|
||
hv = int(ctypes.cast(h, ctypes.c_size_t).value)
|
||
|
||
n = int(handle_count)
|
||
if n <= 0:
|
||
return None
|
||
array_type = SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX * n
|
||
info = ctypes.cast(buf, ctypes.POINTER(SYSTEM_HANDLE_INFORMATION_EX)).contents
|
||
handles = ctypes.cast(ctypes.addressof(info.Handles), ctypes.POINTER(array_type)).contents
|
||
|
||
for ent in handles:
|
||
pid = int(ctypes.cast(ent.UniqueProcessId, ctypes.c_size_t).value)
|
||
if pid != cur_pid:
|
||
continue
|
||
val = int(ctypes.cast(ent.HandleValue, ctypes.c_size_t).value)
|
||
if val == hv:
|
||
return int(ent.ObjectTypeIndex)
|
||
return None
|
||
finally:
|
||
try:
|
||
kernel32.CloseHandle(h)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def find_file_locks(target_path: str, *, max_handles: int = 200000) -> list[FileLock]:
|
||
"""
|
||
返回当前系统中占用 target_path 的 (pid, handle) 列表。
|
||
- 基于 SystemExtendedHandleInformation + DuplicateHandle + GetFinalPathNameByHandleW
|
||
"""
|
||
_try_enable_debug_privilege()
|
||
target_norm = _normalize_path(target_path)
|
||
|
||
buf, n = _query_system_handles()
|
||
if buf is None or n <= 0:
|
||
return []
|
||
|
||
if n <= 0:
|
||
return []
|
||
if n > max_handles:
|
||
n = max_handles
|
||
|
||
# 重新按实际数量解释数组
|
||
array_type = SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX * n
|
||
info = ctypes.cast(buf, ctypes.POINTER(SYSTEM_HANDLE_INFORMATION_EX)).contents
|
||
handles = ctypes.cast(ctypes.addressof(info.Handles), ctypes.POINTER(array_type)).contents
|
||
|
||
cur_proc = kernel32.GetCurrentProcess()
|
||
locks: list[FileLock] = []
|
||
|
||
file_type_index = _get_file_object_type_index(buf, int(info.NumberOfHandles), target_path)
|
||
# 进程句柄缓存：避免每个 handle 都 OpenProcess/CloseHandle
|
||
proc_cache: dict[int, wintypes.HANDLE] = {}
|
||
|
||
for h in handles:
|
||
pid = int(ctypes.cast(h.UniqueProcessId, ctypes.c_size_t).value)
|
||
hv = int(ctypes.cast(h.HandleValue, ctypes.c_size_t).value)
|
||
if pid <= 0 or hv <= 0:
|
||
continue
|
||
if file_type_index is not None and int(h.ObjectTypeIndex) != int(file_type_index):
|
||
continue
|
||
|
||
ph = proc_cache.get(pid)
|
||
if not ph:
|
||
ph = kernel32.OpenProcess(
|
||
PROCESS_DUP_HANDLE | PROCESS_QUERY_LIMITED_INFORMATION, False, pid
|
||
)
|
||
if not ph:
|
||
continue
|
||
proc_cache[pid] = ph
|
||
try:
|
||
dup = wintypes.HANDLE()
|
||
if not kernel32.DuplicateHandle(ph, wintypes.HANDLE(hv), cur_proc, ctypes.byref(dup), 0, False, DUPLICATE_SAME_ACCESS):
|
||
continue
|
||
try:
|
||
# 获取路径
|
||
bufw = ctypes.create_unicode_buffer(4096)
|
||
got = kernel32.GetFinalPathNameByHandleW(dup, bufw, len(bufw), 0)
|
||
if got == 0 or got >= len(bufw):
|
||
continue
|
||
final_norm = _normalize_final_path(bufw.value)
|
||
if final_norm == target_norm:
|
||
locks.append(FileLock(pid=pid, handle=hv, path=final_norm))
|
||
finally:
|
||
kernel32.CloseHandle(dup)
|
||
finally:
|
||
pass
|
||
|
||
for ph in proc_cache.values():
|
||
try:
|
||
kernel32.CloseHandle(ph)
|
||
except Exception:
|
||
pass
|
||
|
||
return locks
|
||
|
||
|
||
def close_remote_handle(pid: int, handle_value: int) -> bool:
|
||
"""
|
||
关闭目标进程中指定 handle（不结束进程）。
|
||
需要足够权限（通常管理员 + SeDebugPrivilege 才更稳）。
|
||
"""
|
||
_try_enable_debug_privilege()
|
||
ph = kernel32.OpenProcess(PROCESS_DUP_HANDLE, False, int(pid))
|
||
if not ph:
|
||
return False
|
||
try:
|
||
dup = wintypes.HANDLE()
|
||
ok = kernel32.DuplicateHandle(
|
||
ph,
|
||
wintypes.HANDLE(int(handle_value)),
|
||
kernel32.GetCurrentProcess(),
|
||
ctypes.byref(dup),
|
||
0,
|
||
False,
|
||
DUPLICATE_CLOSE_SOURCE,
|
||
)
|
||
if ok and dup:
|
||
kernel32.CloseHandle(dup)
|
||
return bool(ok)
|
||
finally:
|
||
kernel32.CloseHandle(ph)
|
||
|