From 88a8b4e0d3ffe11bd18046a25945b63dc729afdc Mon Sep 17 00:00:00 2001 From: Dan Elkouby Date: Tue, 26 Feb 2019 22:27:01 +0000 Subject: [PATCH] Add a timeout for sed --- stdplugins/sed.py | 3 +++ uniborg/util.py | 19 +++++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/stdplugins/sed.py b/stdplugins/sed.py index 04f7320..aeaa8e5 100644 --- a/stdplugins/sed.py +++ b/stdplugins/sed.py @@ -5,6 +5,8 @@ import regex from telethon import events, utils from telethon.tl import types, functions +from uniborg import util + HEADER = "「sed」\n" KNOWN_RE_BOTS = re.compile( r'(regex|moku|BananaButler_|rgx|l4mR)bot', @@ -17,6 +19,7 @@ KNOWN_RE_BOTS = re.compile( last_msgs = defaultdict(lambda: deque(maxlen=10)) +@util.sync_timeout(1) def doit(chat_id, match, original): fr = match.group(1) to = match.group(2) diff --git a/uniborg/util.py b/uniborg/util.py index 39cfc8c..1b93b26 100644 --- a/uniborg/util.py +++ b/uniborg/util.py @@ -2,7 +2,9 @@ # License, v. 2.0. If a copy of the MPL was not distributed with this # file, You can obtain one at http://mozilla.org/MPL/2.0/. +import functools import re +import signal from telethon import events from telethon.tl.functions.messages import GetPeerDialogsRequest @@ -28,3 +30,20 @@ async def is_read(borg, entity, message, is_out=None): dialog = (await borg(GetPeerDialogsRequest([entity]))).dialogs[0] max_id = dialog.read_outbox_max_id if is_out else dialog.read_inbox_max_id return message_id <= max_id + +def _handle_timeout(signum, frame): + raise TimeoutError("Execution took too long") + +def sync_timeout(seconds): + def decorator(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + signal.signal(signal.SIGALRM, _handle_timeout) + signal.setitimer(signal.ITIMER_REAL, seconds) + try: + r = func(*args, **kwargs) + finally: + signal.alarm(0) + return r + return wrapper + return decorator