feat: banning users

Merge pull request #103 from d-Rickyy-b/dev

blackjackbot/__init__.py

@@ -1,8 +1,13 @@
 # -*- coding: utf-8 -*-
+from telegram import Update
 from telegram.ext import CommandHandler, CallbackQueryHandler, MessageHandler, Filters
-from blackjackbot.errors import error_handler
 
 from blackjackbot.commands import game, admin, settings, util
+from blackjackbot.errors import error_handler
+from util import BannedUserHandler, banned_user_callback
+
+# Banned users
+banned_user_handler = BannedUserHandler(callback=banned_user_callback, type=Update)
 
 # User commands
 start_command_handler = CommandHandler("start", game.start_cmd)
@@ -17,6 +22,9 @@
 users_command_handler = CommandHandler("users", admin.users_cmd)
 answer_command_handler = CommandHandler("answer", admin.answer_comment_cmd, Filters.reply)
 kill_command_handler = CommandHandler("kill", admin.kill_game_cmd, Filters.text)
+ban_command_handler = CommandHandler("ban", admin.ban_user_cmd, pass_args=True)
+unban_command_handler = CommandHandler("unban", admin.unban_user_cmd, pass_args=True)
+bans_command_handler = CommandHandler("bans", admin.bans_cmd)
 
 # Callback handlers
 hit_callback_handler = CallbackQueryHandler(game.hit_callback, pattern=r"^hit_[0-9]{7}$")
@@ -26,8 +34,11 @@
 newgame_callback_handler = CallbackQueryHandler(game.newgame_callback, pattern=r"^newgame$")
 language_callback_handler = CallbackQueryHandler(settings.language_callback, pattern=r"^lang_([a-z]{2}(?:-[a-z]{2})?)$")
 
-handlers = [start_command_handler, stop_command_handler, join_callback_handler, hit_callback_handler, stand_callback_handler, start_callback_handler,
-            language_command_handler, stats_command_handler, newgame_callback_handler, reload_lang_command_handler, language_callback_handler,
-            users_command_handler, comment_command_handler, comment_text_command_handler, answer_command_handler]
+handlers = [banned_user_handler,
+            start_command_handler, stop_command_handler, join_callback_handler, hit_callback_handler,
+            stand_callback_handler, start_callback_handler, language_command_handler, stats_command_handler,
+            newgame_callback_handler, reload_lang_command_handler, language_callback_handler, users_command_handler,
+            comment_command_handler, comment_text_command_handler, answer_command_handler, ban_command_handler,
+            unban_command_handler, bans_command_handler]
 
 __all__ = ['handlers', 'error_handler']

blackjackbot/commands/admin/__init__.py

@@ -1,5 +1,7 @@
 # -*- coding: utf-8 -*-
 from .functions import notify_admins
-from .commands import answer_comment_cmd, reload_languages_cmd, users_cmd, kill_game_cmd
+from .commands import answer_comment_cmd, reload_languages_cmd, users_cmd, kill_game_cmd, ban_user_cmd, \
+    unban_user_cmd, bans_cmd
 
-__all__ = ["answer_comment_cmd", "reload_languages_cmd", "users_cmd", "notify_admins", "kill_game_cmd"]
+__all__ = ["answer_comment_cmd", "reload_languages_cmd", "users_cmd", "notify_admins", "kill_game_cmd", "ban_user_cmd",
+           "unban_user_cmd", "bans_cmd"]

blackjackbot/commands/admin/commands.py

@@ -3,6 +3,8 @@
 import logging
 import re
 
+from telegram import ParseMode
+
 from blackjackbot.commands.admin import notify_admins
 from blackjackbot.commands.util.decorators import admin_method
 from blackjackbot.errors import NoActiveGameException
@@ -13,6 +15,53 @@
 logger = logging.getLogger(__name__)
 
 
+@admin_method
+def ban_user_cmd(update, context):
+    """Bans a user from using the bot"""
+    usage_message = r"Please provide a valid userid\. Usage: `/ban <userid>`"
+    # Try to get user_id from command
+    if len(context.args) != 1:
+        update.effective_message.reply_text(usage_message, parse_mode=ParseMode.MARKDOWN_V2)
+        return
+
+    match = re.search(r"^\d+$", context.args[0])
+    if not match:
+        logger.error(f"The user_id did not match. Args: {context.args}")
+        update.effective_message.reply_text(usage_message, parse_mode=ParseMode.MARKDOWN_V2)
+        return
+    user_id = match.group(0)
+
+    db = Database()
+    db.ban_user(user_id=user_id)
+
+    logger.info(f"Admin '{update.effective_user.id}' banned user '{user_id}'!")
+    notify_admins(f"Admin '{update.effective_user.id}' banned user '{user_id}'!", context)
+
+
+@admin_method
+def unban_user_cmd(update, context):
+    """Unbans a user from using the bot"""
+    usage_message = r"Please provide a valid userid\. Usage: `/unban <userid>`"
+
+    # Try to get user_id from command
+    if len(context.args) != 1:
+        update.message.reply_text(usage_message, parse_mode=ParseMode.MARKDOWN_V2)
+        return
+
+    match = re.search(r"^\d+$", context.args[0])
+    if not match:
+        logger.error(f"The user_id did not match. Args: {context.args}")
+        update.effective_message.reply_text(usage_message, parse_mode=ParseMode.MARKDOWN_V2)
+        return
+    user_id = match.group(0)
+
+    db = Database()
+    db.unban_user(user_id=user_id)
+
+    logger.info(f"Admin '{update.effective_user.id}' unbanned user '{user_id}'!")
+    notify_admins(f"Admin '{update.effective_user.id}' unbanned user '{user_id}'!", context)
+
+
 @admin_method
 def kill_game_cmd(update, context):
     """Kills the game for a certain chat/group"""
@@ -103,3 +152,14 @@ def users_cmd(update, context):
     text = "Last 24 hours: {}".format(len(players))
 
     update.message.reply_text(text=text)
+
+
+@admin_method
+def bans_cmd(update, context):
+    """Returns the amount of players in the last 24 hours"""
+    db = Database()
+    banned_users = db.get_banned_users()
+
+    text = f"Banned user count: {len(banned_users)}"
+
+    update.message.reply_text(text=text)

blackjackbot/commands/game/commands.py

@@ -1,5 +1,7 @@
 # -*- coding: utf-8 -*-
 
+from telegram.parsemode import ParseMode
+
 import blackjack.errors as errors
 from blackjack.game import BlackJackGame
 from blackjackbot.commands.util import html_mention, get_game_keyboard, get_join_keyboard, get_start_keyboard, remove_inline_keyboard
@@ -156,21 +158,21 @@ def hit_callback(update, context):
         game.draw_card()
         player_cards = get_cards_string(player, lang_id)
         text = translator("your_cards_are").format(user_mention, player.cardvalue, player_cards)
-        update.effective_message.edit_text(text=text, parse_mode="HTML", reply_markup=get_game_keyboard(game.id, lang_id))
+        update.effective_message.edit_text(text=text, parse_mode=ParseMode.HTML, reply_markup=get_game_keyboard(game.id, lang_id))
     except errors.PlayerBustedException:
         player_cards = get_cards_string(player, lang_id)
         text = (translator("your_cards_are") + "\n\n" + translator("you_busted")).format(user_mention, player.cardvalue, player_cards)
-        update.effective_message.edit_text(text=text, parse_mode="HTML", reply_markup=None)
+        update.effective_message.edit_text(text=text, parse_mode=ParseMode.HTML, reply_markup=None)
         next_player(update, context)
     except errors.PlayerGot21Exception:
         player_cards = get_cards_string(player, lang_id)
         if player.has_blackjack():
             text = (translator("your_cards_are") + "\n\n" + translator("got_blackjack")).format(user_mention, player.cardvalue, player_cards)
-            update.effective_message.edit_text(text=text, parse_mode="HTML", reply_markup=None)
+            update.effective_message.edit_text(text=text, parse_mode=ParseMode.HTML, reply_markup=None)
             next_player(update, context)
         elif player.cardvalue == 21:
             text = (translator("your_cards_are") + "\n\n" + translator("got_21")).format(user_mention, player.cardvalue, player_cards)
-            update.effective_message.edit_text(text=text, parse_mode="HTML", reply_markup=None)
+            update.effective_message.edit_text(text=text, parse_mode=ParseMode.HTML, reply_markup=None)
             next_player(update, context)
 
 

blackjackbot/commands/game/functions.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 import logging
 
-from telegram import InlineKeyboardButton, InlineKeyboardMarkup
+from telegram import InlineKeyboardButton, InlineKeyboardMarkup, ParseMode
 
 from blackjack.errors import NoPlayersLeftException
 from blackjack.game import BlackJackGame
@@ -45,15 +45,15 @@ def players_turn(update, context):
     # We need reply_text here, because we must send a new message (this is the first message for the player)!
     if player.has_blackjack():
         text = (translator("your_cards_are") + "\n\n" + translator("got_blackjack")).format(user_mention, player.cardvalue, player_cards)
-        update.effective_message.reply_text(text=text, parse_mode="HTML", reply_markup=None)
+        update.effective_message.reply_text(text=text, parse_mode=ParseMode.HTML, reply_markup=None)
         next_player(update, context)
     elif player.cardvalue == 21:
         text = (translator("your_cards_are") + "\n\n" + translator("got_21")).format(user_mention, player.cardvalue, player_cards)
-        update.effective_message.reply_text(text=text, parse_mode="HTML", reply_markup=None)
+        update.effective_message.reply_text(text=text, parse_mode=ParseMode.HTML, reply_markup=None)
         next_player(update, context)
     else:
         text = translator("your_cards_are").format(user_mention, player.cardvalue, player_cards)
-        update.effective_message.reply_text(text=text, parse_mode="HTML", reply_markup=get_game_keyboard(game.id, lang_id))
+        update.effective_message.reply_text(text=text, parse_mode=ParseMode.HTML, reply_markup=get_game_keyboard(game.id, lang_id))
 
 
 @needs_active_game
@@ -76,7 +76,7 @@ def next_player(update, context):
         # TODO merge messages
         update.effective_message.reply_text(translator("dealers_cards_are").format(game.dealer.cardvalue,
                                                                                    get_cards_string(game.dealer, lang_id)),
-                                            parse_mode="HTML")
+                                            parse_mode=ParseMode.HTML)
         evaluation_string = generate_evaluation_string(game, lang_id)
 
         newgame_button = InlineKeyboardButton(text=translator("inline_keyboard_newgame"), callback_data="newgame")

blackjackbot/commands/util/commands.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 
-from telegram import ForceReply
+from telegram import ForceReply, ParseMode
 
 from blackjackbot.commands.admin.functions import notify_admins
 from blackjackbot.lang import translate
@@ -10,7 +10,7 @@
 
 
 def stats_cmd(update, context):
-    update.message.reply_text(get_user_stats(update.effective_user.id), parse_mode="HTML")
+    update.message.reply_text(get_user_stats(update.effective_user.id), parse_mode=ParseMode.HTML)
 
 
 def comment_cmd(update, context):

database/database.py

@@ -3,6 +3,7 @@
 import os
 import sqlite3
 from time import time
+
 from util import Cache
 
 
@@ -11,29 +12,35 @@ class Database(object):
 
     _instance = None
     _initialized = False
+    _banned_users = set()
 
     def __new__(cls):
         if not Database._instance:
             Database._instance = super(Database, cls).__new__(cls)
         return Database._instance
 
     def __init__(self):
-        if not self._initialized:
-            database_path = os.path.join(self.dir_path, "users.db")
-            self.logger = logging.getLogger(__name__)
+        if self._initialized:
+            return
+
+        database_path = os.path.join(self.dir_path, "users.db")
+        self.logger = logging.getLogger(__name__)
+
+        if not os.path.exists(database_path):
+            self.logger.debug("File '{}' does not exist! Trying to create one.".format(database_path))
+        try:
+            self.create_database(database_path)
+        except Exception:
+            self.logger.error("An error has occurred while creating the database!")
 
-            if not os.path.exists(database_path):
-                self.logger.debug("File '{}' does not exist! Trying to create one.".format(database_path))
-            try:
-                self.create_database(database_path)
-            except Exception:
-                self.logger.error("An error has occurred while creating the database!")
+        self.connection = sqlite3.connect(database_path)
+        self.connection.row_factory = sqlite3.Row
+        self.connection.text_factory = lambda x: str(x, 'utf-8', "ignore")
+        self.cursor = self.connection.cursor()
 
-            self.connection = sqlite3.connect(database_path)
-            self.connection.text_factory = lambda x: str(x, 'utf-8', "ignore")
-            self.cursor = self.connection.cursor()
+        self.load_banned_users()
 
-            self._initialized = True
+        self._initialized = True
 
     @staticmethod
     def create_database(database_path):
@@ -63,6 +70,7 @@ def create_database(database_path):
                        "'games_won' INTEGER DEFAULT 0,"
                        "'games_tie' INTEGER DEFAULT 0,"
                        "'last_played' INTEGER DEFAULT 0,"
+                       "'banned' INTEGER DEFAULT 0,"
                        "PRIMARY KEY('user_id'));")
 
         cursor.execute("CREATE TABLE IF NOT EXISTS 'chats'"
@@ -72,15 +80,49 @@ def create_database(database_path):
         connection.commit()
         connection.close()
 
+    def load_banned_users(self):
+        """Loads all banned users from the database into a list"""
+        self.cursor.execute("SELECT user_id FROM users WHERE banned=1;")
+        result = self.cursor.fetchall()
+
+        if not result:
+            return
+
+        for row in result:
+            print(int(row["user_id"]))
+            self._banned_users.add(int(row["user_id"]))
+
+    def get_banned_users(self):
+        """Returns a list of all banned user_ids"""
+        return self._banned_users
+
     def get_user(self, user_id):
-        self.cursor.execute("SELECT user_id, first_name, last_name, username, games_played, games_won, games_tie, last_played"
+        self.cursor.execute("SELECT user_id, first_name, last_name, username, games_played, games_won, games_tie, last_played, banned"
                             " FROM users WHERE user_id=?;", [str(user_id)])
 
         result = self.cursor.fetchone()
         if not result or len(result) == 0:
             return None
         return result
 
+    def is_user_banned(self, user_id):
+        """Checks if a user was banned by the admin of the bot from using it"""
+        # user = self.get_user(user_id)
+        # return user is not None and user[8] == 1
+        return int(user_id) in self._banned_users
+
+    def ban_user(self, user_id):
+        """Bans a user from using a the bot"""
+        self.cursor.execute("UPDATE users SET banned=1 WHERE user_id=?;", [str(user_id)])
+        self.connection.commit()
+        self._banned_users.add(int(user_id))
+
+    def unban_user(self, user_id):
+        """Unbans a user from using a the bot"""
+        self.cursor.execute("UPDATE users SET banned=0 WHERE user_id=?;", [str(user_id)])
+        self.connection.commit()
+        self._banned_users.remove(int(user_id))
+
     def get_recent_players(self):
         one_day_in_secs = 60 * 60 * 24
         current_time = int(time())
@@ -141,7 +183,7 @@ def add_user(self, user_id, lang_id, first_name, last_name, username):
 
     def _add_user(self, user_id, lang_id, first_name, last_name, username):
         try:
-            self.cursor.execute("INSERT INTO users VALUES (?, ?, ?, ?, 0, 0, 0, 0);", [str(user_id), first_name, last_name, username])
+            self.cursor.execute("INSERT INTO users VALUES (?, ?, ?, ?, 0, 0, 0, 0, 0);", [str(user_id), first_name, last_name, username])
             self.cursor.execute("INSERT INTO chats VALUES (?, ?);", [str(user_id), lang_id])
             self.connection.commit()
         except sqlite3.IntegrityError:

util/__init__.py

@@ -1,4 +1,6 @@
 # -*- coding: utf-8 -*-
+from .bannedusercallback import banned_user_callback
+from .banneduserhandler import BannedUserHandler
 from .cache import Cache
 
-__all__ = ['Cache']
+__all__ = ["Cache", "BannedUserHandler", "banned_user_callback"]

util/bannedusercallback.py

@@ -0,0 +1,10 @@
+# -*- coding: utf-8 -*-
+
+def banned_user_callback(update, context):
+    """Gets called by the dispatcher when it's found that the user sending the update was banned from using the bot"""
+    banned_text = "You have been banned from using this bot!"
+
+    if update.callback_query:
+        update.callback_query.answer(banned_text)
+    else:
+        update.effective_message.reply_text(banned_text)

util/banneduserhandler.py

@@ -0,0 +1,13 @@
+# -*- coding: utf-8 -*-
+from telegram.ext import TypeHandler
+
+import database
+
+
+class BannedUserHandler(TypeHandler):
+
+    def check_update(self, update):
+        db = database.Database()
+        if db.is_user_banned(update.effective_user.id):
+            return True
+        return False