mirror of
https://github.com/TeamPGM/PagerMaid_Plugins_Pyro.git
synced 2024-11-22 07:08:19 +00:00
267 lines
8.3 KiB
Python
267 lines
8.3 KiB
Python
import contextlib
|
|
import datetime
|
|
|
|
from typing import Optional, List
|
|
|
|
from pagermaid import bot
|
|
from pagermaid.listener import listener
|
|
from pagermaid.scheduler import scheduler
|
|
from pagermaid.single_utils import sqlite, Message
|
|
from pagermaid.utils import alias_command
|
|
|
|
|
|
class SendTask:
|
|
task_id: Optional[int]
|
|
cid: int
|
|
msg: str
|
|
cron: str
|
|
pause: bool
|
|
|
|
def __init__(
|
|
self,
|
|
task_id: int,
|
|
cid: int = 0,
|
|
msg: str = "",
|
|
cron: str = "",
|
|
pause: bool = False,
|
|
):
|
|
self.task_id = task_id
|
|
self.cid = cid
|
|
self.msg = msg
|
|
self.cron = cron
|
|
self.pause = pause
|
|
|
|
def export(self):
|
|
return {
|
|
"task_id": self.task_id,
|
|
"cid": self.cid,
|
|
"msg": self.msg,
|
|
"cron": self.cron,
|
|
"pause": self.pause,
|
|
}
|
|
|
|
def get_job(self):
|
|
return scheduler.get_job(f"send_cron|{self.cid}|{self.task_id}")
|
|
|
|
def remove_job(self):
|
|
if self.get_job():
|
|
scheduler.remove_job(f"send_cron|{self.cid}|{self.task_id}")
|
|
|
|
def export_str(self, show_all: bool = False):
|
|
text = f"<code>{self.task_id}</code> - " f"<code>{self.cron}</code> -"
|
|
if job := self.get_job():
|
|
time: datetime.datetime = job.next_run_time
|
|
text += f"<code>{time.strftime('%Y-%m-%d %H:%M:%S')}</code> - "
|
|
else:
|
|
text += "<code>未运行</code> - "
|
|
if show_all:
|
|
text += f"<code>{self.cid}</code> - "
|
|
text += f"{self.msg}"
|
|
return text
|
|
|
|
@staticmethod
|
|
def parse_cron_kwargs(text: str):
|
|
data = text.split(" ")
|
|
return {
|
|
"second": data[0],
|
|
"minute": data[1],
|
|
"hour": data[2],
|
|
"day": data[3],
|
|
"month": data[4],
|
|
"day_of_week": data[5],
|
|
}
|
|
|
|
@property
|
|
def cron_kwargs(self):
|
|
return self.parse_cron_kwargs(self.cron)
|
|
|
|
def parse_task(self, text: str):
|
|
self.msg = "|".join(text.split("|")[1:]).strip()
|
|
if not self.msg:
|
|
raise ValueError("No message provided")
|
|
text = text.split("|")[0].strip()
|
|
data = text.split(" ")
|
|
if len(data) != 6:
|
|
raise ValueError("Invalid crontab format")
|
|
try:
|
|
scheduler._create_trigger("cron", self.parse_cron_kwargs(text))
|
|
except Exception as e:
|
|
raise ValueError(f"Invalid crontab format: {e}") from e
|
|
self.cron = text
|
|
|
|
|
|
class SendTasks:
|
|
tasks: List[SendTask]
|
|
|
|
def __init__(self):
|
|
self.tasks = []
|
|
|
|
def add(self, task: SendTask):
|
|
for i in self.tasks:
|
|
if i.task_id == task.task_id:
|
|
return
|
|
self.tasks.append(task)
|
|
|
|
def remove(self, task_id: int):
|
|
for task in self.tasks:
|
|
if task.task_id == task_id:
|
|
task.remove_job()
|
|
self.tasks.remove(task)
|
|
return True
|
|
return False
|
|
|
|
def get(self, task_id: int) -> Optional[SendTask]:
|
|
return next((task for task in self.tasks if task.task_id == task_id), None)
|
|
|
|
def get_all(self) -> List[SendTask]:
|
|
return self.tasks
|
|
|
|
def get_all_ids(self) -> List[int]:
|
|
return [task.task_id for task in self.tasks]
|
|
|
|
def print_all_tasks(self, show_all: bool = False, cid: int = 0) -> str:
|
|
return "\n".join(
|
|
task.export_str(show_all)
|
|
for task in self.tasks
|
|
if task.cid == cid or show_all
|
|
)
|
|
|
|
def save_to_file(self):
|
|
data = [task.export() for task in self.tasks]
|
|
sqlite["send_cron_tasks"] = data
|
|
|
|
def load_from_file(self):
|
|
data = sqlite.get("send_cron_tasks", [])
|
|
for i in data:
|
|
self.add(SendTask(**i))
|
|
|
|
def pause_task(self, task_id):
|
|
if task := self.get(task_id):
|
|
task.pause = True
|
|
task.remove_job()
|
|
self.save_to_file()
|
|
return True
|
|
return False
|
|
|
|
@staticmethod
|
|
async def send_message(task: SendTask, _: "SendTasks"):
|
|
with contextlib.suppress(Exception):
|
|
await bot.send_message(task.cid, task.msg)
|
|
|
|
def register_cron_task(self, task: SendTask):
|
|
scheduler.add_job(
|
|
self.send_message,
|
|
"cron",
|
|
id=f"send_cron|{task.cid}|{task.task_id}",
|
|
name=f"send_cron|{task.cid}|{task.task_id}",
|
|
args=[task, self],
|
|
**task.cron_kwargs,
|
|
)
|
|
|
|
def register_single_task(self, task: SendTask):
|
|
if task.pause:
|
|
return
|
|
self.register_cron_task(task)
|
|
|
|
def resume_task(self, task_id: int):
|
|
if task := self.get(task_id):
|
|
task.pause = False
|
|
self.register_single_task(task)
|
|
self.save_to_file()
|
|
return True
|
|
return False
|
|
|
|
def register_all_tasks(self):
|
|
for task in self.tasks:
|
|
self.register_single_task(task)
|
|
|
|
def get_next_task_id(self):
|
|
return max(task.task_id for task in self.tasks) + 1 if self.tasks else 1
|
|
|
|
|
|
send_cron_tasks = SendTasks()
|
|
send_cron_tasks.load_from_file()
|
|
send_cron_tasks.register_all_tasks()
|
|
|
|
send_help_msg = f"""
|
|
定时发送消息。
|
|
,{alias_command("send_cron")} crontab 表达式 | 消息内容
|
|
i.e.
|
|
,{alias_command("send_cron")} 59 59 23 * * * | 又是无所事事的一天呢。
|
|
,{alias_command("send_cron")} 0 * * * * * | 又过去了一分钟。
|
|
|
|
|
|
,{alias_command("send_cron")} rm 2 - 删除某个任务
|
|
,{alias_command("send_cron")} pause 1 - 暂停某个任务
|
|
,{alias_command("send_cron")} resume 1 - 恢复某个任务
|
|
,{alias_command("send_cron")} list <all> - 获取任务列表
|
|
"""
|
|
|
|
|
|
async def from_msg_get_task_id(message: Message):
|
|
uid = -1
|
|
try:
|
|
uid = int(message.parameter[1])
|
|
except ValueError:
|
|
await message.edit("请输入正确的参数")
|
|
message.continue_propagation()
|
|
ids = send_cron_tasks.get_all_ids()
|
|
if uid not in ids:
|
|
await message.edit("该任务不存在")
|
|
message.continue_propagation()
|
|
return uid
|
|
|
|
|
|
@listener(
|
|
command="send_cron",
|
|
parameters="crontab 表达式 | 消息内容",
|
|
need_admin=True,
|
|
description=f"定时发送消息\n请使用 ,{alias_command('send_cron')} h 查看可用命令",
|
|
)
|
|
async def send_cron(message: Message):
|
|
if message.arguments == "h" or len(message.parameter) == 0:
|
|
return await message.edit(send_help_msg)
|
|
if len(message.parameter) == 1:
|
|
if message.parameter[0] != "list":
|
|
return await message.edit("请输入正确的参数")
|
|
if send_cron_tasks.get_all_ids():
|
|
return await message.edit(
|
|
f"已注册的任务：\n\n{send_cron_tasks.print_all_tasks(show_all=False, cid=message.chat.id)}"
|
|
)
|
|
else:
|
|
return await message.edit("没有已注册的任务。")
|
|
if len(message.parameter) == 2:
|
|
if message.parameter[0] == "rm":
|
|
if uid := await from_msg_get_task_id(message):
|
|
send_cron_tasks.remove(uid)
|
|
send_cron_tasks.save_to_file()
|
|
send_cron_tasks.load_from_file()
|
|
return await message.edit(f"已删除任务 {uid}")
|
|
elif message.parameter[0] == "pause":
|
|
if uid := await from_msg_get_task_id(message):
|
|
send_cron_tasks.pause_task(uid)
|
|
return await message.edit(f"已暂停任务 {uid}")
|
|
elif message.parameter[0] == "resume":
|
|
if uid := await from_msg_get_task_id(message):
|
|
send_cron_tasks.resume_task(uid)
|
|
return await message.edit(f"已恢复任务 {uid}")
|
|
elif message.parameter[0] == "list":
|
|
if send_cron_tasks.get_all_ids():
|
|
return await message.edit(
|
|
f"已注册的任务：\n\n{send_cron_tasks.print_all_tasks(show_all=True)}"
|
|
)
|
|
else:
|
|
return await message.edit("没有已注册的任务。")
|
|
# add task
|
|
task = SendTask(send_cron_tasks.get_next_task_id())
|
|
task.cid = message.chat.id
|
|
try:
|
|
task.parse_task(message.arguments)
|
|
except Exception as e:
|
|
return await message.edit(f"参数错误：{e}")
|
|
send_cron_tasks.add(task)
|
|
send_cron_tasks.register_single_task(task)
|
|
send_cron_tasks.save_to_file()
|
|
send_cron_tasks.load_from_file()
|
|
await message.edit(f"已添加任务 {task.task_id}")
|