267 lines
9.3 KiB
Python
267 lines
9.3 KiB
Python
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import subprocess
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
|
||
from app.core.config import settings
|
||
from app.core.redis_client import get_redis
|
||
from app.services.runtime_settings_service import get_runtime_settings
|
||
|
||
|
||
WORKER_CONTROL_CHANNEL = "domain_tool:worker_control"
|
||
WORKER_PENDING_COMMAND_KEY = "domain_tool:worker_pending_command"
|
||
|
||
|
||
def _domain_root() -> Path:
|
||
return Path(settings.domain_root)
|
||
|
||
|
||
def _runtime_config() -> dict:
|
||
return get_runtime_settings()
|
||
|
||
|
||
def _run_powershell(command: str, timeout: int = 20) -> subprocess.CompletedProcess[str]:
|
||
return subprocess.run(
|
||
["powershell", "-NoProfile", "-ExecutionPolicy", "Bypass", "-Command", command],
|
||
capture_output=True,
|
||
text=True,
|
||
timeout=timeout,
|
||
)
|
||
|
||
|
||
def _run_shell(command: list[str], timeout: int = 20) -> subprocess.CompletedProcess[str]:
|
||
return subprocess.run(command, capture_output=True, text=True, timeout=timeout)
|
||
|
||
|
||
def _run_systemctl(command: list[str], timeout: int = 20, require_sudo: bool = True) -> subprocess.CompletedProcess[str]:
|
||
systemctl_command = ["systemctl", *command]
|
||
if os.geteuid() == 0 or not require_sudo:
|
||
return _run_shell(systemctl_command, timeout=timeout)
|
||
return _run_shell(["sudo", "-n", *systemctl_command], timeout=timeout)
|
||
|
||
|
||
def _parse_systemd_timestamp(raw_timestamp: str) -> str:
|
||
raw_timestamp = (raw_timestamp or "").strip()
|
||
if not raw_timestamp:
|
||
return ""
|
||
try:
|
||
parsed = datetime.strptime(raw_timestamp, "%a %Y-%m-%d %H:%M:%S %Z")
|
||
return parsed.strftime("%Y-%m-%d %H:%M:%S")
|
||
except ValueError:
|
||
pass
|
||
try:
|
||
parsed = datetime.strptime(raw_timestamp.rsplit(" ", 1)[0], "%a %Y-%m-%d %H:%M:%S")
|
||
return parsed.strftime("%Y-%m-%d %H:%M:%S")
|
||
except ValueError:
|
||
return ""
|
||
|
||
|
||
def probe_systemd_service(service_name: str, *, mode: str = "linux-systemd") -> dict:
|
||
result = _run_systemctl(
|
||
[
|
||
"show",
|
||
service_name,
|
||
"--no-page",
|
||
"--property=ActiveState,SubState,MainPID,ExecMainStartTimestamp,ActiveEnterTimestamp",
|
||
],
|
||
require_sudo=False,
|
||
)
|
||
output = (result.stdout or result.stderr or "").strip()
|
||
if result.returncode != 0:
|
||
return {
|
||
"mode": mode,
|
||
"service_name": service_name,
|
||
"running": False,
|
||
"process_count": 0,
|
||
"latest_start_time": "",
|
||
"message": output or f"systemd service {service_name} not available",
|
||
}
|
||
|
||
data: dict[str, str] = {}
|
||
for line in output.splitlines():
|
||
if "=" in line:
|
||
key, value = line.split("=", 1)
|
||
data[key] = value
|
||
main_pid = int(data.get("MainPID", "0") or 0)
|
||
active_state = data.get("ActiveState", "")
|
||
sub_state = data.get("SubState", "")
|
||
latest_start_time = ""
|
||
for raw_timestamp in (data.get("ExecMainStartTimestamp", ""), data.get("ActiveEnterTimestamp", "")):
|
||
latest_start_time = _parse_systemd_timestamp(raw_timestamp)
|
||
if latest_start_time:
|
||
break
|
||
return {
|
||
"mode": mode,
|
||
"service_name": service_name,
|
||
"running": active_state == "active",
|
||
"process_count": 1 if main_pid > 0 else 0,
|
||
"latest_start_time": latest_start_time,
|
||
"message": f"{active_state}/{sub_state}" if active_state else "",
|
||
}
|
||
|
||
|
||
def _windows_runtime() -> dict:
|
||
command = """
|
||
$targets = Get-CimInstance Win32_Process -Filter "name='python.exe'" |
|
||
Where-Object { $_.CommandLine -like '*detect_worker.py*' } |
|
||
Select-Object ProcessId, CommandLine
|
||
if (-not $targets) {
|
||
Write-Output '{"running":false,"process_count":0,"latest_start_time":"","mode":"windows-local"}'
|
||
exit 0
|
||
}
|
||
$latest = $null
|
||
foreach ($item in $targets) {
|
||
try {
|
||
$proc = Get-Process -Id $item.ProcessId -ErrorAction Stop
|
||
if (-not $latest -or $proc.StartTime -gt $latest.StartTime) {
|
||
$latest = $proc
|
||
}
|
||
} catch {}
|
||
}
|
||
$payload = @{
|
||
running = $true
|
||
process_count = @($targets).Count
|
||
latest_start_time = if ($latest) { $latest.StartTime.ToString('yyyy-MM-dd HH:mm:ss') } else { '' }
|
||
mode = 'windows-local'
|
||
} | ConvertTo-Json -Compress
|
||
Write-Output $payload
|
||
"""
|
||
result = _run_powershell(command)
|
||
output = (result.stdout or "").strip()
|
||
if result.returncode != 0 or not output:
|
||
return {
|
||
"mode": "windows-local",
|
||
"running": False,
|
||
"process_count": 0,
|
||
"latest_start_time": "",
|
||
"message": (result.stderr or result.stdout or "worker runtime probe failed").strip(),
|
||
}
|
||
try:
|
||
payload = json.loads(output)
|
||
except json.JSONDecodeError:
|
||
return {
|
||
"mode": "windows-local",
|
||
"running": False,
|
||
"process_count": 0,
|
||
"latest_start_time": "",
|
||
"message": output,
|
||
}
|
||
payload.setdefault("message", "")
|
||
return payload
|
||
|
||
|
||
def _linux_runtime() -> dict:
|
||
runtime = _runtime_config()
|
||
service_name = runtime["worker_service_name"]
|
||
return probe_systemd_service(service_name, mode="linux-systemd")
|
||
|
||
|
||
def detect_sync_agent_runtime() -> dict:
|
||
runtime = _runtime_config()
|
||
worker_mode = runtime.get("worker_mode", settings.worker_mode)
|
||
service_name = runtime.get("sync_agent_service_name", settings.sync_agent_service_name)
|
||
if worker_mode != "linux-systemd":
|
||
return {
|
||
"mode": worker_mode,
|
||
"service_name": service_name,
|
||
"running": False,
|
||
"process_count": 0,
|
||
"latest_start_time": "",
|
||
"message": "sync-agent 仅在 Linux systemd 多机部署中使用",
|
||
}
|
||
return probe_systemd_service(service_name, mode="linux-systemd")
|
||
|
||
|
||
def detect_worker_runtime() -> dict:
|
||
runtime = _runtime_config()
|
||
worker_mode = runtime["worker_mode"]
|
||
if worker_mode == "linux-systemd":
|
||
return _linux_runtime()
|
||
if os.name == "nt":
|
||
return _windows_runtime()
|
||
return {
|
||
"mode": worker_mode,
|
||
"running": False,
|
||
"process_count": 0,
|
||
"latest_start_time": "",
|
||
"message": f"unsupported worker_mode: {worker_mode}",
|
||
}
|
||
|
||
|
||
def start_worker() -> tuple[bool, str]:
|
||
runtime = _runtime_config()
|
||
worker_mode = runtime["worker_mode"]
|
||
service_name = runtime["worker_service_name"]
|
||
if worker_mode == "linux-systemd":
|
||
result = _run_systemctl(["start", service_name], timeout=30)
|
||
if result.returncode != 0:
|
||
return False, (result.stderr or result.stdout or "启动 Linux Worker 失败").strip()
|
||
return True, f"Linux Worker 启动命令已发送: {service_name}"
|
||
|
||
if os.name != "nt":
|
||
return False, "当前仅实现 Windows 本地 Worker 启动，Linux 请将 worker_mode 设为 linux-systemd。"
|
||
|
||
script_path = _domain_root() / "start_worker.ps1"
|
||
if not script_path.exists():
|
||
return False, f"未找到启动脚本: {script_path}"
|
||
|
||
command = (
|
||
"Start-Process powershell "
|
||
f"-ArgumentList '-ExecutionPolicy Bypass -File \"{script_path}\"' "
|
||
f"-WorkingDirectory '{_domain_root()}'"
|
||
)
|
||
result = _run_powershell(command)
|
||
if result.returncode != 0:
|
||
return False, (result.stderr or result.stdout or "启动检测端失败").strip()
|
||
return True, "检测端启动命令已发送"
|
||
|
||
|
||
def stop_worker() -> tuple[bool, str]:
|
||
runtime = _runtime_config()
|
||
worker_mode = runtime["worker_mode"]
|
||
service_name = runtime["worker_service_name"]
|
||
if worker_mode == "linux-systemd":
|
||
result = _run_systemctl(["stop", service_name], timeout=30)
|
||
if result.returncode != 0:
|
||
return False, (result.stderr or result.stdout or "停止 Linux Worker 失败").strip()
|
||
return True, f"Linux Worker 停止命令已发送: {service_name}"
|
||
|
||
if os.name != "nt":
|
||
return False, "当前仅实现 Windows 本地 Worker 停止，Linux 请将 worker_mode 设为 linux-systemd。"
|
||
|
||
command = """
|
||
$targets = Get-CimInstance Win32_Process -Filter "name='python.exe'" |
|
||
Where-Object { $_.CommandLine -like '*detect_worker.py*' } |
|
||
Select-Object -ExpandProperty ProcessId
|
||
if (-not $targets) {
|
||
Write-Output 'NO_PROCESS'
|
||
exit 0
|
||
}
|
||
$targets | ForEach-Object { Stop-Process -Id $_ -Force }
|
||
Write-Output ('STOPPED:' + (($targets | Measure-Object).Count))
|
||
"""
|
||
result = _run_powershell(command)
|
||
output = (result.stdout or result.stderr or "").strip()
|
||
if result.returncode != 0:
|
||
return False, output or "停止检测端失败"
|
||
if "NO_PROCESS" in output:
|
||
return True, "当前没有运行中的检测端进程"
|
||
return True, output or "检测端已停止"
|
||
|
||
|
||
def send_worker_command(action: str, payload: dict | None = None) -> tuple[bool, str]:
|
||
try:
|
||
redis_client = get_redis()
|
||
command_payload = {"action": action}
|
||
if payload:
|
||
command_payload.update(payload)
|
||
serialized = json.dumps(command_payload, ensure_ascii=False)
|
||
redis_client.set(WORKER_PENDING_COMMAND_KEY, serialized, ex=120)
|
||
redis_client.publish(WORKER_CONTROL_CHANNEL, serialized)
|
||
return True, f"已发送 Worker 控制指令: {action}"
|
||
except Exception as exc:
|
||
return False, f"发送 Worker 控制指令失败: {exc}"
|