PagerMaid_Plugins_Pyro/send_cron/main.py
xtaodada fb5a3f5baf
All checks were successful
Github commit to telegram / build (push) Successful in 13s
♻️ Reformat All Plugins
2024-09-28 22:35:08 +08:00

267 lines
8.3 KiB
Python

import contextlib
import datetime
from typing import Optional, List
from pagermaid.dependence import scheduler, sqlite
from pagermaid.enums import Message
from pagermaid.listener import listener
from pagermaid.services import bot
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}")