Add files via upload

repthon/sql_helper/__init__.py

@@ -0,0 +1,34 @@
+import os
+
+from sqlalchemy import create_engine
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import scoped_session, sessionmaker
+
+# the secret configuration specific things
+from ..Config import Config
+from ..core.logger import logging
+
+LOGS = logging.getLogger(__name__)
+
+
+def start() -> scoped_session:
+    database_url = (
+        Config.DB_URI.replace("postgres:", "postgresql:")
+        if "postgres://" in Config.DB_URI
+        else Config.DB_URI
+    )
+    engine = create_engine(database_url)
+    BASE.metadata.bind = engine
+    BASE.metadata.create_all(engine)
+    return scoped_session(sessionmaker(bind=engine, autoflush=False))
+
+
+try:
+    BASE = declarative_base()
+    SESSION = start()
+except AttributeError as e:
+    # this is a dirty way for the work-around required for #23
+    LOGS.error(
+        "DB_URI is not configured. Features depending on the database might have issues."
+    )
+    LOGS.error(str(e))

repthon/sql_helper/antiflood_sql.py

@@ -0,0 +1,98 @@
+import threading
+
+from sqlalchemy import Column, Integer, String
+
+from . import BASE, SESSION
+
+DEF_COUNT = 0
+DEF_LIMIT = 0
+DEF_OBJ = (None, DEF_COUNT, DEF_LIMIT)
+
+
+class FloodControl(BASE):
+    __tablename__ = "antiflood"
+    chat_id = Column(String(14), primary_key=True)
+    user_id = Column(Integer)
+    count = Column(Integer, default=DEF_COUNT)
+    limit = Column(Integer, default=DEF_LIMIT)
+
+    def __init__(self, chat_id):
+        self.chat_id = str(chat_id)  # ensure string
+
+    def __repr__(self):
+        return "<flood control for %s>" % self.chat_id
+
+
+FloodControl.__table__.create(bind=SESSION.get_bind(), checkfirst=True)
+
+INSERTION_LOCK = threading.RLock()
+
+
+class ANTIFLOOD_SQL:
+    def __init__(self):
+        self.CHAT_FLOOD = {}
+
+
+ANTIFLOOD_SQL_ = ANTIFLOOD_SQL()
+
+
+def set_flood(chat_id, amount):
+    with INSERTION_LOCK:
+        flood = SESSION.query(FloodControl).get(str(chat_id))
+        if not flood:
+            flood = FloodControl(str(chat_id))
+
+        flood.user_id = None
+        flood.limit = amount
+
+        ANTIFLOOD_SQL_.CHAT_FLOOD[str(chat_id)] = (None, DEF_COUNT, amount)
+
+        SESSION.add(flood)
+        SESSION.commit()
+
+
+def update_flood(chat_id: str, user_id) -> bool:
+    if str(chat_id) not in ANTIFLOOD_SQL_.CHAT_FLOOD:
+        return
+    curr_user_id, count, limit = ANTIFLOOD_SQL_.CHAT_FLOOD.get(str(chat_id), DEF_OBJ)
+    if limit == 0:  # no antiflood
+        return False
+    if user_id != curr_user_id or user_id is None:  # other user
+        ANTIFLOOD_SQL_.CHAT_FLOOD[str(chat_id)] = (user_id, DEF_COUNT + 1, limit)
+        return False
+
+    count += 1
+    if count > limit:  # too many msgs, kick
+        ANTIFLOOD_SQL_.CHAT_FLOOD[str(chat_id)] = (None, DEF_COUNT, limit)
+        return True
+
+    # default -> update
+    ANTIFLOOD_SQL_.CHAT_FLOOD[str(chat_id)] = (user_id, count, limit)
+    return False
+
+
+def get_flood_limit(chat_id):
+    return ANTIFLOOD_SQL_.CHAT_FLOOD.get(str(chat_id), DEF_OBJ)[2]
+
+
+def migrate_chat(old_chat_id, new_chat_id):
+    with INSERTION_LOCK:
+        if flood := SESSION.query(FloodControl).get(str(old_chat_id)):
+            ANTIFLOOD_SQL_.CHAT_FLOOD[str(new_chat_id)] = ANTIFLOOD_SQL_.CHAT_FLOOD.get(
+                str(old_chat_id), DEF_OBJ
+            )
+            flood.chat_id = str(new_chat_id)
+            SESSION.commit()
+
+        SESSION.close()
+
+
+def __load_flood_settings():
+    try:
+        all_chats = SESSION.query(FloodControl).all()
+        ANTIFLOOD_SQL_.CHAT_FLOOD = {
+            chat.chat_id: (None, DEF_COUNT, chat.limit) for chat in all_chats
+        }
+    finally:
+        SESSION.close()
+    return ANTIFLOOD_SQL_.CHAT_FLOOD

repthon/sql_helper/autopost_sql.py

@@ -0,0 +1,84 @@
+import threading
+from sqlalchemy import Column, String, UnicodeText, distinct, func
+from . import BASE, SESSION
+
+
+class Post(BASE):
+    __tablename__ = "post"
+    target_chat_id = Column(String(14), primary_key=True)
+    to_post_chat_id = Column(String(14), primary_key=True, nullable=False)
+
+    def __init__(self, target_chat_id, to_post_chat_id):
+        self.to_post_chat_id = str(to_post_chat_id)
+        self.target_chat_id = str(target_chat_id)
+
+    def __repr__(self):
+        return "<Auto post filter '%s' for %s>" % (self.target_chat_id, self.to_post_chat_id)
+
+    def __eq__(self, other):
+        return bool(
+            isinstance(other, Post)
+            and self.target_chat_id == other.target_chat_id
+            and self.to_post_chat_id == other.to_post_chat_id
+        )
+
+
+
+
+Post.__table__.create(bind=SESSION.get_bind(), checkfirst=True)
+
+POST_FILTER_INSERTION_LOCK = threading.RLock()
+
+CHAT_POSTS = {}
+
+def add_post(target_chat_id: str, to_post_chat_id: str):
+    with POST_FILTER_INSERTION_LOCK:
+        blacklist_filt = Post(str(target_chat_id), str(to_post_chat_id))
+
+        SESSION.merge(blacklist_filt)  
+        SESSION.commit()
+        CHAT_POSTS.setdefault(str(target_chat_id), set()).add(str(to_post_chat_id))
+
+
+def get_all_post(target_chat_id: str):
+    return CHAT_POSTS.get(str(target_chat_id), set())
+
+
+def is_post(target_chat_id, to_post_chat_id):
+    with POST_FILTER_INSERTION_LOCK:
+        broadcast_group = SESSION.query(Post).get((str(target_chat_id), str(to_post_chat_id)))
+        return bool(broadcast_group)
+
+
+def remove_post(target_chat_id, to_post_chat_id):
+    with POST_FILTER_INSERTION_LOCK:
+        blacklist_filt = SESSION.query(Post).get((str(target_chat_id), str(to_post_chat_id)))
+        if blacklist_filt:
+            if str(to_post_chat_id) in CHAT_POSTS.get(str(target_chat_id), set()):  # sanity check
+                CHAT_POSTS.get(str(target_chat_id), set()).remove(str(to_post_chat_id))
+
+            SESSION.delete(blacklist_filt)
+            SESSION.commit()
+            return True
+
+        SESSION.close()
+        return False
+
+def __load_chat_channels():
+    global CHAT_POSTS
+    try:
+        chats = SESSION.query(Post.target_chat_id).distinct().all()
+        for (target_chat_id,) in chats: 
+            CHAT_POSTS[target_chat_id] = []
+
+        all_filters = SESSION.query(Post).all()
+        for x in all_filters:
+            CHAT_POSTS[x.target_chat_id] += [x.to_post_chat_id]
+
+        CHAT_POSTS = {x: set(y) for x, y in CHAT_POSTS.items()}
+
+    finally:
+        SESSION.close()
+
+
+__load_chat_channels()

repthon/sql_helper/blacklist_sql.py

@@ -0,0 +1,115 @@
+import threading
+
+from sqlalchemy import Column, String, UnicodeText, distinct, func
+
+from . import BASE, SESSION
+
+
+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(bind=SESSION.get_bind(), checkfirst=True)
+
+BLACKLIST_FILTER_INSERTION_LOCK = threading.RLock()
+
+
+class BLACKLIST_SQL:
+    def __init__(self):
+        self.CHAT_BLACKLISTS = {}
+
+
+BLACKLIST_SQL_ = BLACKLIST_SQL()
+
+
+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()
+        BLACKLIST_SQL_.CHAT_BLACKLISTS.setdefault(str(chat_id), set()).add(trigger)
+
+
+def rm_from_blacklist(chat_id, trigger):
+    with BLACKLIST_FILTER_INSERTION_LOCK:
+        if blacklist_filt := SESSION.query(BlackListFilters).get(
+            (str(chat_id), trigger)
+        ):
+            if trigger in BLACKLIST_SQL_.CHAT_BLACKLISTS.get(
+                str(chat_id), set()
+            ):  # sanity check
+                BLACKLIST_SQL_.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 BLACKLIST_SQL_.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():
+    try:
+        chats = SESSION.query(BlackListFilters.chat_id).distinct().all()
+        for (chat_id,) in chats:  # remove tuple by ( ,)
+            BLACKLIST_SQL_.CHAT_BLACKLISTS[chat_id] = []
+
+        all_filters = SESSION.query(BlackListFilters).all()
+        for x in all_filters:
+            BLACKLIST_SQL_.CHAT_BLACKLISTS[x.chat_id] += [x.trigger]
+
+        BLACKLIST_SQL_.CHAT_BLACKLISTS = {
+            x: set(y) for x, y in BLACKLIST_SQL_.CHAT_BLACKLISTS.items()
+        }
+
+    finally:
+        SESSION.close()
+
+
+__load_chat_blacklists()