From a39eb3483f4f669d1f230e85611cf621e42182c9 Mon Sep 17 00:00:00 2001
From: almost_random_username <gerogerigegege@localdomain>
Date: Sat, 19 Oct 2024 22:18:05 +0300
Subject: [PATCH] Allow to set up custom ban duration via
 CAPTCHABOT_BAN_DURATION option

---
 .gitignore              | 1 +
 src/constants.py        | 6 ++++++
 src/join_captcha_bot.py | 9 ++++++++-
 src/settings.py         | 6 +++++-
 4 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/.gitignore b/.gitignore
index 144ecff6..39b19c34 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,6 +1,7 @@
 **/.vscode/
 **/__pycache__/
 **/data/
+**/.idea/
 
 **/*.pyc
 **/*.pkl
diff --git a/src/constants.py b/src/constants.py
index 8f73da4c..8068aaae 100644
--- a/src/constants.py
+++ b/src/constants.py
@@ -220,6 +220,12 @@
         int(os_getenv("CAPTCHABOT_MAX_FAIL_BAN_POLL",
                       SETTINGS["CAPTCHABOT_MAX_FAIL_BAN_POLL"])),
 
+    # Duration of ban (in seconds, negative values mean indefinite ban).
+    # Useful if you want to ban someone temporarily.
+    "BAN_DURATION":
+        int(os_getenv("CAPTCHABOT_BAN_DURATION",
+                      SETTINGS["CAPTCHABOT_BAN_DURATION"])),
+
     # Last session restorable RAM data backup file path
     "F_SESSION": SCRIPT_PATH + "/session.pkl",
 
diff --git a/src/join_captcha_bot.py b/src/join_captcha_bot.py
index ef489567..9098f60c 100644
--- a/src/join_captcha_bot.py
+++ b/src/join_captcha_bot.py
@@ -32,6 +32,9 @@
 # Collections Data Types Library
 from collections import OrderedDict
 
+# Date and Time Library
+from datetime import datetime, timedelta, timezone
+
 # JSON Library
 from json import dumps as json_dumps
 
@@ -873,7 +876,11 @@ async def captcha_fail_member_kick(bot, chat_id, user_id, user_name):
         logger.info("[%s] Captcha Fail - Ban - %s (%s)",
                     chat_id, user_name, user_id)
         # Try to ban the user and notify Admins
-        ban_result = await tlg_ban_user(bot, chat_id, user_id)
+        if CONST["BAN_DURATION"] >= 0:
+            ban_until_date = datetime.now(timezone.utc) + timedelta(seconds=CONST["BAN_DURATION"])
+        else:
+            ban_until_date = None
+        ban_result = await tlg_ban_user(bot, chat_id, user_id, until_date=ban_until_date)
         if ban_result["error"] == "":
             # Ban success
             banned = True
diff --git a/src/settings.py b/src/settings.py
index c1e4b7fe..e46cdec5 100644
--- a/src/settings.py
+++ b/src/settings.py
@@ -168,5 +168,9 @@
     # Maximum number of times a user fail to solve a Poll captcha.
     # If a user don't solve the captcha after this, it will be ban
     # instead kick
-    "CAPTCHABOT_MAX_FAIL_BAN_POLL": 3
+    "CAPTCHABOT_MAX_FAIL_BAN_POLL": 3,
+
+    # Duration of ban (in seconds, negative values mean indefinite ban).
+    # Useful if you want to ban someone temporarily.
+    "CAPTCHABOT_BAN_DURATION": -1,
 }