txlyre 1 month ago
parent
commit
9dd4b5e68e
2 changed files with 39 additions and 23 deletions
  1. 1 0
      .gitignore
  2. 38 23
      main.py

+ 1 - 0
.gitignore

@@ -1,4 +1,5 @@
 __pycache__
 .venv
 *.session
+*.session-journal
 .env

+ 38 - 23
main.py

@@ -47,15 +47,13 @@ MAX_CAPTCHA_ATTEMPTS = int(os.getenv("MAX_CAPTCHA_ATTEMPTS", 8))
 USE_LOLS_API = os.getenv("USE_LOLS_API", "yes") == "yes"
 
 bot = TelegramClient("captcha_bot", API_ID, API_HASH).start(bot_token=BOT_TOKEN)
-pending = {}
-
+pending = {peer_id: {} for peer_id in ALLOWED_GROUPS}
+whitelist = {}
 
 async def lols_check(user_id):
     try:
         async with aiohttp.ClientSession() as session:
-            async with session.get(
-                f"https://api.lols.bot/account?id={user_id}"
-            ) as response:
+            async with session.get(f"https://api.lols.bot/account?id={user_id}") as response:
                 data = await response.json()
 
                 banned = data.get("banned", False)
@@ -67,7 +65,6 @@ async def lols_check(user_id):
 
     return True
 
-
 class PendingUser:
     def __init__(self, peer_id, user_id):
         self.peer_id = peer_id
@@ -132,7 +129,7 @@ class PendingUser:
         captcha = captcha.replace(" ", "").strip().upper()
 
         if captcha == self.captcha:
-            del pending[self.user_id]
+            del pending[self.peer_id][self.user_id]
 
             try:
                 await bot.delete_messages(self.peer_id, self.message_id)
@@ -156,7 +153,7 @@ class PendingUser:
         return False
 
     async def remove(self):
-        del pending[self.user_id]
+        del pending[self.peer_id][self.user_id]
 
         if self.message_id is not None:
             try:
@@ -165,7 +162,7 @@ class PendingUser:
                 pass
 
     async def kick(self, forever=False):
-        if self.user_id not in pending:
+        if self.user_id not in pending[self.peer_id]:
             return
 
         await self.remove()
@@ -176,9 +173,7 @@ class PendingUser:
                     self.peer_id,
                     self.user_id,
                     ChatBannedRights(
-                        until_date=None
-                        if forever
-                        else datetime.now() + timedelta(minutes=5),
+                        until_date=None if forever else datetime.now() + timedelta(minutes=5),
                         view_messages=True,
                     ),
                 )
@@ -186,23 +181,27 @@ class PendingUser:
         except Exception:
             pass
 
-
 @bot.on(events.ChatAction)
 async def handler(event):
     if event.chat_id not in ALLOWED_GROUPS:
         return
 
-    if (event.user_left or event.user_kicked) and event.user_id in pending:
-        await pending[event.user_id].remove()
+    peer_id = event.chat_id
+
+    if (event.user_left or event.user_kicked) and event.user_id in pending[peer_id]:
+        await pending[peer_id][event.user_id].remove()
 
         return
 
     if not event.user_joined:
         return
 
-    pending[event.user_id] = PendingUser(peer_id=event.chat_id, user_id=event.user_id)
+    if event.user_id in whitelist:
+        return
+
+    pending[peer_id][event.user_id] = PendingUser(peer_id=event.chat_id, user_id=event.user_id)
 
-    await pending[event.user_id].start(event)
+    await pending[peer_id][event.user_id].start(event)
 
 
 @bot.on(events.NewMessage)
@@ -214,7 +213,7 @@ async def handler(event):
     if not event.sender:
         return
 
-    pending_user = pending.get(event.sender.id)
+    pending_user = pending[peer_id].get(event.sender.id)
     if not pending_user:
         return
 
@@ -228,21 +227,37 @@ async def handler(event):
 
     text = getattr(event, "text", "")
 
-    await pending_user.check_captcha(text)
+    if await pending_user.check_captcha(text):
+        whitelist[event.sender.id] = time.time()
+
+        for peer_id in pending:
+            if event.sender.id in pending[peer_id]:
+                await pending[peer_id][event.sender.id].remove()
 
 
 async def check_pending():
     while True:
-        for user_id in dict(pending):
-            pending_user = pending[user_id]
+        for peer_id in dict(pending):
+            for user_id in pending[peer_id]:
+                pending_user = pending[peer_id][user_id]
 
-            if pending_user.expired:
-                await pending_user.kick()
+                if pending_user.expired:
+                    await pending_user.kick()
 
         await asyncio.sleep(1)
 
+async def cleanup_whitelist():
+    while True:
+        for user_id in dict(whitelist):
+            ts = whitelist[user_id]
+
+            if time.time() - ts >= 30*60:
+                del whitelist[user_id]
+
+        await asyncio.sleep(60)
 
 if __name__ == "__main__":
     bot.loop.create_task(check_pending())
+    bot.loop.create_task(cleanup_whitelist())
     bot.start()
     bot.run_until_disconnected()