nonebot2/commands/scheduler.py

353 lines
13 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import os
import re
from functools import wraps
import pytz
import requests
from apscheduler.schedulers.background import BackgroundScheduler
from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore
from apscheduler.executors.pool import ProcessPoolExecutor
from apscheduler.jobstores.base import JobLookupError
from command import CommandRegistry, hub as cmdhub
from command import CommandNotExistsError, CommandScopeError, CommandPermissionError
from commands import core
from little_shit import get_db_dir, get_command_args_start_flags, get_target
_db_url = 'sqlite:///' + os.path.join(get_db_dir(), 'scheduler.sqlite')
_scheduler = BackgroundScheduler(
jobstores={
'default': SQLAlchemyJobStore(url=_db_url)
},
executors={
'default': ProcessPoolExecutor(max_workers=5)
},
timezone=pytz.timezone('Asia/Shanghai')
)
_command_args_start_flags = get_command_args_start_flags()
_args_split_sep = '[ \n\t]'
_job_id_suffix_start = '@'
def _init():
_scheduler.start()
__registry__ = cr = CommandRegistry(init_func=_init)
class _InvalidTriggerArgsError(Exception):
pass
class _IncompleteArgsError(Exception):
pass
def _call_commands(job_id, command_list, ctx_msg, internal=False):
for command in command_list:
try:
cmdhub.call(command[0], command[1], ctx_msg)
except CommandNotExistsError:
core.echo('没有找到计划任务 %s 中的命令 %s' % (job_id, command[0]), ctx_msg, internal)
except CommandPermissionError:
core.echo('你没有权限执行计划任务 %s 中的命令 %s' % (job_id, command[0]), ctx_msg, internal)
except CommandScopeError as se:
core.echo(
'计划任务 %s 中的命令 %s 不支持 %s' % (job_id, command[0], se.msg_type),
ctx_msg, internal
)
def _check_target(func):
@wraps(func)
def wrapper(args_text, ctx_msg, internal=False, *args, **kwargs):
target = get_target(ctx_msg)
if not target:
_send_fail_to_get_target_msg(ctx_msg, internal)
return None
else:
return func(args_text, ctx_msg, internal, *args, **kwargs)
return wrapper
@cr.register('cron_check', 'cron-check', 'cron_test', 'cron-test')
def cron_check(args_text, ctx_msg):
cron = args_text.strip()
if not cron:
core.echo('请指定要检查的 Cron 时间表达式', ctx_msg)
return
resp = requests.post('http://tool.lu/crontab/ajax.html', data={'expression': cron})
if resp.status_code == 200:
data = resp.json()
if data.get('status') and 'dates' in data:
reply = '接下来 7 次的执行时间:\n' + '\n'.join(data['dates'])
core.echo(reply, ctx_msg)
return
core.echo('检查失败,可能因为表达式格式错误或服务器连接不上', ctx_msg)
@cr.register('add_job', 'add-job', 'add')
@cr.restrict(full_command_only=True, group_admin_only=True)
@_check_target
def add_job(args_text, ctx_msg, internal=False):
if args_text.strip() in ('', 'help', '-h', '--help') and not internal:
_send_add_job_help_msg(ctx_msg, internal)
return
args_text = args_text.lstrip()
try:
# Parse trigger args
trigger_args = {}
if args_text.startswith('-'):
# options mode
key_dict = {
'-M': 'minute',
'-H': 'hour',
'-d': 'day',
'-m': 'month',
'-w': 'day_of_week'
}
while args_text.startswith('-') and not args_text.startswith('--'):
try:
option, value, args_text = re.split(_args_split_sep, args_text, 2)
trigger_args[key_dict[option]] = value
args_text = args_text.lstrip()
except (ValueError, KeyError):
# Split failed or get key failed, which means format is not correct
raise _InvalidTriggerArgsError
else:
# cron mode
try:
trigger_args['minute'], \
trigger_args['hour'], \
trigger_args['day'], \
trigger_args['month'], \
trigger_args['day_of_week'], \
args_text = re.split(_args_split_sep, args_text, 5)
args_text = args_text.lstrip()
except ValueError:
# Split failed, which means format is not correct
raise _InvalidTriggerArgsError
# Parse '--multi' option
multi = False
if args_text.startswith('--multi '):
multi = True
tmp = re.split(_args_split_sep, args_text, 1)
if len(tmp) < 2:
raise _IncompleteArgsError
args_text = tmp[1].lstrip()
tmp = re.split(_args_split_sep, args_text, 1)
if len(tmp) < 2:
raise _IncompleteArgsError
job_id_without_suffix, command_raw = tmp
job_id = job_id_without_suffix + _job_id_suffix_start + get_target(ctx_msg)
command_list = []
if multi:
command_raw_list = command_raw.split('\n')
for cmd_raw in command_raw_list:
cmd_raw = cmd_raw.lstrip()
if not cmd_raw:
continue
tmp = re.split('|'.join(_command_args_start_flags), cmd_raw, 1)
if len(tmp) < 2:
tmp.append('')
command_list.append(tuple(tmp))
else:
command_raw = command_raw.lstrip()
tmp = re.split('|'.join(_command_args_start_flags), command_raw, 1)
if len(tmp) < 2:
tmp.append('')
command_list.append(tuple(tmp))
job_args = {
'job_id': job_id_without_suffix,
'command_list': command_list,
'ctx_msg': ctx_msg
}
job = _scheduler.add_job(_call_commands, kwargs=job_args, trigger='cron', **trigger_args,
id=job_id, replace_existing=True, misfire_grace_time=30)
_send_text('成功添加计划任务 ' + job_id_without_suffix, ctx_msg, internal)
if job:
job.id = job_id_without_suffix
return job
except _InvalidTriggerArgsError:
_send_add_job_trigger_args_invalid_msg(ctx_msg, internal)
except _IncompleteArgsError:
_send_add_job_incomplete_args_msg(ctx_msg, internal)
@cr.register('remove_job', 'remove-job', 'remove')
@cr.restrict(full_command_only=True, group_admin_only=True)
@_check_target
def remove_job(args_text, ctx_msg, internal=False):
job_id_without_suffix = args_text.strip()
if not job_id_without_suffix:
_send_text('请指定计划任务的 ID', ctx_msg, internal)
return False
job_id = job_id_without_suffix + _job_id_suffix_start + get_target(ctx_msg)
try:
_scheduler.remove_job(job_id, 'default')
_send_text('成功删除计划任务 ' + job_id_without_suffix, ctx_msg, internal)
return True
except JobLookupError:
_send_text('没有找到计划任务 ' + job_id_without_suffix, ctx_msg, internal)
return False
@cr.register('get_job', 'get-job', 'get')
@cr.restrict(full_command_only=True)
@_check_target
def get_job(args_text, ctx_msg, internal=False):
job_id_without_suffix = args_text.strip()
if not job_id_without_suffix:
_send_text('请指定计划任务的 ID', ctx_msg, internal)
return None
job_id = job_id_without_suffix + _job_id_suffix_start + get_target(ctx_msg)
job = _scheduler.get_job(job_id, 'default')
if internal:
if job:
job.id = job_id_without_suffix
return job
if not job:
core.echo('没有找到该计划任务,请指定正确的计划任务 ID', ctx_msg, internal)
return
reply = '找到计划任务如下:\n'
reply += 'ID' + job_id_without_suffix + '\n'
reply += '下次触发时间:\n%s\n' % job.next_run_time.strftime('%Y-%m-%d %H:%M')
reply += '命令:\n'
command_list = job.kwargs['command_list']
reply += convert_command_list_to_str(command_list)
_send_text(reply, ctx_msg, internal)
@cr.register('list_jobs', 'list-jobs', 'list')
@cr.restrict(full_command_only=True)
@_check_target
def list_jobs(_, ctx_msg, internal=False):
target = get_target(ctx_msg)
job_id_suffix = _job_id_suffix_start + target
jobs = list(filter(lambda j: j.id.endswith(job_id_suffix), _scheduler.get_jobs('default')))
if internal:
for job in jobs:
job.id = job.id[:-len(job_id_suffix)]
return jobs
for job in jobs:
job_id = job.id[:-len(job_id_suffix)]
command_list = job.kwargs['command_list']
reply = 'ID' + job_id + '\n'
reply += '下次触发时间:\n%s\n' % job.next_run_time.strftime('%Y-%m-%d %H:%M')
reply += '命令:\n'
reply += convert_command_list_to_str(command_list)
_send_text(reply, ctx_msg, internal)
if len(jobs):
_send_text('以上', ctx_msg, internal)
else:
_send_text('还没有添加计划任务', ctx_msg, internal)
@cr.register('execute_job', 'execute-job', 'execute', 'exec', 'trigger', 'do')
@cr.restrict(full_command_only=True, group_admin_only=True)
@_check_target
def execute_job(args_text, ctx_msg, internal=False):
job = get_job(args_text, ctx_msg, internal=True)
if not job:
core.echo('没有找到该计划任务,请指定正确的计划任务 ID', ctx_msg, internal)
return
job_id_suffix = _job_id_suffix_start + get_target(ctx_msg)
job_id = job.id[:-len(job_id_suffix)]
_call_commands(job_id, job.kwargs['command_list'], job.kwargs['ctx_msg'], internal)
def convert_command_list_to_str(command_list):
s = ''
if len(command_list) > 1:
for c in command_list:
s += c[0] + (' ' + c[1] if c[1] else '') + '\n'
s = s.rstrip('\n')
else:
s = command_list[0][0] + ' ' + command_list[0][1]
return s
def _send_text(text, ctx_msg, internal):
if not internal:
core.echo(text, ctx_msg)
def _send_fail_to_get_target_msg(ctx_msg, internal):
_send_text(
'无法获取 target可能因为不支持当前消息类型不支持微信群组消息'
'或由于延迟还没能加载到用户的固定 ID微信号',
ctx_msg,
internal
)
def _send_add_job_help_msg(ctx_msg, internal):
_send_text(
'此为高级命令!如果你不知道自己在做什么,请不要使用此命令。\n\n'
'使用方法:\n'
'/scheduler.add_job options|cron [--multi] job_id command\n'
'说明:\n'
'options 和 cron 用来表示触发参数,有且只能有其一,格式分别如下:\n'
'options\n'
' -M 分0 到 59\n'
' -H 时0 到 23\n'
' -d 日1 到 31\n'
' -m 月1 到 12\n'
' -w 星期0 到 6其中 0 表示星期一6 表示星期天\n'
' 以上选项的值的表示法和下面的 cron 模式相同\n'
'cron\n'
' 此模式和 Linux 的 crontab 文件的格式、顺序相同(除了星期是从 0 到 6一共 5 个用空格隔开的参数\n'
'\n'
'剩下三个参数见下一条',
ctx_msg,
internal
)
_send_text(
'--multi 为可选项,表示读取多条命令\n'
'job_id 为必填项,允许使用符合正则 [_\-a-zA-Z0-9] 的字符,作为计划任务的唯一标识,如果指定重复的 ID则会覆盖原先已有的\n'
'command 为必填项,从 job_id 之后第一个非空白字符开始,如果加了 --multi 选项,则每行算一条命令,否则一直到消息结束算作一整条命令(注意这里的命令不要加 / 前缀)\n'
'\n'
'例 1\n'
'以下命令将添加计划在每天晚上 10 点推送当天的知乎日报,并发送一条鼓励的消息:\n'
'/scheduler.add_job 0 22 * * * --multi zhihu-daily-job\n'
'zhihu\n'
'echo 今天又是很棒的一天哦!\n'
'例 2\n'
'以下命令将每 5 分钟发送一条提示:\n'
'/scheduler.add_job -M */5 tip-job echo 提示内容',
ctx_msg,
internal
)
def _send_add_job_trigger_args_invalid_msg(ctx_msg, internal):
_send_text(
'触发参数的格式不正确\n'
'如需帮助,请发送如下命令:\n'
'/scheduler.add_job --help',
ctx_msg,
internal
)
def _send_add_job_incomplete_args_msg(ctx_msg, internal):
_send_text(
'缺少必须的参数\n'
'如需帮助,请发送如下命令:\n'
'/scheduler.add_job --help',
ctx_msg,
internal
)