254 lines
9.9 KiB
Python
254 lines
9.9 KiB
Python
"""按需提权的 TUN 控制通道；GUI 与管理员工作进程共用同一个入口。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ctypes
|
|
import hmac
|
|
import json
|
|
import os
|
|
import secrets
|
|
import socket
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
from typing import Any, Callable
|
|
|
|
|
|
WORKER_FLAG = "--rrpp-tun-worker"
|
|
_MAX_MESSAGE = 1024 * 1024
|
|
|
|
|
|
class TunModeError(RuntimeError):
|
|
pass
|
|
|
|
|
|
def _send(sock: socket.socket, value: dict[str, Any], lock: threading.Lock | None = None) -> None:
|
|
payload = (json.dumps(value, ensure_ascii=False, separators=(",", ":")) + "\n").encode("utf-8")
|
|
if len(payload) > _MAX_MESSAGE:
|
|
raise TunModeError("TUN 控制消息过大")
|
|
if lock is None:
|
|
sock.sendall(payload)
|
|
else:
|
|
with lock:
|
|
sock.sendall(payload)
|
|
|
|
|
|
def _receive(reader) -> dict[str, Any]:
|
|
line = reader.readline(_MAX_MESSAGE + 1)
|
|
if not line or len(line) > _MAX_MESSAGE or not line.endswith(b"\n"):
|
|
raise TunModeError("TUN 工作进程连接已断开或消息无效")
|
|
value = json.loads(line.decode("utf-8"))
|
|
if not isinstance(value, dict):
|
|
raise TunModeError("TUN 控制消息格式错误")
|
|
return value
|
|
|
|
|
|
class TunModeController:
|
|
"""普通权限 GUI 仅在首次开启 TUN 时启动管理员工作进程。"""
|
|
|
|
def __init__(self, command: list[str], log: Callable[[str], None],
|
|
detail_log: Callable[[str], None] | None = None):
|
|
self.command = command
|
|
self.log = log
|
|
self.detail_log = detail_log or log
|
|
self._lock = threading.RLock()
|
|
self._socket: socket.socket | None = None
|
|
self._reader = None
|
|
|
|
def _launch(self, port: int, token: str) -> None:
|
|
args = [*self.command[1:], WORKER_FLAG, str(port), token]
|
|
if os.name == "nt":
|
|
# PyInstaller onefile 的提权子进程必须独立解包，以便 GUI 退出后仍可清理网络。
|
|
previous = os.environ.get("PYINSTALLER_RESET_ENVIRONMENT")
|
|
os.environ["PYINSTALLER_RESET_ENVIRONMENT"] = "1"
|
|
try:
|
|
result = ctypes.windll.shell32.ShellExecuteW(
|
|
None, "runas", self.command[0], subprocess.list2cmdline(args),
|
|
os.path.dirname(os.path.abspath(self.command[0])), 0,
|
|
)
|
|
finally:
|
|
if previous is None:
|
|
os.environ.pop("PYINSTALLER_RESET_ENVIRONMENT", None)
|
|
else:
|
|
os.environ["PYINSTALLER_RESET_ENVIRONMENT"] = previous
|
|
if result <= 32:
|
|
raise TunModeError("管理员授权已取消或无法启动 TUN 工作进程")
|
|
else:
|
|
# 非 Windows 仅供无网卡的协议测试；真实 TUN 后端仍限定 Windows。
|
|
subprocess.Popen([self.command[0], *args], stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL)
|
|
|
|
def start(self, payload: dict[str, Any], cancel_event=None) -> None:
|
|
with self._lock:
|
|
if self._socket is not None:
|
|
raise TunModeError("TUN 工作进程已经启动")
|
|
token = secrets.token_hex(32)
|
|
listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
listener.bind(("127.0.0.1", 0))
|
|
listener.listen(2)
|
|
deadline = time.monotonic() + 90
|
|
try:
|
|
self._launch(listener.getsockname()[1], token)
|
|
while True:
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
raise TunModeError("用户取消了开启代理")
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
raise TunModeError("等待 TUN 管理员工作进程超时")
|
|
listener.settimeout(min(0.5, remaining))
|
|
try:
|
|
connection, _peer = listener.accept()
|
|
except socket.timeout:
|
|
continue
|
|
connection.settimeout(3)
|
|
reader = connection.makefile("rb")
|
|
try:
|
|
hello = _receive(reader)
|
|
valid = (hello.get("type") == "hello" and
|
|
hmac.compare_digest(str(hello.get("token") or ""), token))
|
|
except Exception:
|
|
valid = False
|
|
if valid:
|
|
connection.settimeout(120)
|
|
self._socket, self._reader = connection, reader
|
|
break
|
|
reader.close()
|
|
connection.close()
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
raise TunModeError("用户取消了开启代理")
|
|
self._request("start", payload)
|
|
except Exception:
|
|
self.close()
|
|
raise
|
|
finally:
|
|
listener.close()
|
|
|
|
def _request(self, command: str, payload: dict[str, Any] | None = None) -> Any:
|
|
if self._socket is None or self._reader is None:
|
|
raise TunModeError("TUN 工作进程未连接")
|
|
try:
|
|
_send(self._socket, {"command": command, "payload": payload or {}})
|
|
while True:
|
|
message = _receive(self._reader)
|
|
if message.get("type") == "log":
|
|
callback = self.detail_log if message.get("detail") else self.log
|
|
callback(str(message.get("message") or ""))
|
|
continue
|
|
if message.get("type") != "result":
|
|
raise TunModeError("TUN 工作进程响应无效")
|
|
if not message.get("ok"):
|
|
raise TunModeError(str(message.get("error") or "TUN 操作失败"))
|
|
return message.get("value")
|
|
except TunModeError:
|
|
raise
|
|
except OSError as exc:
|
|
raise TunModeError(f"TUN 工作进程连接中止: {exc}") from exc
|
|
|
|
def status(self) -> dict[str, Any]:
|
|
with self._lock:
|
|
if self._socket is None:
|
|
return {"running": False, "pid": None, "exit_code": None}
|
|
try:
|
|
return self._request("status")
|
|
except (OSError, TunModeError, ValueError):
|
|
self.close()
|
|
return {"running": False, "pid": None, "exit_code": None}
|
|
|
|
def traffic(self) -> tuple[int, int, int, int]:
|
|
with self._lock:
|
|
if self._socket is None:
|
|
return 0, 0, 0, 0
|
|
try:
|
|
return tuple(self._request("traffic"))
|
|
except (OSError, TunModeError, ValueError):
|
|
return 0, 0, 0, 0
|
|
|
|
def stop(self) -> None:
|
|
with self._lock:
|
|
try:
|
|
if self._socket is not None:
|
|
self._request("stop")
|
|
finally:
|
|
self.close()
|
|
|
|
def close(self) -> None:
|
|
with self._lock:
|
|
if self._socket is not None:
|
|
try:
|
|
self._request("quit")
|
|
except (OSError, TunModeError, ValueError):
|
|
pass
|
|
reader, connection = self._reader, self._socket
|
|
self._reader = self._socket = None
|
|
if reader is not None:
|
|
reader.close()
|
|
if connection is not None:
|
|
connection.close()
|
|
|
|
|
|
def run_tun_worker(port: int, token: str, backend_factory,
|
|
start_backend: Callable[[Any, dict[str, Any]], None]) -> None:
|
|
"""管理员进程入口。主 GUI 断开后，必定清理本进程建立的网络状态。"""
|
|
connection = socket.create_connection(("127.0.0.1", int(port)), timeout=15)
|
|
connection.settimeout(None)
|
|
reader = connection.makefile("rb")
|
|
send_lock = threading.Lock()
|
|
_send(connection, {"type": "hello", "token": token}, send_lock)
|
|
|
|
def emit(message: str, detail: bool = False) -> None:
|
|
try:
|
|
_send(connection, {"type": "log", "message": str(message),
|
|
"detail": bool(detail)}, send_lock)
|
|
except OSError:
|
|
pass
|
|
|
|
class _WorkerLogger:
|
|
def __call__(self, message: str) -> None:
|
|
emit(message)
|
|
|
|
def detail(self, message: str) -> None:
|
|
emit(message, detail=True)
|
|
|
|
logger = _WorkerLogger()
|
|
|
|
backend = None
|
|
try:
|
|
while True:
|
|
try:
|
|
request = _receive(reader)
|
|
except TunModeError:
|
|
break
|
|
command = request.get("command")
|
|
try:
|
|
if command == "start":
|
|
if not isinstance(request.get("payload"), dict):
|
|
raise TunModeError("TUN 启动参数无效")
|
|
# 延迟创建后端，使初始化失败也能通过控制协议返回具体原因，
|
|
# 避免 GUI 只看到 Windows 的 10053 连接中止错误。
|
|
if backend is None:
|
|
backend = backend_factory(logger)
|
|
start_backend(backend, request["payload"])
|
|
value = None
|
|
elif command == "stop":
|
|
backend.stop()
|
|
value = None
|
|
elif command == "status":
|
|
value = backend.status()
|
|
elif command == "traffic":
|
|
value = backend.traffic()
|
|
elif command == "quit":
|
|
backend.stop()
|
|
value = None
|
|
else:
|
|
raise TunModeError("未知的 TUN 控制命令")
|
|
_send(connection, {"type": "result", "ok": True, "value": value}, send_lock)
|
|
except Exception as exc:
|
|
_send(connection, {"type": "result", "ok": False, "error": str(exc)}, send_lock)
|
|
if command == "quit":
|
|
break
|
|
finally:
|
|
if backend is not None:
|
|
backend.stop()
|
|
reader.close()
|
|
connection.close()
|