350 lines
13 KiB
Python
350 lines
13 KiB
Python
#!/usr/bin/env python3
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import math
|
||
import os
|
||
import re
|
||
import socket
|
||
import sys
|
||
import urllib.error
|
||
import urllib.request
|
||
from datetime import datetime, timedelta, timezone
|
||
from typing import Any, NoReturn
|
||
|
||
try:
|
||
from zoneinfo import ZoneInfo
|
||
except ImportError: # pragma: no cover - Python 3.8 fallback
|
||
ZoneInfo = None # type: ignore[assignment,misc]
|
||
|
||
sys.stderr = sys.stdout
|
||
|
||
SCHEDULE_TYPES = ("delay_once", "daily", "weekly", "cn_workday")
|
||
MAX_DELAY_SECONDS = 24 * 60 * 60
|
||
|
||
|
||
class AmbiguousCreateError(RuntimeError):
|
||
"""The POST may have reached the server, so retrying could create a duplicate."""
|
||
|
||
|
||
class SkillArgumentParser(argparse.ArgumentParser):
|
||
def error(self, message: str) -> NoReturn:
|
||
raise ValueError(f"参数错误：{message}")
|
||
|
||
|
||
def _shanghai_timezone():
|
||
if ZoneInfo is not None:
|
||
try:
|
||
return ZoneInfo("Asia/Shanghai")
|
||
except Exception:
|
||
pass
|
||
return timezone(timedelta(hours=8), name="Asia/Shanghai")
|
||
|
||
|
||
SHANGHAI_TZ = _shanghai_timezone()
|
||
|
||
|
||
def _parse_args(argv: list[str]) -> argparse.Namespace:
|
||
parser = SkillArgumentParser(description="创建当前微信会话的定时任务")
|
||
parser.add_argument("--name", required=True)
|
||
parser.add_argument(
|
||
"--schedule-type",
|
||
"--schedule_type",
|
||
dest="schedule_type",
|
||
choices=SCHEDULE_TYPES,
|
||
required=True,
|
||
)
|
||
parser.add_argument("--content", default="")
|
||
parser.add_argument("--ai-prompt", "--ai_prompt", dest="ai_prompt", default="")
|
||
parser.add_argument("--time", default="")
|
||
parser.add_argument("--weekday", action="append", default=[])
|
||
parser.add_argument("--weekdays", action="append", default=[])
|
||
|
||
delay = parser.add_mutually_exclusive_group()
|
||
delay.add_argument("--delay-seconds", "--delay_seconds", dest="delay_seconds", type=int)
|
||
delay.add_argument("--delay-minutes", "--delay_minutes", dest="delay_minutes", type=int)
|
||
delay.add_argument("--delay-hours", "--delay_hours", dest="delay_hours", type=int)
|
||
delay.add_argument("--run-at", "--run_at", dest="run_at")
|
||
|
||
parser.add_argument("--dry-run", action="store_true")
|
||
return parser.parse_args(argv)
|
||
|
||
|
||
def _require_env(name: str) -> str:
|
||
value = os.environ.get(name, "").strip()
|
||
if not value:
|
||
raise ValueError(f"环境变量 {name} 未配置")
|
||
return value
|
||
|
||
|
||
def _conversation_identity() -> tuple[dict[str, str], dict[str, str], str]:
|
||
from_wechat_id = _require_env("ROBOT_FROM_WX_ID")
|
||
sender_wechat_id = os.environ.get("ROBOT_SENDER_WX_ID", "").strip()
|
||
|
||
if from_wechat_id.endswith("@chatroom"):
|
||
if not sender_wechat_id:
|
||
raise ValueError("当前群聊缺少消息发送人微信 ID")
|
||
target = {"wechat_id": from_wechat_id, "type": "chat_room"}
|
||
creator = {
|
||
"type": "chat_room",
|
||
"wechat_id": sender_wechat_id,
|
||
"chat_room_id": from_wechat_id,
|
||
}
|
||
return target, creator, "当前群聊"
|
||
|
||
creator_wechat_id = sender_wechat_id or from_wechat_id
|
||
target = {"wechat_id": from_wechat_id, "type": "friend"}
|
||
creator = {"type": "friend", "wechat_id": creator_wechat_id, "chat_room_id": ""}
|
||
return target, creator, "当前私聊"
|
||
|
||
|
||
def _validate_clock(value: str) -> str:
|
||
if not re.fullmatch(r"(?:[01]\d|2[0-3]):[0-5]\d", value):
|
||
raise ValueError("执行时间必须使用 HH:mm 格式，例如 08:30")
|
||
return value
|
||
|
||
|
||
def _parse_weekdays(values: list[str]) -> list[int]:
|
||
parsed: list[int] = []
|
||
for raw in values:
|
||
text = raw.strip()
|
||
if not text:
|
||
continue
|
||
if text.startswith("["):
|
||
try:
|
||
items = json.loads(text)
|
||
except json.JSONDecodeError as exc:
|
||
raise ValueError(f"weekdays JSON 格式错误: {exc.msg}") from exc
|
||
if not isinstance(items, list):
|
||
raise ValueError("weekdays 必须是 JSON 数组或逗号分隔列表")
|
||
else:
|
||
items = [item.strip() for item in text.split(",") if item.strip()]
|
||
|
||
for item in items:
|
||
if isinstance(item, bool):
|
||
raise ValueError("星期必须是 1 到 7 的整数")
|
||
try:
|
||
weekday = int(item)
|
||
except (TypeError, ValueError) as exc:
|
||
raise ValueError("星期必须是 1 到 7 的整数") from exc
|
||
if weekday < 1 or weekday > 7:
|
||
raise ValueError("星期必须在 1 到 7 之间，1 代表周一，7 代表周日")
|
||
parsed.append(weekday)
|
||
|
||
result = sorted(set(parsed))
|
||
if not result:
|
||
raise ValueError("每周任务至少需要一个星期")
|
||
return result
|
||
|
||
|
||
def _parse_run_at(value: str, now: datetime) -> int:
|
||
text = value.strip()
|
||
parsed: datetime | None = None
|
||
for pattern in ("%Y-%m-%d %H:%M", "%Y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M", "%Y-%m-%dT%H:%M:%S"):
|
||
try:
|
||
parsed = datetime.strptime(text, pattern).replace(tzinfo=SHANGHAI_TZ)
|
||
break
|
||
except ValueError:
|
||
continue
|
||
if parsed is None:
|
||
raise ValueError("run-at 必须使用 YYYY-MM-DD HH:mm 或 YYYY-MM-DD HH:mm:ss 格式")
|
||
|
||
delay_seconds = math.ceil((parsed - now).total_seconds())
|
||
if delay_seconds < 1:
|
||
raise ValueError("一次性任务的执行时间必须晚于当前时间")
|
||
if delay_seconds > MAX_DELAY_SECONDS:
|
||
raise ValueError("一次性任务最多只能设置到未来 24 小时内")
|
||
return delay_seconds
|
||
|
||
|
||
def _delay_seconds(args: argparse.Namespace, now: datetime) -> int:
|
||
if args.delay_seconds is not None:
|
||
seconds = args.delay_seconds
|
||
elif args.delay_minutes is not None:
|
||
seconds = args.delay_minutes * 60
|
||
elif args.delay_hours is not None:
|
||
seconds = args.delay_hours * 60 * 60
|
||
elif args.run_at is not None:
|
||
seconds = _parse_run_at(args.run_at, now)
|
||
else:
|
||
raise ValueError(
|
||
"一次性任务必须提供 delay-seconds、delay-minutes、delay-hours 或 run-at"
|
||
)
|
||
|
||
if seconds < 1 or seconds > MAX_DELAY_SECONDS:
|
||
raise ValueError("一次性任务延时必须在 1 秒到 24 小时之间")
|
||
return seconds
|
||
|
||
|
||
def _build_schedule_config(args: argparse.Namespace, now: datetime) -> dict[str, Any]:
|
||
has_delay = any(
|
||
value is not None
|
||
for value in (args.delay_seconds, args.delay_minutes, args.delay_hours, args.run_at)
|
||
)
|
||
weekday_values = [*args.weekday, *args.weekdays]
|
||
|
||
if args.schedule_type == "delay_once":
|
||
if args.time or weekday_values:
|
||
raise ValueError("一次性任务不能同时设置 time 或 weekdays")
|
||
return {"delay_seconds": _delay_seconds(args, now)}
|
||
|
||
if has_delay:
|
||
raise ValueError(f"{args.schedule_type} 任务不能设置延时参数")
|
||
clock = _validate_clock(args.time.strip())
|
||
|
||
if args.schedule_type == "weekly":
|
||
return {"time": clock, "weekdays": _parse_weekdays(weekday_values)}
|
||
if weekday_values:
|
||
raise ValueError(f"{args.schedule_type} 任务不能设置 weekdays")
|
||
return {"time": clock}
|
||
|
||
|
||
def _build_payload(args: argparse.Namespace) -> tuple[dict[str, Any], str]:
|
||
name = args.name.strip()
|
||
content = args.content.strip()
|
||
ai_prompt = args.ai_prompt.strip()
|
||
|
||
if not name:
|
||
raise ValueError("任务名称不能为空")
|
||
if len(name) > 100:
|
||
raise ValueError("任务名称不能超过 100 个字符")
|
||
if len(content) > 500:
|
||
raise ValueError("固定提醒文本不能超过 500 个字符")
|
||
if not content and not ai_prompt:
|
||
raise ValueError("content 和 ai-prompt 至少需要提供一项")
|
||
|
||
target, creator, target_label = _conversation_identity()
|
||
schedule_config = _build_schedule_config(args, datetime.now(SHANGHAI_TZ))
|
||
payload: dict[str, Any] = {
|
||
"name": name,
|
||
"enabled": True,
|
||
"schedule_type": args.schedule_type,
|
||
"schedule_config": schedule_config,
|
||
"targets": [target],
|
||
"fixed_text": content,
|
||
"images": [],
|
||
"ai_prompt": ai_prompt,
|
||
"creator": creator,
|
||
}
|
||
return payload, target_label
|
||
|
||
|
||
def _post_json(url: str, payload: dict[str, Any], timeout: int = 30) -> dict[str, Any]:
|
||
body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
|
||
request = urllib.request.Request(
|
||
url,
|
||
data=body,
|
||
headers={"Content-Type": "application/json"},
|
||
method="POST",
|
||
)
|
||
|
||
try:
|
||
with urllib.request.urlopen(request, timeout=timeout) as response:
|
||
response_text = response.read().decode("utf-8", errors="replace")
|
||
except urllib.error.HTTPError as exc:
|
||
error_body = exc.read().decode("utf-8", errors="replace")
|
||
raise RuntimeError(f"接口返回 HTTP {exc.code}: {error_body}") from exc
|
||
except (urllib.error.URLError, TimeoutError, socket.timeout, ConnectionError) as exc:
|
||
raise AmbiguousCreateError(
|
||
f"请求定时任务接口失败，结果可能未知：{exc}。请先在定时任务列表核对，勿直接重试"
|
||
) from exc
|
||
|
||
if not response_text.strip():
|
||
raise AmbiguousCreateError(
|
||
"定时任务接口返回空响应，任务可能已创建。请先在定时任务列表核对，勿直接重试"
|
||
)
|
||
try:
|
||
result = json.loads(response_text)
|
||
except json.JSONDecodeError as exc:
|
||
raise AmbiguousCreateError(
|
||
"定时任务接口返回了无效 JSON，任务可能已创建。请先在定时任务列表核对，勿直接重试"
|
||
) from exc
|
||
if not isinstance(result, dict):
|
||
raise AmbiguousCreateError(
|
||
"定时任务接口响应不是 JSON 对象，任务可能已创建。请先在定时任务列表核对，勿直接重试"
|
||
)
|
||
return result
|
||
|
||
|
||
def _unwrap_api_response(response: dict[str, Any]) -> dict[str, Any]:
|
||
if response.get("code") != 200:
|
||
message = str(response.get("message") or "接口返回未知业务错误")
|
||
if "任务已保存，但刷新调度器失败" in message:
|
||
raise AmbiguousCreateError(f"{message}。任务可能已经入库，请先核对任务列表，勿直接重试")
|
||
raise RuntimeError(message)
|
||
|
||
data = response.get("data")
|
||
if not isinstance(data, dict):
|
||
raise AmbiguousCreateError(
|
||
"接口成功响应中缺少任务数据，任务可能已创建。请先在定时任务列表核对，勿直接重试"
|
||
)
|
||
return data
|
||
|
||
|
||
def _format_next_run(timestamp: Any) -> str | None:
|
||
if timestamp is None:
|
||
return None
|
||
try:
|
||
value = int(timestamp)
|
||
except (TypeError, ValueError):
|
||
return None
|
||
return datetime.fromtimestamp(value, SHANGHAI_TZ).strftime("%Y-%m-%d %H:%M:%S Asia/Shanghai")
|
||
|
||
|
||
def _success_output(data: dict[str, Any], target_label: str) -> dict[str, Any]:
|
||
targets = data.get("targets")
|
||
first_target = targets[0] if isinstance(targets, list) and targets else {}
|
||
if not isinstance(first_target, dict):
|
||
first_target = {}
|
||
return {
|
||
"ok": True,
|
||
"message": "定时任务创建成功",
|
||
"task": {
|
||
"id": data.get("id"),
|
||
"name": data.get("name"),
|
||
"schedule_type": data.get("schedule_type"),
|
||
"schedule_summary": data.get("schedule_summary"),
|
||
"next_run_at": data.get("next_run_at"),
|
||
"next_run_time": _format_next_run(data.get("next_run_at")),
|
||
"fixed_text": data.get("fixed_text"),
|
||
"uses_ai_prompt": bool(data.get("ai_prompt")),
|
||
"target_label": target_label,
|
||
"target_type": first_target.get("type"),
|
||
},
|
||
}
|
||
|
||
|
||
def main(argv: list[str]) -> int:
|
||
try:
|
||
args = _parse_args(argv)
|
||
payload, target_label = _build_payload(args)
|
||
|
||
if args.dry_run:
|
||
print(
|
||
json.dumps(
|
||
{"ok": True, "dry_run": True, "target_label": target_label, "payload": payload},
|
||
ensure_ascii=False,
|
||
indent=2,
|
||
)
|
||
)
|
||
return 0
|
||
|
||
client_port = _require_env("ROBOT_WECHAT_CLIENT_PORT")
|
||
url = f"http://127.0.0.1:{client_port}/api/v1/robot/scheduled-tasks"
|
||
response = _post_json(url, payload)
|
||
data = _unwrap_api_response(response)
|
||
print(json.dumps(_success_output(data, target_label), ensure_ascii=False, indent=2))
|
||
return 0
|
||
except (ValueError, RuntimeError) as exc:
|
||
print(f"创建定时任务失败：{exc}")
|
||
return 1
|
||
except Exception as exc: # Keep script failures visible to the agent without hiding them.
|
||
print(f"创建定时任务失败：{type(exc).__name__}: {exc}")
|
||
return 1
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main(sys.argv[1:]))
|