feat(lite): merge TUN mode with IPv6 routing
This commit is contained in:
RoyirinRosmontis
@@ -0,0 +1,232 @@
|
||||
"""按需提权的 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]):
|
||||
self.command = command
|
||||
self.log = 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 工作进程未连接")
|
||||
_send(self._socket, {"command": command, "payload": payload or {}})
|
||||
while True:
|
||||
message = _receive(self._reader)
|
||||
if message.get("type") == "log":
|
||||
self.log(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")
|
||||
|
||||
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) -> None:
|
||||
try:
|
||||
_send(connection, {"type": "log", "message": str(message)}, send_lock)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
backend = None
|
||||
try:
|
||||
backend = backend_factory(emit)
|
||||
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 启动参数无效")
|
||||
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()
|
||||
Reference in New Issue
Block a user