getDomain/domain-api/app/services/worker_control_service.py

from __future__ import annotations

import json
import os
import subprocess
from pathlib import Path

from app.core.config import settings
from app.services.runtime_settings_service import get_runtime_settings


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 _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"]
    result = _run_shell(["systemctl", "show", service_name, "--no-page", "--property=ActiveState,SubState,MainPID"])
    output = (result.stdout or result.stderr or "").strip()
    if result.returncode != 0:
        return {
            "mode": "linux-systemd",
            "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", "")
    return {
        "mode": "linux-systemd",
        "running": active_state == "active",
        "process_count": 1 if main_pid > 0 else 0,
        "latest_start_time": "",
        "message": f"{active_state}/{sub_state}" if active_state else "",
    }


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_shell(["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_shell(["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 "检测端已停止"