nonebot2/commands/scheduler.py

import os
import re
from functools import reduce, wraps

import pytz
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]'


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):
    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)
        except CommandPermissionError:
            core.echo('你没有权限执行计划任务 %s 中的命令 %s' % (job_id, command[0]), ctx_msg)
        except CommandScopeError as se:
            core.echo(
                '计划任务 %s 中的命令 %s 不支持 %s' % (job_id, command[0], se.msg_type),
                ctx_msg
            )


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('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)

    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 + '_' + 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)
        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 + '_' + 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 + '_' + get_target(ctx_msg)
    job = _scheduler.get_job(job_id, 'default')
    if internal:
        return job
    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']
    if len(command_list) > 1:
        reply += reduce(lambda x, y: x[0] + ' ' + x[1] + '\n' + y[0] + ' ' + y[1], command_list)
    else:
        reply += command_list[0][0] + ' ' + command_list[0][1]
    _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 = '_' + target
    jobs = list(filter(lambda j: j.id.endswith(job_id_suffix), _scheduler.get_jobs('default')))
    if internal:
        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'
        if len(command_list) > 1:
            reply += reduce(lambda x, y: x[0] + ' ' + x[1] + '\n' + y[0] + ' ' + y[1], command_list)
        else:
            reply += command_list[0][0] + ' ' + command_list[0][1]
        _send_text(reply, ctx_msg, internal)
    if len(jobs):
        _send_text('以上', ctx_msg, internal)
    else:
        _send_text('还没有添加计划任务', ctx_msg, internal)


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
    )