346 lines
11 KiB
Python
346 lines
11 KiB
Python
"""注册、查看和管理场外邮件 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()
|