From 46fc925cf4332a6db5c229eb7a4a142ced760d05 Mon Sep 17 00:00:00 2001
From: Richard Chien <richardchienthebest@gmail.com>
Date: Sun, 14 Oct 2018 22:53:23 +0800
Subject: [PATCH] Add session timeout and privileged command
---
none/command.py | 64 ++++++++++++++++++++++++++++++++++++-------------
1 file changed, 48 insertions(+), 16 deletions(-)
diff --git a/none/command.py b/none/command.py
index 4681321f..5dada7d8 100644
--- a/none/command.py
+++ b/none/command.py
@@ -5,10 +5,9 @@ from typing import (
Tuple, Union, Callable, Iterable, Dict, Any, Optional, Sequence
)
-from aiocqhttp.message import Message
-
from . import NoneBot, permission as perm
from .log import logger
+from .message import Message
from .expression import render
from .helpers import context_id, send_expr
from .session import BaseSession
@@ -28,14 +27,15 @@ _sessions = {}
class Command:
__slots__ = ('name', 'func', 'permission',
- 'only_to_me', 'args_parser_func')
+ 'only_to_me', 'privileged', 'args_parser_func')
def __init__(self, *, name: Tuple[str], func: Callable, permission: int,
- only_to_me: bool):
+ only_to_me: bool, privileged: bool):
self.name = name
self.func = func
self.permission = permission
self.only_to_me = only_to_me
+ self.privileged = privileged
self.args_parser_func = None
async def run(self, session, *,
@@ -74,7 +74,8 @@ class Command:
def on_command(name: Union[str, Tuple[str]], *,
aliases: Iterable = (),
permission: int = perm.EVERYBODY,
- only_to_me: bool = True) -> Callable:
+ only_to_me: bool = True,
+ privileged: bool = False) -> Callable:
"""
Decorator to register a function as a command.
@@ -82,6 +83,7 @@ def on_command(name: Union[str, Tuple[str]], *,
:param aliases: aliases of command name, for convenient access
:param permission: permission required by the command
:param only_to_me: only handle messages to me
+ :param privileged: can be run even when there is already a session
"""
def deco(func: Callable) -> Callable:
@@ -96,7 +98,7 @@ def on_command(name: Union[str, Tuple[str]], *,
current_parent[parent_key] = current_parent.get(parent_key) or {}
current_parent = current_parent[parent_key]
cmd = Command(name=cmd_name, func=func, permission=permission,
- only_to_me=only_to_me)
+ only_to_me=only_to_me, privileged=privileged)
current_parent[cmd_name[-1]] = cmd
for alias in aliases:
_aliases[alias] = cmd_name
@@ -400,16 +402,25 @@ async def handle_command(bot: NoneBot, ctx: Dict[str, Any]) -> bool:
:param ctx: message context
:return: the message is handled as a command
"""
+ cmd, current_arg = parse_command(bot, str(ctx['message']).lstrip())
+ is_privileged_cmd = cmd and cmd.privileged
+ disable_interaction = is_privileged_cmd
+
+ if is_privileged_cmd:
+ logger.debug(f'Command {cmd.name} is a privileged command')
+
ctx_id = context_id(ctx)
- # wait for 1.5 seconds (at most) if the current session is running
- retry = 5
- while retry > 0 and _sessions.get(ctx_id) and _sessions[ctx_id].running:
- retry -= 1
- await asyncio.sleep(0.3)
+ if not is_privileged_cmd:
+ # wait for 1.5 seconds (at most) if the current session is running
+ retry = 5
+ while retry > 0 and \
+ _sessions.get(ctx_id) and _sessions[ctx_id].running:
+ retry -= 1
+ await asyncio.sleep(0.3)
check_perm = True
- session = _sessions.get(ctx_id)
+ session = _sessions.get(ctx_id) if not is_privileged_cmd else None
if session:
if session.running:
logger.warning(f'There is a session of command '
@@ -432,7 +443,6 @@ async def handle_command(bot: NoneBot, ctx: Dict[str, Any]) -> bool:
session = None
if not session:
- cmd, current_arg = parse_command(bot, str(ctx['message']).lstrip())
if not cmd:
logger.debug('Not a known command, ignored')
return False
@@ -442,7 +452,8 @@ async def handle_command(bot: NoneBot, ctx: Dict[str, Any]) -> bool:
session = CommandSession(bot, ctx, cmd, current_arg=current_arg)
logger.debug(f'New session of command {session.cmd.name} created')
- return await _real_run_command(session, ctx_id, check_perm=check_perm)
+ return await _real_run_command(session, ctx_id, check_perm=check_perm,
+ disable_interaction=disable_interaction)
async def call_command(bot: NoneBot, ctx: Dict[str, Any],
@@ -490,8 +501,16 @@ async def _real_run_command(session: CommandSession,
try:
logger.debug(f'Running command {session.cmd.name}')
session.running = True
- res = await session.cmd.run(session, **kwargs)
- raise _FinishException(res)
+ future = asyncio.ensure_future(session.cmd.run(session, **kwargs))
+ timeout = None
+ if session.bot.config.SESSION_RUN_TIMEOUT:
+ timeout = session.bot.config.SESSION_RUN_TIMEOUT.total_seconds()
+ try:
+ await asyncio.wait_for(future, timeout)
+ raise _FinishException(future.result())
+ except asyncio.TimeoutError:
+ # if timeout happens, we think the session is finished
+ raise _FinishException(True)
except _FurtherInteractionNeeded:
session.running = False
if disable_interaction:
@@ -522,3 +541,16 @@ async def _real_run_command(session: CommandSession,
logger.debug(f'Session of command {session.cmd.name} switching, '
f'new context message: {e.new_ctx_message}')
raise e # this is intended to be propagated to handle_message()
+
+
+def kill_current_session(bot: NoneBot, ctx: Dict[str, Any]) -> None:
+ """
+ Force kill current session of the given context,
+ despite whether it is running or not.
+
+ :param bot: NoneBot instance
+ :param ctx: message context
+ """
+ ctx_id = context_id(ctx)
+ if ctx_id in _sessions:
+ del _sessions[ctx_id]