Add a rate limit for bots in contrib_bots
To prevent bots from accidently entering an infinite message loop, where they send messages as a reacting to their own messages, this commit adds the RateLimit class to run.py. It specifies how many messages can be sent in a given time interval. If this rate is exceeded, run.py exits with an error. Fixes #3210.
This commit is contained in:
parent
775df0f30a
commit
26fa2a5fc5
|
@ -6,6 +6,7 @@ import logging
|
||||||
import optparse
|
import optparse
|
||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
|
import time
|
||||||
|
|
||||||
our_dir = os.path.dirname(os.path.abspath(__file__))
|
our_dir = os.path.dirname(os.path.abspath(__file__))
|
||||||
|
|
||||||
|
@ -15,11 +16,27 @@ if os.path.exists(os.path.join(our_dir, '../api/zulip')):
|
||||||
|
|
||||||
from zulip import Client
|
from zulip import Client
|
||||||
|
|
||||||
|
class RateLimit(object):
|
||||||
|
def __init__(self, message_limit, interval_limit):
|
||||||
|
self.message_limit = message_limit
|
||||||
|
self.interval_limit = interval_limit
|
||||||
|
self.message_list = []
|
||||||
|
|
||||||
|
def is_legal(self):
|
||||||
|
self.message_list.append(time.time())
|
||||||
|
if len(self.message_list) > self.message_limit:
|
||||||
|
self.message_list.pop(0)
|
||||||
|
time_diff = self.message_list[-1] - self.message_list[0]
|
||||||
|
return time_diff >= self.interval_limit
|
||||||
|
else:
|
||||||
|
return True
|
||||||
|
|
||||||
class RestrictedClient(object):
|
class RestrictedClient(object):
|
||||||
def __init__(self, client):
|
def __init__(self, client):
|
||||||
# Only expose a subset of our Client's functionality
|
# Only expose a subset of our Client's functionality
|
||||||
user_profile = client.get_profile()
|
user_profile = client.get_profile()
|
||||||
self.send_message = client.send_message
|
self.rate_limit = RateLimit(20, 5)
|
||||||
|
self.client = client
|
||||||
try:
|
try:
|
||||||
self.full_name = user_profile['full_name']
|
self.full_name = user_profile['full_name']
|
||||||
self.email = user_profile['email']
|
self.email = user_profile['email']
|
||||||
|
@ -28,6 +45,15 @@ class RestrictedClient(object):
|
||||||
' up the zuliprc file correctly.')
|
' up the zuliprc file correctly.')
|
||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
|
|
||||||
|
def send_message(self, *args, **kwargs):
|
||||||
|
if self.rate_limit.is_legal():
|
||||||
|
self.client.send_message(*args, **kwargs)
|
||||||
|
else:
|
||||||
|
logging.error('-----> !*!*!*MESSAGE RATE LIMIT REACHED, EXITING*!*!*! <-----\n'
|
||||||
|
'Is your bot trapped in an infinite loop by reacting to'
|
||||||
|
' its own messages?')
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
def get_lib_module(lib_fn):
|
def get_lib_module(lib_fn):
|
||||||
lib_fn = os.path.abspath(lib_fn)
|
lib_fn = os.path.abspath(lib_fn)
|
||||||
if not os.path.dirname(lib_fn).startswith(os.path.join(our_dir, 'lib')):
|
if not os.path.dirname(lib_fn).startswith(os.path.join(our_dir, 'lib')):
|
||||||
|
|
Loading…
Reference in a new issue