import threading from sqlalchemy import func, distinct, Column, String, UnicodeText, Integer from emilia.modules.sql 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) class BlacklistSettings(BASE): __tablename__ = "blacklist_settings" chat_id = Column(String(14), primary_key=True) blacklist_type = Column(Integer, default=1) value = Column(UnicodeText, default="0") def __init__(self, chat_id, blacklist_type=1, value="0"): self.chat_id = str(chat_id) self.blacklist_type = blacklist_type self.value = value def __repr__(self): return "<{} will executing {} for blacklist trigger.>".format(self.chat_id, self.blacklist_type) BlackListFilters.__table__.create(checkfirst=True) BlacklistSettings.__table__.create(checkfirst=True) BLACKLIST_FILTER_INSERTION_LOCK = threading.RLock() BLACKLIST_SETTINGS_INSERTION_LOCK = threading.RLock() CHAT_BLACKLISTS = {} CHAT_SETTINGS_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() global CHAT_BLACKLISTS if CHAT_BLACKLISTS.get(str(chat_id), set()) == set(): CHAT_BLACKLISTS[str(chat_id)] = {trigger} else: CHAT_BLACKLISTS.get(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 set_blacklist_strength(chat_id, blacklist_type, value): # for blacklist_type # 0 = nothing # 1 = delete # 2 = warn # 3 = mute # 4 = kick # 5 = ban # 6 = tban # 7 = tmute with BLACKLIST_SETTINGS_INSERTION_LOCK: global CHAT_SETTINGS_BLACKLISTS curr_setting = SESSION.query(BlacklistSettings).get(str(chat_id)) if not curr_setting: curr_setting = BlacklistSettings(chat_id, blacklist_type=int(blacklist_type), value=value) curr_setting.blacklist_type = int(blacklist_type) curr_setting.value = str(value) CHAT_SETTINGS_BLACKLISTS[str(chat_id)] = {'blacklist_type': int(blacklist_type), 'value': value} SESSION.add(curr_setting) SESSION.commit() def get_blacklist_setting(chat_id): try: setting = CHAT_SETTINGS_BLACKLISTS.get(str(chat_id)) if setting: return setting['blacklist_type'], setting['value'] else: return 1, "0" 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() def __load_chat_settings_blacklists(): global CHAT_SETTINGS_BLACKLISTS try: chats_settings = SESSION.query(BlacklistSettings).all() for x in chats_settings: # remove tuple by ( ,) CHAT_SETTINGS_BLACKLISTS[x.chat_id] = {'blacklist_type': x.blacklist_type, 'value': x.value} finally: SESSION.close() def migrate_chat(old_chat_id, new_chat_id): with BLACKLIST_FILTER_INSERTION_LOCK: chat_filters = SESSION.query(BlackListFilters).filter(BlackListFilters.chat_id == str(old_chat_id)).all() for filt in chat_filters: filt.chat_id = str(new_chat_id) SESSION.commit() __load_chat_blacklists() __load_chat_settings_blacklists()