"""注册、查看和管理场外邮件 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""" 奶龙基金场外申购赎回邮件 Worker 自动启动任务 <{trigger}> true {user_id} {logon_type} HighestAvailable IgnoreNew false false true true true false false true true false false false true false PT0S 7 PT1M 3 {command} -m app.worker {working_dir} """ 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()