zapret-kvn/xray_fluent/platform/windows/subprocess_utils.py

from __future__ import annotations

from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeoutError
import locale
import os
import subprocess
import time
from pathlib import Path
from typing import Any


CREATE_NO_WINDOW = 0x08000000 if os.name == "nt" else 0
_SUBPROCESS_EXECUTOR = ThreadPoolExecutor(max_workers=4, thread_name_prefix="xray_fluent_subprocess")


def decode_output(data: bytes | None) -> str:
    if not data:
        return ""
    encodings: list[str] = ["utf-8"]
    if os.name == "nt":
        encodings.extend(["cp866", locale.getpreferredencoding(False), "cp1251"])
    for encoding in dict.fromkeys(encodings):
        try:
            return data.decode(encoding)
        except (LookupError, UnicodeDecodeError):
            continue
    return data.decode("utf-8", errors="replace")


def result_output_text(result: subprocess.CompletedProcess[bytes]) -> str:
    return decode_output(result.stdout or result.stderr or b"")


def run_text(
    command: list[str],
    *,
    timeout: float,
    check: bool = False,
    creationflags: int | None = None,
) -> subprocess.CompletedProcess[bytes]:
    kwargs: dict[str, int | bool | float] = {
        "capture_output": True,
        "text": False,
        "timeout": timeout,
        "check": check,
    }
    # The app is a windowed program: a console child without CREATE_NO_WINDOW
    # would pop up its own console window.
    kwargs["creationflags"] = CREATE_NO_WINDOW if creationflags is None else creationflags
    return subprocess.run(command, **kwargs)


# Pumped helpers below (pump_qt_events, sleep_with_events, wait_for_qprocess_*,
# wait_for_future_with_events, run_text_pumped) re-enter the Qt event loop.
# Connection transitions never use them: they run as async_steps generators.
# Remaining callers are the synchronous shutdown driver (run_steps_blocking),
# version probes on settings pages and the legacy Xray TUN route manager.


def pump_qt_events() -> None:
    try:
        from PyQt6.QtCore import QThread
        from PyQt6.QtWidgets import QApplication
    except Exception:
        return

    app = QApplication.instance()
    if app is not None and QThread.currentThread() == app.thread():
        app.processEvents()


def sleep_with_events(duration_sec: float, *, step_sec: float = 0.05) -> None:
    deadline = time.monotonic() + max(0.0, duration_sec)
    while True:
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            return
        pump_qt_events()
        time.sleep(min(step_sec, remaining))


def _wait_for_qprocess_call(
    process: Any,
    method_name: str,
    timeout_ms: int,
    *,
    slice_ms: int = 50,
) -> bool:
    remaining = max(0, int(timeout_ms))
    waiter = getattr(process, method_name)

    while remaining > 0:
        step = min(slice_ms, remaining)
        if waiter(step):
            return True
        pump_qt_events()
        remaining -= step

    pump_qt_events()
    return False


def wait_for_qprocess_started(process: Any, timeout_ms: int) -> bool:
    return _wait_for_qprocess_call(process, "waitForStarted", timeout_ms)


def wait_for_qprocess_finished(process: Any, timeout_ms: int) -> bool:
    return _wait_for_qprocess_call(process, "waitForFinished", timeout_ms)


def wait_for_qprocess_ready_read(process: Any, timeout_ms: int) -> bool:
    return _wait_for_qprocess_call(process, "waitForReadyRead", timeout_ms)


def wait_for_future_with_events(future: Any, timeout_sec: float, *, slice_sec: float = 0.05) -> Any:
    deadline = time.monotonic() + max(0.0, float(timeout_sec))
    while True:
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            raise TimeoutError("Future did not complete before timeout")
        try:
            return future.result(timeout=min(slice_sec, remaining))
        except FutureTimeoutError:
            pump_qt_events()


def run_text_pumped(
    command: list[str],
    *,
    timeout: float,
    check: bool = False,
    creationflags: int | None = None,
) -> subprocess.CompletedProcess[bytes]:
    future = _SUBPROCESS_EXECUTOR.submit(
        run_text,
        command,
        timeout=timeout,
        check=check,
        creationflags=creationflags,
    )
    return wait_for_future_with_events(future, timeout + 0.5)


def is_same_path(left: str | Path | None, right: str | Path | None) -> bool:
    if not left or not right:
        return False
    try:
        return Path(left).resolve() == Path(right).resolve()
    except Exception:
        return False


def _kill_processes_native(target: Path, timeout: float) -> bool:
    """WinAPI-снимок процессов вместо PowerShell/WMI: миллисекунды, а не секунды.

    Тот же отбор, что у PowerShell-варианта: имя и полный путь exe совпадают.
    Завершённые процессы ждём до ``timeout``, чтобы порты и адаптер освободились.
    """
    from zapret_update import processes as win_processes

    pids = win_processes.processes_at(target)
    for pid in pids:
        win_processes.terminate(pid)
    deadline = time.monotonic() + max(0.0, timeout)
    for pid in pids:
        win_processes.wait_for_exit(pid, max(0.0, deadline - time.monotonic()))
    return bool(pids)


def kill_processes_by_path(
    process_name: str,
    executable_path: str | Path,
    *,
    timeout: float = 5.0,
    pump: bool = True,
) -> bool:
    # pump=False — вызов уже выполняется в worker-потоке (async_steps.run_in_worker):
    # прокачка событий там не нужна, а вложенный сабмит в _SUBPROCESS_EXECUTOR
    # рискует дедлоком при исчерпании пула.
    if os.name != "nt":
        return False
    target = Path(executable_path)
    try:
        return _kill_processes_native(target, timeout)
    except Exception:
        pass  # WinAPI-путь не сработал — прежний PowerShell ниже.
    target_text = str(target).replace("'", "''")
    script = (
        "$matches = @(Get-CimInstance Win32_Process | "
        f"Where-Object {{ $_.Name -eq '{process_name}' -and $_.ExecutablePath -eq '{target_text}' }}); "
        "$matches | ForEach-Object { Stop-Process -Id $_.ProcessId -Force }; "
        "Write-Output $matches.Count"
    )
    command = ["powershell", "-NoProfile", "-NonInteractive", "-Command", script]
    if pump:
        result = run_text_pumped(command, timeout=timeout, creationflags=CREATE_NO_WINDOW)
    else:
        result = run_text(command, timeout=timeout, creationflags=CREATE_NO_WINDOW)
    if result.returncode != 0:
        return False
    return result_output_text(result).strip() not in {"", "0"}