File size: 2,906 Bytes
44aada1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | import threading
from sqlalchemy import func, distinct, Column, String, UnicodeText
from telebot.plugins.sql_helper import SESSION, BASE
class BlackListFilters(BASE):
__tablename__ = "blacklist"
chat_id = Column(String(14), primary_key=True)
trigger = Column(UnicodeText, primary_key=True, nullable=False)
def __init__(self, chat_id, trigger):
self.chat_id = str(chat_id) # ensure string
self.trigger = trigger
def __repr__(self):
return "<Blacklist filter '%s' for %s>" % (self.trigger, self.chat_id)
def __eq__(self, other):
return bool(isinstance(other, BlackListFilters)
and self.chat_id == other.chat_id
and self.trigger == other.trigger)
BlackListFilters.__table__.create(checkfirst=True)
BLACKLIST_FILTER_INSERTION_LOCK = threading.RLock()
CHAT_BLACKLISTS = {}
def add_to_blacklist(chat_id, trigger):
with BLACKLIST_FILTER_INSERTION_LOCK:
blacklist_filt = BlackListFilters(str(chat_id), trigger)
SESSION.merge(blacklist_filt) # merge to avoid duplicate key issues
SESSION.commit()
CHAT_BLACKLISTS.setdefault(str(chat_id), set()).add(trigger)
def rm_from_blacklist(chat_id, trigger):
with BLACKLIST_FILTER_INSERTION_LOCK:
blacklist_filt = SESSION.query(
BlackListFilters).get((str(chat_id), trigger))
if blacklist_filt:
if trigger in CHAT_BLACKLISTS.get(
str(chat_id), set()): # sanity check
CHAT_BLACKLISTS.get(str(chat_id), set()).remove(trigger)
SESSION.delete(blacklist_filt)
SESSION.commit()
return True
SESSION.close()
return False
def get_chat_blacklist(chat_id):
return CHAT_BLACKLISTS.get(str(chat_id), set())
def num_blacklist_filters():
try:
return SESSION.query(BlackListFilters).count()
finally:
SESSION.close()
def num_blacklist_chat_filters(chat_id):
try:
return SESSION.query(BlackListFilters.chat_id).filter(
BlackListFilters.chat_id == str(chat_id)).count()
finally:
SESSION.close()
def num_blacklist_filter_chats():
try:
return SESSION.query(
func.count(
distinct(
BlackListFilters.chat_id))).scalar()
finally:
SESSION.close()
def __load_chat_blacklists():
global CHAT_BLACKLISTS
try:
chats = SESSION.query(BlackListFilters.chat_id).distinct().all()
for (chat_id,) in chats: # remove tuple by ( ,)
CHAT_BLACKLISTS[chat_id] = []
all_filters = SESSION.query(BlackListFilters).all()
for x in all_filters:
CHAT_BLACKLISTS[x.chat_id] += [x.trigger]
CHAT_BLACKLISTS = {x: set(y) for x, y in CHAT_BLACKLISTS.items()}
finally:
SESSION.close()
__load_chat_blacklists()
|