Update __init__.py

tg_bot/modules/helper_funcs/telethn/__init__.py

@@ -1,14 +1,92 @@
-from tg_bot import (
-    telethn,
-    SUDO_USERS,
-    WHITELIST_USERS,
-    SUPPORT_USERS,
-    SARDEGNA_USERS,
-    DEV_USERS,
-)
+import threading
 
-HIGHER_AUTH = SUDO_USERS + DEV_USERS
+from sqlalchemy import BigInteger, Boolean, Column, UnicodeText
 
-HIGHER_AUTH = list(SUDO_USERS) + list(DEV_USERS)
+from tg_bot.modules.sql import BASE, SESSION
 
-HIGHER_AUTH.append(1087968824)
+
+class AFK(BASE):
+    __tablename__ = "afk_users"
+
+    user_id = Column(BigInteger, primary_key=True)
+    is_afk = Column(Boolean)
+    reason = Column(UnicodeText)
+
+    def __init__(self, user_id, reason="", is_afk=True):
+        self.user_id = user_id
+        self.reason = reason
+        self.is_afk = is_afk
+
+    def __repr__(self):
+        return "afk_status for {}".format(self.user_id)
+
+
+AFK.__table__.create(checkfirst=True)
+INSERTION_LOCK = threading.RLock()
+
+AFK_USERS = {}
+
+
+def is_afk(user_id):
+    return user_id in AFK_USERS
+
+
+def check_afk_status(user_id):
+    try:
+        return SESSION.query(AFK).get(user_id)
+    finally:
+        SESSION.close()
+
+
+def set_afk(user_id, reason=""):
+    with INSERTION_LOCK:
+        curr = SESSION.query(AFK).get(user_id)
+        if not curr:
+            curr = AFK(user_id, reason, True)
+        else:
+            curr.is_afk = True
+
+        AFK_USERS[user_id] = reason
+
+        SESSION.add(curr)
+        SESSION.commit()
+
+
+def rm_afk(user_id):
+    with INSERTION_LOCK:
+        curr = SESSION.query(AFK).get(user_id)
+        if curr:
+            if user_id in AFK_USERS:  # sanity check
+                del AFK_USERS[user_id]
+
+            SESSION.delete(curr)
+            SESSION.commit()
+            return True
+
+        SESSION.close()
+        return False
+
+
+def toggle_afk(user_id, reason=""):
+    with INSERTION_LOCK:
+        curr = SESSION.query(AFK).get(user_id)
+        if not curr:
+            curr = AFK(user_id, reason, True)
+        elif curr.is_afk:
+            curr.is_afk = False
+        elif not curr.is_afk:
+            curr.is_afk = True
+        SESSION.add(curr)
+        SESSION.commit()
+
+
+def __load_afk_users():
+    global AFK_USERS
+    try:
+        all_afk = SESSION.query(AFK).all()
+        AFK_USERS = {user.user_id: user.reason for user in all_afk if user.is_afk}
+    finally:
+        SESSION.close()
+
+
+__load_afk_users()