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