diff --git a/requirements.txt b/requirements.txt index bdab894..3c0199f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,6 +10,7 @@ howdoi pillow psutil speedtest-cli +sqlalchemy telethon Pafy youtube-dl \ No newline at end of file diff --git a/ubot/database.py b/ubot/database.py index dd2102d..7b1dddd 100644 --- a/ubot/database.py +++ b/ubot/database.py @@ -1,30 +1,30 @@ # SPDX-License-Identifier: GPL-2.0-or-later import json -import sqlite3 + +from sqlalchemy import Column, Integer, MetaData, String, Table, create_engine class Database(): - db_conn = sqlite3.connect("database.sqlite") + engine = create_engine('sqlite:///database.sqlite') + metadata = MetaData() - def __init__(self): - cur = self.db_conn.cursor() - cur.execute( - """CREATE TABLE IF NOT EXISTS chats ( - id integer PRIMARY KEY, - disabled_commands text NOT NULL - );""" - ) + chats = Table( + "chats", + metadata, + Column("id", Integer, primary_key=True), + Column("disabled_commands", String, nullable=False) + ) - def ensure_chat_table(self, chat_id: int): - cur = self.db_conn.cursor() - cur.execute("INSERT OR IGNORE INTO chats (id, disabled_commands) VALUES (?, ?);", [str(chat_id), "[]"]) - self.db_conn.commit() + metadata.create_all(engine) + conn = engine.connect() + + def ensure_disabled_commands_table(self, chat_id: int): + self.conn.execute(f"INSERT OR IGNORE INTO chats (id, disabled_commands) VALUES (?, ?);", (chat_id, "[]")) def get_disabled_commands(self, chat_id: int) -> list: - self.ensure_chat_table(chat_id) - cur = self.db_conn.cursor() - disabled_raw = cur.execute("SELECT disabled_commands FROM chats WHERE id = ?;", [str(chat_id)]).fetchone() + self.ensure_disabled_commands_table(chat_id) + disabled_raw = self.conn.execute(f"SELECT disabled_commands FROM chats WHERE id = ?;", (chat_id, )).fetchone() return json.loads(disabled_raw[0] if disabled_raw else "[]") def disable_command(self, chat_id: int, command: str): @@ -33,9 +33,7 @@ class Database(): if command not in disabled_commands: disabled_commands.append(command) new_disabled_commands = json.dumps(disabled_commands) - cur = self.db_conn.cursor() - cur.execute("UPDATE chats SET disabled_commands = ? WHERE id = ?;", [new_disabled_commands, str(chat_id)]) - self.db_conn.commit() + self.conn.execute("UPDATE chats SET disabled_commands = ? WHERE id = ?;", (new_disabled_commands, chat_id)) def enable_command(self, chat_id: int, command: str): disabled_commands = self.get_disabled_commands(chat_id) @@ -43,6 +41,4 @@ class Database(): if command in disabled_commands: disabled_commands.remove(command) new_disabled_commands = json.dumps(disabled_commands) - cur = self.db_conn.cursor() - cur.execute("UPDATE chats SET disabled_commands = ? WHERE id = ?;", [new_disabled_commands, str(chat_id)]) - self.db_conn.commit() + self.conn.execute(f"UPDATE chats SET disabled_commands = ? WHERE id = ?;", (new_disabled_commands, chat_id)) diff --git a/ubot/loader.py b/ubot/loader.py index 9f30e7f..b6a9aa1 100644 --- a/ubot/loader.py +++ b/ubot/loader.py @@ -64,18 +64,19 @@ class Loader(): def add(self, pattern: str = None, **args): pattern = args.get("pattern", pattern) pattern_extra = args.get("pattern_extra", "") + not_disableable = args.get("no_disable", False) or args.get('owner', False) or args.get('sudo', False) or args.get('admin', False) def decorator(func): if args.get("hide_help", False): if func.__module__.split(".")[-1] in self.help_hidden_dict: - self.help_hidden_dict[func.__module__.split(".")[-1]] += [[pattern, args.get('help', None)]] + self.help_hidden_dict[func.__module__.split(".")[-1]] += [[pattern, args.get('help', None), not_disableable]] else: - self.help_hidden_dict[func.__module__.split(".")[-1]] = [[pattern, args.get('help', None)]] + self.help_hidden_dict[func.__module__.split(".")[-1]] = [[pattern, args.get('help', None), not_disableable]] else: if func.__module__.split(".")[-1] in self.help_dict: - self.help_dict[func.__module__.split(".")[-1]] += [[pattern, args.get('help', None)]] + self.help_dict[func.__module__.split(".")[-1]] += [[pattern, args.get('help', None), not_disableable]] else: - self.help_dict[func.__module__.split(".")[-1]] = [[pattern, args.get('help', None)]] + self.help_dict[func.__module__.split(".")[-1]] = [[pattern, args.get('help', None), not_disableable]] self.command_handler.incoming_commands.append({ "pattern": pattern, @@ -104,19 +105,20 @@ class Loader(): def add_list(self, pattern: list = None, **args): pattern_list = args.get("pattern", pattern) pattern_extra = args.get("pattern_extra", "") + not_disableable = args.get("no_disable", False) or args.get('owner', False) or args.get('sudo', False) or args.get('admin', False) def decorator(func): for pattern in pattern_list: if args.get("hide_help", False): if func.__module__.split(".")[-1] in self.help_hidden_dict: - self.help_hidden_dict[func.__module__.split(".")[-1]] += [[pattern, args.get('help', None)]] + self.help_hidden_dict[func.__module__.split(".")[-1]] += [[pattern, args.get('help', None), not_disableable]] else: - self.help_hidden_dict[func.__module__.split(".")[-1]] = [[pattern, args.get('help', None)]] + self.help_hidden_dict[func.__module__.split(".")[-1]] = [[pattern, args.get('help', None), not_disableable]] else: if func.__module__.split(".")[-1] in self.help_dict: - self.help_dict[func.__module__.split(".")[-1]] += [[pattern, args.get('help', None)]] + self.help_dict[func.__module__.split(".")[-1]] += [[pattern, args.get('help', None), not_disableable]] else: - self.help_dict[func.__module__.split(".")[-1]] = [[pattern, args.get('help', None)]] + self.help_dict[func.__module__.split(".")[-1]] = [[pattern, args.get('help', None), not_disableable]] self.command_handler.incoming_commands.append({ "pattern": pattern, @@ -145,19 +147,20 @@ class Loader(): def add_dict(self, pattern: dict = None, **args): pattern_dict = args.get("pattern", pattern) pattern_extra = args.get("pattern_extra", "") + not_disableable = args.get("no_disable", False) or args.get('owner', False) or args.get('sudo', False) or args.get('admin', False) def decorator(func): for pattern, extra in pattern_dict.items(): if args.get("hide_help", False): if func.__module__.split(".")[-1] in self.help_hidden_dict: - self.help_hidden_dict[func.__module__.split(".")[-1]] += [[pattern, args.get('help', None)]] + self.help_hidden_dict[func.__module__.split(".")[-1]] += [[pattern, args.get('help', None), not_disableable]] else: - self.help_hidden_dict[func.__module__.split(".")[-1]] = [[pattern, args.get('help', None)]] + self.help_hidden_dict[func.__module__.split(".")[-1]] = [[pattern, args.get('help', None), not_disableable]] else: if func.__module__.split(".")[-1] in self.help_dict: - self.help_dict[func.__module__.split(".")[-1]] += [[pattern, args.get('help', None)]] + self.help_dict[func.__module__.split(".")[-1]] += [[pattern, args.get('help', None), not_disableable]] else: - self.help_dict[func.__module__.split(".")[-1]] = [[pattern, args.get('help', None)]] + self.help_dict[func.__module__.split(".")[-1]] = [[pattern, args.get('help', None), not_disableable]] self.command_handler.incoming_commands.append({ "pattern": pattern, diff --git a/ubot/modules/system.py b/ubot/modules/system.py index 9728947..6f4a8bc 100644 --- a/ubot/modules/system.py +++ b/ubot/modules/system.py @@ -79,10 +79,14 @@ async def bot_repo(event): async def disable_command(event): if event.args: for value in ldr.help_dict.values(): - for info in [i[0] for i in value]: - if event.args == info: - await event.reply(f"Disabling **{info}** in chat **{event.chat.id}**!") - ldr.db.disable_command(event.chat.id, info) + for info in value: + if event.args == info[0]: + if info[2]: + await event.reply(f"**{info[0]}** cannot be disabled!") + return + + await event.reply(f"Disabling **{info[0]}** in chat **{event.chat.id}**!") + ldr.db.disable_command(event.chat.id, info[0]) return await event.reply(f"**{event.args}** is not a command!")