group_fqcd_jr/tools/install_offsite_worker_task.py

"""注册、查看和管理场外邮件 Worker 的 Windows 自启动任务。"""

from __future__ import annotations

import argparse
import os
import platform
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path
from xml.sax.saxutils import escape


DEFAULT_TASK_NAME = "NanfangFund-OffsiteWorker"
TRUE_VALUES = {"1", "true", "yes", "on", "y"}
REQUIRED_TRUE_FLAGS = (
    "OFFSITE_MAIL_WORKER_ENABLED",
    "OFFSITE_IMAP_ENABLED",
    "OFFSITE_OCR_ENABLED",
    "OFFSITE_DEEPSEEK_ENABLED",
)
REQUIRED_NON_EMPTY = (
    "OFFSITE_IMAP_HOST",
    "OFFSITE_IMAP_USERNAME",
    "OFFSITE_IMAP_PASSWORD",
    "OFFSITE_WORKER_USER_ID",
)


def main() -> None:
    parser = _build_parser()
    args = parser.parse_args()
    project_dir = Path(args.project_dir).resolve()
    task_name = str(args.task_name)

    if platform.system().lower() != "windows":
        raise SystemExit("本脚本只支持 Windows 任务计划程序。")

    if args.command == "install":
        install_task(
            task_name=task_name,
            project_dir=project_dir,
            python_exe=Path(args.python).resolve(),
            run_as=str(args.run_as),
            start_now=bool(args.start_now),
            skip_env_check=bool(args.skip_env_check),
        )
        return
    if args.command == "run":
        run_worker(project_dir)
        return
    if args.command == "uninstall":
        run_schtasks(["/Delete", "/TN", task_name, "/F"], check=True)
        print(f"已删除 Windows 任务：{task_name}")
        return
    if args.command == "status":
        run_schtasks(["/Query", "/TN", task_name, "/FO", "LIST", "/V"], check=True)
        return
    if args.command == "start":
        run_schtasks(["/Run", "/TN", task_name], check=True)
        print(f"已请求启动 Windows 任务：{task_name}")
        return
    if args.command == "stop":
        run_schtasks(["/End", "/TN", task_name], check=True)
        print(f"已请求停止 Windows 任务：{task_name}")
        return
    parser.print_help()


def _build_parser() -> argparse.ArgumentParser:
    project_dir = Path(__file__).resolve().parents[1]
    parser = argparse.ArgumentParser(
        description="将场外邮件 Worker 注册为 Windows 开机自启动任务。"
    )
    parser.add_argument(
        "command",
        choices=("install", "run", "uninstall", "status", "start", "stop"),
        help=(
            "install 注册或更新；run 由计划任务调用 Worker；status 查看；"
            "start 启动；stop 停止；uninstall 删除。"
        ),
    )
    parser.add_argument(
        "--task-name",
        default=DEFAULT_TASK_NAME,
        help=f"Windows 任务名称，默认：{DEFAULT_TASK_NAME}",
    )
    parser.add_argument(
        "--project-dir",
        default=str(project_dir),
        help="后端项目目录，必须包含 .env 和 app/worker。",
    )
    parser.add_argument(
        "--python",
        default=sys.executable,
        help="运行 Worker 的 Python 路径，默认使用当前 Python。",
    )
    parser.add_argument(
        "--run-as",
        choices=("system", "current"),
        default="system",
        help="system 表示开机即启动，通常需要管理员权限；current 表示当前用户登录后启动。",
    )
    parser.add_argument(
        "--start-now",
        action="store_true",
        help="注册成功后立即启动任务。",
    )
    parser.add_argument(
        "--skip-env-check",
        action="store_true",
        help="跳过 .env 开关检查，仅注册任务。",
    )
    return parser


def install_task(
    *,
    task_name: str,
    project_dir: Path,
    python_exe: Path,
    run_as: str,
    start_now: bool,
    skip_env_check: bool,
) -> None:
    validate_project(project_dir, python_exe)
    if not skip_env_check:
        validate_env(project_dir / ".env")
    if run_as == "system":
        install_system_task(
            task_name=task_name,
            project_dir=project_dir,
            python_exe=python_exe,
            start_now=start_now,
        )
        return
    xml = build_task_xml(project_dir=project_dir, python_exe=python_exe, run_as=run_as)
    with tempfile.NamedTemporaryFile("w", suffix=".xml", delete=False, encoding="utf-16") as file:
        file.write(xml)
        xml_path = Path(file.name)
    try:
        run_schtasks(["/Create", "/TN", task_name, "/XML", str(xml_path), "/F"], check=True)
    finally:
        with contextlib_suppress_os_error():
            xml_path.unlink()
    print(f"已注册或更新 Windows 任务：{task_name}")
    print(f"工作目录：{project_dir}")
    print(f"Python：{python_exe}")
    print("任务动作：python -m app.worker")
    if run_as == "system":
        print("启动方式：系统启动时自动运行。")
    else:
        print("启动方式：当前用户登录时自动运行。")
    if start_now:
        run_schtasks(["/Run", "/TN", task_name], check=True)
        print(f"已请求立即启动任务：{task_name}")


def install_system_task(
    *,
    task_name: str,
    project_dir: Path,
    python_exe: Path,
    start_now: bool,
) -> None:
    task_run = build_system_task_command(project_dir=project_dir, python_exe=python_exe)
    run_schtasks(
        [
            "/Create",
            "/TN",
            task_name,
            "/SC",
            "ONSTART",
            "/RU",
            "SYSTEM",
            "/RL",
            "HIGHEST",
            "/TR",
            task_run,
            "/F",
        ],
        check=True,
    )
    print(f"已注册或更新 Windows 任务：{task_name}")
    print(f"工作目录：{project_dir}")
    print(f"Python：{python_exe}")
    print("任务动作：python -m app.worker")
    print("启动方式：系统启动时自动运行。")
    if start_now:
        run_schtasks(["/Run", "/TN", task_name], check=True)
        print(f"已请求立即启动任务：{task_name}")


def build_system_task_command(*, project_dir: Path, python_exe: Path) -> str:
    """生成不依赖任务计划程序工作目录设置的脚本启动命令。"""
    launcher = Path(__file__).resolve()
    return f'"{python_exe}" "{launcher}" run --project-dir "{project_dir}"'


def run_worker(project_dir: Path) -> None:
    """由 Windows 计划任务调用，切换目录后启动既有 Worker 入口。"""
    os.chdir(project_dir)
    sys.path.insert(0, str(project_dir))
    from app.worker.__main__ import main as worker_main

    sys.argv = ["app.worker"]
    worker_main()


def validate_project(project_dir: Path, python_exe: Path) -> None:
    missing = []
    if not project_dir.exists():
        missing.append(f"项目目录不存在：{project_dir}")
    if not (project_dir / "app" / "worker" / "__main__.py").exists():
        missing.append("未找到 app/worker/__main__.py")
    if not (project_dir / ".env").exists():
        missing.append("未找到 .env")
    if not python_exe.exists():
        missing.append(f"Python 不存在：{python_exe}")
    if shutil.which("schtasks.exe") is None:
        missing.append("未找到 schtasks.exe")
    if missing:
        raise SystemExit("\n".join(missing))


def validate_env(env_path: Path) -> None:
    values = parse_env(env_path)
    errors = []
    for key in REQUIRED_TRUE_FLAGS:
        if values.get(key, "").strip().strip('"').strip("'").lower() not in TRUE_VALUES:
            errors.append(f"{key} 必须为 true")
    for key in REQUIRED_NON_EMPTY:
        if not values.get(key, "").strip().strip('"').strip("'"):
            errors.append(f"{key} 不能为空")
    if errors:
        lines = [
            ".env 未满足场外 Worker 自动启动条件，已中止注册：",
            *[f"- {item}" for item in errors],
            "如只是预注册任务，可追加 --skip-env-check。",
        ]
        raise SystemExit("\n".join(lines))


def parse_env(env_path: Path) -> dict[str, str]:
    values: dict[str, str] = {}
    for raw_line in env_path.read_text(encoding="utf-8").splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        values[key.strip()] = value.strip()
    return values


def build_task_xml(*, project_dir: Path, python_exe: Path, run_as: str) -> str:
    command = escape(str(python_exe))
    working_dir = escape(str(project_dir))
    user_id = "SYSTEM" if run_as == "system" else escape(os.environ.get("USERNAME", ""))
    logon_type = "ServiceAccount" if run_as == "system" else "InteractiveToken"
    trigger = "BootTrigger" if run_as == "system" else "LogonTrigger"
    return f"""<?xml version="1.0" encoding="UTF-16"?>
<Task version="1.4" xmlns="http://schemas.microsoft.com/windows/2004/02/mit/task">
  <RegistrationInfo>
    <Description>南方基金场外申购赎回邮件 Worker 自动启动任务</Description>
  </RegistrationInfo>
  <Triggers>
    <{trigger}>
      <Enabled>true</Enabled>
    </{trigger}>
  </Triggers>
  <Principals>
    <Principal id="Author">
      <UserId>{user_id}</UserId>
      <LogonType>{logon_type}</LogonType>
      <RunLevel>HighestAvailable</RunLevel>
    </Principal>
  </Principals>
  <Settings>
    <MultipleInstancesPolicy>IgnoreNew</MultipleInstancesPolicy>
    <DisallowStartIfOnBatteries>false</DisallowStartIfOnBatteries>
    <StopIfGoingOnBatteries>false</StopIfGoingOnBatteries>
    <AllowHardTerminate>true</AllowHardTerminate>
    <StartWhenAvailable>true</StartWhenAvailable>
    <RunOnlyIfNetworkAvailable>true</RunOnlyIfNetworkAvailable>
    <IdleSettings>
      <StopOnIdleEnd>false</StopOnIdleEnd>
      <RestartOnIdle>false</RestartOnIdle>
    </IdleSettings>
    <AllowStartOnDemand>true</AllowStartOnDemand>
    <Enabled>true</Enabled>
    <Hidden>false</Hidden>
    <RunOnlyIfIdle>false</RunOnlyIfIdle>
    <DisallowStartOnRemoteAppSession>false</DisallowStartOnRemoteAppSession>
    <UseUnifiedSchedulingEngine>true</UseUnifiedSchedulingEngine>
    <WakeToRun>false</WakeToRun>
    <ExecutionTimeLimit>PT0S</ExecutionTimeLimit>
    <Priority>7</Priority>
    <RestartOnFailure>
      <Interval>PT1M</Interval>
      <Count>3</Count>
    </RestartOnFailure>
  </Settings>
  <Actions Context="Author">
    <Exec>
      <Command>{command}</Command>
      <Arguments>-m app.worker</Arguments>
      <WorkingDirectory>{working_dir}</WorkingDirectory>
    </Exec>
  </Actions>
</Task>
"""


def run_schtasks(args: list[str], *, check: bool) -> subprocess.CompletedProcess[str]:
    command = ["schtasks.exe", *args]
    result = subprocess.run(command, check=False, text=True, capture_output=True)
    if result.stdout.strip():
        print(result.stdout.strip())
    if result.stderr.strip():
        print(result.stderr.strip(), file=sys.stderr)
    if check and result.returncode != 0:
        hint = ""
        if "/Create" in args:
            hint = "\n如果使用 --run-as system，请用管理员身份运行 PowerShell 或终端。"
        raise SystemExit(f"schtasks 执行失败，退出码：{result.returncode}{hint}")
    return result


class contextlib_suppress_os_error:
    def __enter__(self) -> None:
        return None

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        traceback: object | None,
    ) -> bool:
        return exc_type is not None and issubclass(exc_type, OSError)


if __name__ == "__main__":
    main()