192 lines
8.2 KiB
Python
192 lines
8.2 KiB
Python
"""One desktop application per Windows user and interactive logon session.
|
|
|
|
Windows owns the mutex lifetime, so an abnormal process exit cannot leave a
|
|
stale lock file behind. The activation event is created *before* competing for
|
|
the mutex: another launch can request activation while the first is still
|
|
initialising its login dialog. Only the primary instance consumes that event.
|
|
|
|
The default identity deliberately excludes installation paths and versions.
|
|
``namespace`` is an override for isolated tests, not a per-installation setting.
|
|
Acquire and close a primary guard on the same (normally GUI) thread.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import ctypes
|
|
from ctypes import wintypes
|
|
import hashlib
|
|
import os
|
|
import threading
|
|
|
|
|
|
_APPLICATION_ID = "ZhenAI.WeComAssistant.Desktop"
|
|
_DEFAULT_NAMESPACE = "main-desktop"
|
|
_WAIT_OBJECT_0 = 0
|
|
_WAIT_ABANDONED = 0x80
|
|
_WAIT_TIMEOUT = 0x102
|
|
_TOKEN_QUERY = 0x0008
|
|
_TOKEN_USER = 1
|
|
_ERROR_INSUFFICIENT_BUFFER = 122
|
|
_process_owners: set[str] = set()
|
|
_process_lock = threading.RLock()
|
|
|
|
|
|
class _SidAndAttributes(ctypes.Structure):
|
|
_fields_ = [("Sid", wintypes.LPVOID), ("Attributes", wintypes.DWORD)]
|
|
|
|
|
|
class _TokenUser(ctypes.Structure):
|
|
_fields_ = [("User", _SidAndAttributes)]
|
|
|
|
|
|
class _WindowsObjects:
|
|
def __init__(self) -> None:
|
|
if os.name != "nt":
|
|
raise OSError("Desktop single-instance protection requires Windows")
|
|
self.kernel = ctypes.WinDLL("kernel32", use_last_error=True)
|
|
self.advapi = ctypes.WinDLL("advapi32", use_last_error=True)
|
|
signatures = {
|
|
"CreateMutexW": ([wintypes.LPVOID, wintypes.BOOL, wintypes.LPCWSTR], wintypes.HANDLE),
|
|
"CreateEventW": ([wintypes.LPVOID, wintypes.BOOL, wintypes.BOOL, wintypes.LPCWSTR], wintypes.HANDLE),
|
|
"WaitForSingleObject": ([wintypes.HANDLE, wintypes.DWORD], wintypes.DWORD),
|
|
"ReleaseMutex": ([wintypes.HANDLE], wintypes.BOOL),
|
|
"SetEvent": ([wintypes.HANDLE], wintypes.BOOL),
|
|
"CloseHandle": ([wintypes.HANDLE], wintypes.BOOL),
|
|
"GetCurrentProcess": ([], wintypes.HANDLE),
|
|
}
|
|
for name, (args, result) in signatures.items():
|
|
function = getattr(self.kernel, name)
|
|
function.argtypes, function.restype = args, result
|
|
self.advapi.OpenProcessToken.argtypes = [wintypes.HANDLE, wintypes.DWORD, ctypes.POINTER(wintypes.HANDLE)]
|
|
self.advapi.OpenProcessToken.restype = wintypes.BOOL
|
|
self.advapi.GetTokenInformation.argtypes = [wintypes.HANDLE, ctypes.c_int, wintypes.LPVOID, wintypes.DWORD, ctypes.POINTER(wintypes.DWORD)]
|
|
self.advapi.GetTokenInformation.restype = wintypes.BOOL
|
|
self.advapi.GetLengthSid.argtypes = [wintypes.LPVOID]
|
|
self.advapi.GetLengthSid.restype = wintypes.DWORD
|
|
|
|
@staticmethod
|
|
def error(operation: str) -> OSError:
|
|
code = ctypes.get_last_error()
|
|
return OSError(code, f"{operation}: {ctypes.FormatError(code)}")
|
|
|
|
def user_identity(self) -> str:
|
|
"""Use the security identifier, independent of profile or account name."""
|
|
token = wintypes.HANDLE()
|
|
if not self.advapi.OpenProcessToken(self.kernel.GetCurrentProcess(), _TOKEN_QUERY, ctypes.byref(token)):
|
|
raise self.error("OpenProcessToken")
|
|
try:
|
|
size = wintypes.DWORD()
|
|
self.advapi.GetTokenInformation(token, _TOKEN_USER, None, 0, ctypes.byref(size))
|
|
if ctypes.get_last_error() != _ERROR_INSUFFICIENT_BUFFER or not 0 < size.value <= 65536:
|
|
raise self.error("GetTokenInformation(size)")
|
|
data = ctypes.create_string_buffer(size.value)
|
|
if not self.advapi.GetTokenInformation(token, _TOKEN_USER, data, size.value, ctypes.byref(size)):
|
|
raise self.error("GetTokenInformation")
|
|
sid = ctypes.cast(data, ctypes.POINTER(_TokenUser)).contents.User.Sid
|
|
length = self.advapi.GetLengthSid(sid)
|
|
if not length:
|
|
raise self.error("GetLengthSid")
|
|
return hashlib.sha256(ctypes.string_at(sid, length)).hexdigest()[:24]
|
|
finally:
|
|
self.kernel.CloseHandle(token)
|
|
|
|
|
|
class DesktopInstanceGuard:
|
|
"""A process-lifetime primary guard and a coalescing activation notification.
|
|
|
|
``acquire()`` returns False for another launch. That launch should call
|
|
``request_activation()`` then ``close()`` and exit. A GUI timer in the
|
|
primary calls ``consume_activation()`` only when it can display its login
|
|
dialog or main window. Several clicks may coalesce into one activation.
|
|
|
|
Win32 errors raise OSError rather than silently permitting two primaries.
|
|
The Windows ``Local`` object namespace supplies session isolation; the SID
|
|
digest additionally separates users in the same session.
|
|
"""
|
|
|
|
def __init__(self, *, namespace: str = _DEFAULT_NAMESPACE) -> None:
|
|
if not isinstance(namespace, str) or not namespace.strip():
|
|
raise ValueError("namespace must be a nonempty string")
|
|
self._api = _WindowsObjects()
|
|
scope = hashlib.sha256(namespace.encode("utf-8")).hexdigest()[:24]
|
|
name = f"Local\\{_APPLICATION_ID}.{self._api.user_identity()}.{scope}"
|
|
self._mutex_name = name + ".Mutex"
|
|
self._event_name = name + ".Activate"
|
|
self._mutex = None
|
|
self._event = None
|
|
self._owned = False
|
|
self._owner_thread = None
|
|
self._closed = False
|
|
self._lock = threading.RLock()
|
|
|
|
def _ensure_handles(self) -> None:
|
|
if self._closed:
|
|
raise RuntimeError("Desktop instance guard is closed")
|
|
if self._mutex is not None:
|
|
return
|
|
event = self._api.kernel.CreateEventW(None, False, False, self._event_name)
|
|
if not event:
|
|
raise self._api.error("CreateEventW")
|
|
mutex = self._api.kernel.CreateMutexW(None, False, self._mutex_name)
|
|
if not mutex:
|
|
error = self._api.error("CreateMutexW")
|
|
self._api.kernel.CloseHandle(event)
|
|
raise error
|
|
self._event, self._mutex = event, mutex
|
|
|
|
def acquire(self) -> bool:
|
|
with self._lock, _process_lock:
|
|
self._ensure_handles()
|
|
if self._owned:
|
|
return True
|
|
# A Windows mutex is recursive on its owning thread. Do not let a
|
|
# second guard in that same process mistake recursion for ownership.
|
|
if self._mutex_name in _process_owners:
|
|
return False
|
|
result = self._api.kernel.WaitForSingleObject(self._mutex, 0)
|
|
if result == _WAIT_TIMEOUT:
|
|
return False
|
|
if result not in (_WAIT_OBJECT_0, _WAIT_ABANDONED):
|
|
raise self._api.error("WaitForSingleObject(mutex)")
|
|
self._owned = True
|
|
self._owner_thread = threading.get_ident()
|
|
_process_owners.add(self._mutex_name)
|
|
return True
|
|
|
|
def request_activation(self) -> bool:
|
|
with self._lock:
|
|
self._ensure_handles()
|
|
if not self._api.kernel.SetEvent(self._event):
|
|
raise self._api.error("SetEvent")
|
|
return True
|
|
|
|
def consume_activation(self) -> bool:
|
|
with self._lock:
|
|
if not self._owned or self._closed:
|
|
return False
|
|
result = self._api.kernel.WaitForSingleObject(self._event, 0)
|
|
if result == _WAIT_TIMEOUT:
|
|
return False
|
|
if result != _WAIT_OBJECT_0:
|
|
raise self._api.error("WaitForSingleObject(event)")
|
|
return True
|
|
|
|
def close(self) -> None:
|
|
with self._lock, _process_lock:
|
|
if self._closed:
|
|
return
|
|
if self._owned and self._owner_thread != threading.get_ident():
|
|
raise RuntimeError("Close the primary guard on its acquiring thread")
|
|
error = None
|
|
if self._owned:
|
|
if not self._api.kernel.ReleaseMutex(self._mutex):
|
|
error = self._api.error("ReleaseMutex")
|
|
_process_owners.discard(self._mutex_name)
|
|
for handle in (self._mutex, self._event):
|
|
if handle is not None:
|
|
self._api.kernel.CloseHandle(handle)
|
|
self._owned = False
|
|
self._mutex = self._event = None
|
|
self._closed = True
|
|
if error:
|
|
raise error
|