procproxy 0.1.0__py3-none-any.whl

proxycore/proc.py

@@ -0,0 +1,342 @@
+"""pid → 名称/路径/会话/用户（跨平台）。
+
+字段约定：
+* `name` 不带路径（"chrome.exe"）；
+* `path` 可能为空（受保护进程，如 Windows 的 System/PID 4、带保护的 EDR 进程）；
+* `session_id` Windows 为会话 ID，Linux/macOS 为 0；
+* `user` 可空。
+
+红线：平台差异只允许出现在本文件与
+`backend/` 内部，核心层不得出现 `if windows` 分支。
+"""
+
+from __future__ import annotations
+
+import contextlib
+import os
+import sys
+from pathlib import Path
+from typing import Any
+
+from .types import ProcInfo
+
+IS_WINDOWS = os.name == "nt"
+IS_LINUX = sys.platform.startswith("linux")
+
+# 受保护/系统进程的兜底名称（取不到路径时至少要有个名字，否则规则无法按名称匹配）
+_WELL_KNOWN_NAMES = {0: "unknown", 4: "System"}
+
+
+# ------------------------------------------------------------------ 公共 ---
+def current_pid() -> int:
+    return os.getpid()
+
+
+def is_alive(pid: int) -> bool:
+    """进程是否存活（用于粘性表清理）。"""
+    if pid <= 0:
+        return False
+    if IS_WINDOWS:
+        return _win_is_alive(pid)
+    try:
+        os.kill(pid, 0)
+    except ProcessLookupError:
+        return False
+    except PermissionError:
+        return True
+    except OSError:
+        return False
+    return True
+
+
+def get_proc_info(pid: int, *, path: str | None = None, name: str | None = None) -> ProcInfo:
+    """查询进程信息。任何字段查询失败都退化为空串，绝不抛异常。
+
+    `path`/`name` 可显式传入（后端已经拿到元数据时避免重复查询，契约要求
+    "进程身份以驱动上报为准，用户态回传的 PID 只用于对账"）。
+    """
+    if pid <= 0:
+        return ProcInfo(pid=pid, name=name or _WELL_KNOWN_NAMES.get(pid, "unknown"), path=path or "")
+
+    resolved_name, resolved_path, session_id, user = "", "", 0, ""
+    if path is not None:
+        resolved_path = path
+        resolved_name = name if name is not None else basename_of(path)
+    elif IS_WINDOWS:
+        resolved_path, resolved_name = _win_name_path(pid)
+        session_id = _win_session_id(pid)
+        user = _win_user(pid)
+    else:
+        resolved_path, resolved_name = _posix_name_path(pid)
+        session_id = 0
+        user = _posix_user(pid)
+
+    if name is not None and resolved_name != name:
+        resolved_name = name
+    if not resolved_name:
+        resolved_name = basename_of(resolved_path) or _WELL_KNOWN_NAMES.get(pid, "unknown")
+    return ProcInfo(pid=pid, name=resolved_name, path=resolved_path, session_id=session_id, user=user)
+
+
+def basename_of(path: str) -> str:
+    """从路径取进程名（Windows 反斜杠与 POSIX 斜杠都支持）。"""
+    if not path:
+        return ""
+    # 不能用 os.path.basename：跨平台（Linux 上带 `\\` 的 Windows 路径会整段返回）
+    normalized = path.replace("\\", "/").rstrip("/")
+    return normalized.rsplit("/", 1)[-1] if normalized else ""
+
+
+def self_pids() -> set[int]:
+    """环路防护第一重：自身进程 PID 白名单（服务 + 拦截器 + 上游代理进程）。
+
+    上游代理进程如果跑在本机，会把它的 PID 也交给 `GuardSet` 一起白名单。
+    """
+    pids = {os.getpid()}
+    if hasattr(os, "getppid"):
+        with contextlib.suppress(OSError):  # pragma: no cover - 平台差异
+            pids.add(os.getppid())
+    return pids
+
+
+def iter_pids() -> list[int]:
+    """枚举本机 PID（粘性表清理、`cli status` 用）。"""
+    if IS_WINDOWS:
+        return _win_list_pids()
+    out: list[int] = []
+    for entry in os.listdir("/proc"):
+        if entry.isdigit():
+            out.append(int(entry))
+    return out
+
+
+# ------------------------------------------------------------ Windows ------
+def _win_kernel32() -> Any:
+    """kernel32 + 显式签名。
+
+    64 位下必须声明 `argtypes`/`restype`，否则 ctypes 会把 HANDLE 当 32 位
+    并有符号扩展，导致 `OverflowError: int too long to convert`。
+    """
+    import ctypes
+    from ctypes import wintypes
+
+    k32 = ctypes.WinDLL("kernel32", use_last_error=True)
+    k32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
+    k32.OpenProcess.restype = wintypes.HANDLE
+    k32.CloseHandle.argtypes = [wintypes.HANDLE]
+    k32.CloseHandle.restype = wintypes.BOOL
+    k32.QueryFullProcessImageNameW.argtypes = [
+        wintypes.HANDLE,
+        wintypes.DWORD,
+        wintypes.LPWSTR,
+        ctypes.POINTER(wintypes.DWORD),
+    ]
+    k32.QueryFullProcessImageNameW.restype = wintypes.BOOL
+    k32.ProcessIdToSessionId.argtypes = [wintypes.DWORD, ctypes.POINTER(wintypes.DWORD)]
+    k32.ProcessIdToSessionId.restype = wintypes.BOOL
+    k32.GetExitCodeProcess.argtypes = [wintypes.HANDLE, ctypes.POINTER(wintypes.DWORD)]
+    k32.GetExitCodeProcess.restype = wintypes.BOOL
+    return k32
+
+
+def _win_psapi() -> Any:
+    """`EnumProcesses` 在 psapi.dll（部分版本也在 kernel32 转发）。"""
+    import ctypes
+    from ctypes import wintypes
+
+    psapi = ctypes.WinDLL("psapi", use_last_error=True)
+    psapi.EnumProcesses.argtypes = [ctypes.POINTER(wintypes.DWORD), wintypes.DWORD, ctypes.POINTER(wintypes.DWORD)]
+    psapi.EnumProcesses.restype = wintypes.BOOL
+    return psapi
+
+
+def _win_name_path(pid: int) -> tuple[str, str]:
+    import ctypes
+    from ctypes import wintypes
+
+    PROCESS_QUERY_LIMITED_INFORMATION = 0x1000
+    k32 = _win_kernel32()
+    handle = k32.OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, False, pid)
+    if not handle:
+        # 受保护进程：路径为空（拿不到路径时的降级场景）
+        return "", _WELL_KNOWN_NAMES.get(pid, "")
+    try:
+        buf = ctypes.create_unicode_buffer(32768)
+        size = wintypes.DWORD(len(buf))
+        ok = k32.QueryFullProcessImageNameW(handle, 0, buf, ctypes.byref(size))
+        if not ok:
+            return "", _WELL_KNOWN_NAMES.get(pid, "")
+        full = buf.value
+        return full, basename_of(full)
+    finally:
+        k32.CloseHandle(handle)
+
+
+def _win_session_id(pid: int) -> int:
+    import ctypes
+    from ctypes import wintypes
+
+    k32 = _win_kernel32()
+    handle = k32.OpenProcess(0x0400, False, pid)  # PROCESS_QUERY_INFORMATION
+    if not handle:
+        return 0
+    try:
+        sid = wintypes.DWORD(0)
+        if not k32.ProcessIdToSessionId(wintypes.DWORD(pid), ctypes.byref(sid)):
+            return 0
+        return int(sid.value)
+    finally:
+        k32.CloseHandle(handle)
+
+
+def _win_advapi() -> Any:
+    """advapi32 + 显式签名（SID 指针必须按 64 位传，否则 OverflowError）。"""
+    import ctypes
+    from ctypes import wintypes
+
+    advapi = ctypes.WinDLL("advapi32", use_last_error=True)
+    advapi.OpenProcessToken.argtypes = [wintypes.HANDLE, wintypes.DWORD, ctypes.POINTER(wintypes.HANDLE)]
+    advapi.OpenProcessToken.restype = wintypes.BOOL
+    advapi.GetTokenInformation.argtypes = [
+        wintypes.HANDLE,
+        ctypes.c_int,
+        wintypes.LPVOID,
+        wintypes.DWORD,
+        ctypes.POINTER(wintypes.DWORD),
+    ]
+    advapi.GetTokenInformation.restype = wintypes.BOOL
+    advapi.LookupAccountSidW.argtypes = [
+        wintypes.LPCWSTR,
+        wintypes.LPVOID,
+        wintypes.LPWSTR,
+        ctypes.POINTER(wintypes.DWORD),
+        wintypes.LPWSTR,
+        ctypes.POINTER(wintypes.DWORD),
+        ctypes.POINTER(wintypes.DWORD),
+    ]
+    advapi.LookupAccountSidW.restype = wintypes.BOOL
+    return advapi
+
+
+def _win_user(pid: int) -> str:
+    import ctypes
+    from ctypes import wintypes
+
+    advapi = _win_advapi()
+    k32 = _win_kernel32()
+    TOKEN_QUERY = 0x0008
+    TokenUser = 1
+    handle = k32.OpenProcess(0x0400, False, pid)  # PROCESS_QUERY_INFORMATION
+    if not handle:
+        return ""
+    try:
+        token = wintypes.HANDLE()
+        if not advapi.OpenProcessToken(handle, TOKEN_QUERY, ctypes.byref(token)):
+            return ""
+        try:
+            size = wintypes.DWORD(0)
+            advapi.GetTokenInformation(token, TokenUser, None, 0, ctypes.byref(size))
+            if size.value == 0:
+                return ""
+            buf = ctypes.create_string_buffer(size.value)
+            if not advapi.GetTokenInformation(token, TokenUser, buf, size.value, ctypes.byref(size)):
+                return ""
+            # TOKEN_USER.Sid 是缓冲区的首指针（64 位下必须按 c_void_p 读出）
+            sid_ptr = ctypes.cast(buf, ctypes.POINTER(ctypes.c_void_p))[0]
+            if not sid_ptr:
+                return ""
+            return _win_sid_to_string(advapi, sid_ptr)
+        finally:
+            k32.CloseHandle(token)
+    finally:
+        k32.CloseHandle(handle)
+
+
+def _win_sid_to_string(advapi: Any, sid_ptr: Any) -> str:
+    import ctypes
+    from ctypes import wintypes
+
+    name = ctypes.create_unicode_buffer(256)
+    domain = ctypes.create_unicode_buffer(256)
+    name_len = wintypes.DWORD(256)
+    domain_len = wintypes.DWORD(256)
+    use = wintypes.DWORD(0)
+    ok = advapi.LookupAccountSidW(
+        None,
+        ctypes.c_void_p(sid_ptr),
+        name,
+        ctypes.byref(name_len),
+        domain,
+        ctypes.byref(domain_len),
+        ctypes.byref(use),
+    )
+    if ok:
+        return f"{domain.value}\\{name.value}" if domain.value else name.value
+    return ""
+
+
+def _win_is_alive(pid: int) -> bool:
+    import ctypes
+    from ctypes import wintypes
+
+    k32 = _win_kernel32()
+    handle = k32.OpenProcess(0x1000, False, pid)  # PROCESS_QUERY_LIMITED_INFORMATION
+    if not handle:
+        # 可能是权限不足而非进程不存在；用是否 ERROR_INVALID_PARAMETER 粗判
+        return ctypes.get_last_error() != 87
+    try:
+        code = wintypes.DWORD(0)
+        if k32.GetExitCodeProcess(handle, ctypes.byref(code)):
+            return code.value == 259  # STILL_ACTIVE
+        return True
+    finally:
+        k32.CloseHandle(handle)
+
+
+def _win_list_pids() -> list[int]:
+    import ctypes
+    from ctypes import wintypes
+
+    psapi = _win_psapi()
+    count = 4096
+    while True:
+        buf = (wintypes.DWORD * count)()
+        needed = wintypes.DWORD(0)
+        if not psapi.EnumProcesses(buf, ctypes.sizeof(buf), ctypes.byref(needed)):
+            return []
+        n = needed.value // ctypes.sizeof(wintypes.DWORD)
+        if n < count:
+            return [int(buf[i]) for i in range(n) if buf[i]]
+        count *= 2
+
+
+# ------------------------------------------------------------- POSIX ------
+def _posix_name_path(pid: int) -> tuple[str, str]:
+    # Linux：/proc/<pid>/exe 是最准的；退化为 /proc/<pid>/comm
+    exe_link = f"/proc/{pid}/exe"
+    try:
+        raw = os.readlink(exe_link)
+        if raw.endswith(" (deleted)"):
+            raw = raw[: -len(" (deleted)")]
+        return raw, basename_of(raw)
+    except OSError:
+        pass
+    try:
+        name = Path(f"/proc/{pid}/comm").read_text(encoding="utf-8").strip()
+        return "", name
+    except OSError:
+        return "", ""
+
+
+def _posix_user(pid: int) -> str:
+    try:
+        st = os.stat(f"/proc/{pid}")
+    except OSError:
+        return ""
+    lookup = getattr(__import__("pwd"), "getpwuid", None) if IS_LINUX else None
+    if lookup is None:  # pragma: no cover - 非 Linux
+        return str(st.st_uid)
+    try:
+        return str(lookup(st.st_uid).pw_name)
+    except KeyError:  # pragma: no cover - uid 无对应用户名
+        return str(st.st_uid)