Add chat blacklist and whitelist (Redone)

README.md

@@ -1,7 +1,7 @@
 # tgbot
 A modular telegram Python bot running on python3 with an sqlalchemy database.
 
-Originally a simple group management bot with multiple admin features, it has evolved, becoming extremely modular and 
+Originally a simple group management bot with multiple admin features, it has evolved, becoming extremely modular and
 simple to use.
 
 Can be found on telegram as [Marie](https://t.me/BanhammerMarie_bot).
@@ -31,11 +31,11 @@ This is because markdown parsing is done by iterating through a dict, which are
 There are two possible ways of configuring your bot: a config.py file, or ENV variables.
 
 The prefered version is to use a `config.py` file, as it makes it easier to see all your settings grouped together.
-This file should be placed in your `tg_bot` folder, alongside the `__main__.py` file . 
-This is where your bot token will be loaded from, as well as your database URI (if you're using a database), and most of 
+This file should be placed in your `tg_bot` folder, alongside the `__main__.py` file .
+This is where your bot token will be loaded from, as well as your database URI (if you're using a database), and most of
 your other settings.
 
-It is recommended to import sample_config and extend the Config class, as this will ensure your config contains all 
+It is recommended to import sample_config and extend the Config class, as this will ensure your config contains all
 defaults set in the sample_config, hence making it easier to upgrade.
 
 An example `config.py` file could be:
@@ -64,7 +64,7 @@ The following env variables are supported:
  - `OWNER_USERNAME`: Your username
 
  - `DATABASE_URL`: Your database URL
- - `MESSAGE_DUMP`: optional: a chat where your replied saved messages are stored, to stop people deleting their old 
+ - `MESSAGE_DUMP`: optional: a chat where your replied saved messages are stored, to stop people deleting their old
  - `LOAD`: Space separated list of modules you would like to load
  - `NO_LOAD`: Space separated list of modules you would like NOT to load
  - `WEBHOOK`: Setting this to ANYTHING will enable webhooks when in env mode
@@ -75,13 +75,15 @@ The following env variables are supported:
  - `SUPPORT_USERS`: A space separated list of user_ids which should be considered support users (can gban/ungban,
  nothing else)
  - `WHITELIST_USERS`: A space separated list of user_ids which should be considered whitelisted - they can't be banned.
+ - `WHITELIST_CHATS`: A space separated list of chat_ids which should be considered whitelisted - the bot will only join those chats
+ - `BLACKLIST_CHATS`: A space separated list of chat_ids which should be considered blacklisted - the bot will not join those chats. **NOTE** that if both whitelist and blacklist will be used, only the blacklisted chats will be taken into consideration.
  - `DONATION_LINK`: Optional: link where you would like to receive donations.
  - `CERT_PATH`: Path to your webhook certificate
  - `PORT`: Port to use for your webhooks
  - `DEL_CMDS`: Whether to delete commands from users which don't have rights to use that command
  - `STRICT_GBAN`: Enforce gbans across new groups as well as old groups. When a gbanned user talks, he will be banned.
  - `WORKERS`: Number of threads to use. 8 is the recommended (and default) amount, but your experience may vary.
- __Note__ that going crazy with more threads wont necessarily speed up your bot, given the large amount of sql data 
+ __Note__ that going crazy with more threads wont necessarily speed up your bot, given the large amount of sql data
  accesses, and the way python asynchronous calls work.
  - `BAN_STICKER`: Which sticker to use when banning people.
  - `ALLOW_EXCL`: Whether to allow using exclamation marks ! for commands as well as /.
@@ -166,8 +168,8 @@ commands will allow the bot to load it and add the documentation for
 your module to the `/help` command. Setting the `__mod_name__` variable will also allow you to use a nicer, user
 friendly name for a module.
 
-The `__migrate__()` function is used for migrating chats - when a chat is upgraded to a supergroup, the ID changes, so 
+The `__migrate__()` function is used for migrating chats - when a chat is upgraded to a supergroup, the ID changes, so
 it is necessary to migrate it in the db.
 
-The `__stats__()` function is for retrieving module statistics, eg number of users, number of chats. This is accessed 
+The `__stats__()` function is for retrieving module statistics, eg number of users, number of chats. This is accessed
 through the `/stats` command, which is only available to the bot owner.

tg_bot/__init__.py

@@ -43,6 +43,16 @@
     except ValueError:
         raise Exception("Your whitelisted users list does not contain valid integers.")
 
+    try:
+        WHITELIST_CHATS = set(int(x) for x in os.environ.get("WHITELIST_CHATS", "").split())
+    except ValueError:
+        raise Exception("Your whitelisted users list does not contain valid integers.")
+
+    try:
+        BLACKLIST_CHATS = set(int(x) for x in os.environ.get("BLACKLIST_CHATS", "").split())
+    except ValueError:
+        raise Exception("Your whitelisted users list does not contain valid integers.")
+
     WEBHOOK = bool(os.environ.get('WEBHOOK', False))
     URL = os.environ.get('URL', "")  # Does not contain token
     PORT = int(os.environ.get('PORT', 5000))
@@ -84,6 +94,16 @@
     except ValueError:
         raise Exception("Your whitelisted users list does not contain valid integers.")
 
+    try:
+        WHITELIST_CHATS = set(int(x) for x in Config.WHITELIST_CHATS or [])
+    except ValueError:
+        raise Exception("Your whitelisted chats list does not contain valid integers.")
+
+    try:
+        BLACKLIST_CHATS = set(int(x) for x in Config.BLACKLIST_CHATS or [])
+    except ValueError:
+        raise Exception("Your blacklisted chats list does not contain valid integers.")
+
     WEBHOOK = Config.WEBHOOK
     URL = Config.URL
     PORT = Config.PORT

tg_bot/__main__.py

@@ -11,7 +11,7 @@
 from telegram.utils.helpers import escape_markdown
 
 from tg_bot import dispatcher, updater, TOKEN, WEBHOOK, OWNER_ID, DONATION_LINK, CERT_PATH, PORT, URL, LOGGER, \
-    ALLOW_EXCL
+    ALLOW_EXCL, BLACKLIST_CHATS, WHITELIST_CHATS
 # needed to dynamically load modules
 # NOTE: Module order is not guaranteed, specify that in the config file!
 from tg_bot.modules import ALL_MODULES
@@ -415,6 +415,32 @@ def migrate_chats(bot: Bot, update: Update):
     raise DispatcherHandlerStop
 
 
+def is_chat_allowed(bot, update):
+    if BLACKLIST_CHATS:
+        chat_id = update.effective_message.chat_id
+        if chat_id in BLACKLIST_CHATS:
+            bot.send_message(chat_id=update.message.chat_id,
+                             text='Unallowed chat! Leaving...')
+            try:
+                bot.leave_chat(chat_id)
+
+            finally:
+                raise DispatcherHandlerStop
+            return
+
+    if WHITELIST_CHATS:
+        chat_id = update.effective_message.chat_id
+        if chat_id not in WHITELIST_CHATS:
+            bot.send_message(chat_id=update.message.chat_id,
+                             text='Unallowed chat! Leaving...')
+            try:
+                bot.leave_chat(chat_id)
+
+            finally:
+                raise DispatcherHandlerStop
+            return
+
+
 def main():
     test_handler = CommandHandler("test", test)
     start_handler = CommandHandler("start", start, pass_args=True)
@@ -427,6 +453,7 @@ def main():
 
     donate_handler = CommandHandler("donate", donate)
     migrate_handler = MessageHandler(Filters.status_update.migrate, migrate_chats)
+    blacklist_whitelist_handler = MessageHandler(Filters.group, is_chat_allowed)
 
     # dispatcher.add_handler(test_handler)
     dispatcher.add_handler(start_handler)
@@ -435,6 +462,7 @@ def main():
     dispatcher.add_handler(help_callback_handler)
     dispatcher.add_handler(settings_callback_handler)
     dispatcher.add_handler(migrate_handler)
+    dispatcher.add_handler(blacklist_whitelist_handler)
     dispatcher.add_handler(donate_handler)
 
     # dispatcher.add_error_handler(error_callback)