diff --git a/BotBase/database/query.py b/BotBase/database/query.py
index 6e87840..a1ef9c6 100644
--- a/BotBase/database/query.py
+++ b/BotBase/database/query.py
@@ -1,11 +1,11 @@
import sqlite3.dbapi2 as sqlite3
-from ..config import DB_GET_USERS, DB_GET_USER, DB_RELPATH, DB_SET_USER, DB_GET_IMEI, DB_SET_IMEI, DB_GET_API_DATE, \
- DB_SET_API_DATE, DB_SET_IMEI_DATA, DB_GET_IMEI_DATA
+from ..config import DB_GET_USERS, DB_GET_USER, DB_RELPATH, DB_SET_USER
import logging
import time
from types import FunctionType
import os
+
def create_database(path: str, query: str):
if os.path.exists(path):
logging.warning(f"Database file exists at {path}, running query")
@@ -24,7 +24,6 @@ def create_database(path: str, query: str):
except sqlite3.Error as query_error:
logging.info(f"An error has occurred while executing query: {query_error}")
-
def get_user(tg_id: int):
try:
database = sqlite3.connect(DB_RELPATH)
@@ -38,7 +37,7 @@ def get_user(tg_id: int):
return query.fetchone()
except sqlite3.Error as query_error:
logging.error(f"An error has occurred while executing DB_GET_USER query: {query_error}")
-
+ return query_error
def get_users():
try:
@@ -53,6 +52,7 @@ def get_users():
return query.fetchall()
except sqlite3.Error as query_error:
logging.error(f"An error has occurred while executing DB_GET_USERS query: {query_error}")
+ return query_error
def set_user(tg_id: int, uname: str):
@@ -69,96 +69,4 @@ def set_user(tg_id: int, uname: str):
return True
except sqlite3.Error as query_error:
logging.error(f"An error has occurred while executing DB_GET_USERS query: {query_error}")
-
-
-def set_imei(tg_id: int, imei: str):
- try:
- database = sqlite3.connect(DB_RELPATH)
- except sqlite3.Error as connection_error:
- logging.error(f"An error has occurred while connecting to database: {connection_error}")
- else:
- try:
- with database:
- cursor = database.cursor()
- cursor.execute(DB_SET_IMEI, (imei, tg_id))
- cursor.close()
- return True
- except sqlite3.Error as query_error:
- logging.error(f"An error has occurred while executing DB_SET_IMEI query: {query_error}")
-
-
-def get_imei(tg_id: int):
- try:
- database = sqlite3.connect(DB_RELPATH)
- except sqlite3.Error as connection_error:
- logging.error(f"An error has occurred while connecting to database: {connection_error}")
- else:
- try:
- with database:
- cursor = database.cursor()
- query = cursor.execute(DB_GET_IMEI, (tg_id,))
- return query.fetchone()
- except sqlite3.Error as query_error:
- logging.error(f"An error has occurred while executing DB_GET_IMEI query: {query_error}")
-
-
-def set_api_date(tg_id: int, timer: FunctionType = time.time):
- try:
- database = sqlite3.connect(DB_RELPATH)
- except sqlite3.Error as connection_error:
- logging.error(f"An error has occurred while connecting to database: {connection_error}")
- else:
- try:
- with database:
- cursor = database.cursor()
- cursor.execute(DB_SET_API_DATE, (int(timer()), tg_id))
- cursor.close()
- return True
- except sqlite3.Error as query_error:
- logging.error(f"An error has occurred while executing DB_SET_API_DATE query: {query_error}")
-
-
-def get_api_date(tg_id: int):
- try:
- database = sqlite3.connect(DB_RELPATH)
- except sqlite3.Error as connection_error:
- logging.error(f"An error has occurred while connecting to database: {connection_error}")
- else:
- try:
- with database:
- cursor = database.cursor()
- query = cursor.execute(DB_GET_API_DATE, (tg_id,))
- return query.fetchone()
- except sqlite3.Error as query_error:
- logging.error(f"An error has occurred while executing DB_GET_API_DATE query: {query_error}")
-
-
-def set_imei_data(imei: int, json_data: str):
- try:
- database = sqlite3.connect(DB_RELPATH)
- except sqlite3.Error as connection_error:
- logging.error(f"An error has occurred while connecting to database: {connection_error}")
- else:
- try:
- with database:
- cursor = database.cursor()
- cursor.execute(DB_SET_IMEI_DATA, (imei, json_data))
- cursor.close()
- return True
- except sqlite3.Error as query_error:
- logging.error(f"An error has occurred while executing DB_SET_IMEI_DATA query: {query_error}")
-
-
-def get_imei_data(imei: int):
- try:
- database = sqlite3.connect(DB_RELPATH)
- except sqlite3.Error as connection_error:
- logging.error(f"An error has occurred while connecting to database: {connection_error}")
- else:
- try:
- with database:
- cursor = database.cursor()
- query = cursor.execute(DB_GET_IMEI_DATA, (imei,))
- return query.fetchone()
- except sqlite3.Error as query_error:
- logging.error(f"An error has occurred while executing DB_GET_IMEI_DATA query: {query_error}")
+ return query_error
diff --git a/BotBase/modules/admin.py b/BotBase/modules/admin.py
index ef448b3..1d00f7e 100644
--- a/BotBase/modules/admin.py
+++ b/BotBase/modules/admin.py
@@ -26,7 +26,11 @@ def get_user_info(client, message):
if user:
logging.warning(f"Admin with id {message.from_user.id} sent /getuser {message.command[1]}")
_, uid, uname, date, banned = user
- text = USER_INFO.format(uid=uid, uname='@' + uname if uname != 'null' else uname, date=date, status='User' if not admin else 'Admin')
+ admin = uid in ADMINS
+ text = USER_INFO.format(uid=uid, uname='@' + uname if uname != 'null' else uname, date=date,
+ status='✅' if banned else '❌',
+ admin='❌' if not admin else '✅'),
+ )
send_message(client, message.chat.id, text)
else:
send_message(client, message.chat.id, f"{ERROR}: {ID_MISSING.format(uid=message.command[1])}")
@@ -43,9 +47,11 @@ def get_random_user(client, message):
send_message(client, message.chat.id, f"{INVALID_SYNTAX}: {NO_PARAMETERS.format(command='/getranduser')}")
else:
user = random.choice(get_users())
- rowid, uid, uname, date, admin = get_user(*user)
+ rowid, uid, uname, date, banned = get_user(*user)
+ admin = uid in ADMINS
text = USER_INFO.format(uid=uid, uname='@' + uname if uname != 'null' else uname, date=date,
- status='User' if not admin else 'Admin',
+ status='✅' if banned else '❌',
+ admin='❌' if not admin else '✅'),
)
send_message(client, message.chat.id, text)
@@ -59,9 +65,16 @@ def global_message(client, message):
count = 0
for uid in get_users():
count += 1
- if not send_message(client, *uid, msg): # Returns False if an error gets raised
+ if isinstance(send_message(client, *uid, msg), Exception):
missed += 1
send_message(client, message.chat.id, GLOBAL_MESSAGE_STATS.format(count=count, success=(count - missed), msg=msg))
else:
send_message(client, message.chat.id, f"{INVALID_SYNTAX}: Use /global message
"
"\n🍮 Note that the /global
command supports markdown and html styling")
+
+@Client.on_message(Filters.command("whisper") & ADMINS_FILTER & Filters.private & ~BANNED_USERS)
+def whisper(client, message):
+ if len(message.command) > 1:
+ msg = " ".join(message.command[1:])
+ logging.warning(f"Admin with id {message.from_user.id} sent /whisper to {message.command}
+
diff --git a/BotBase/modules/start.py b/BotBase/modules/start.py
index 4cb8dd1..8252f2b 100644
--- a/BotBase/modules/start.py
+++ b/BotBase/modules/start.py
@@ -22,8 +22,15 @@ def start_handler(client, message):
if message.from_user.id not in itertools.chain(*get_users()):
logging.warning(f"New user detected ({message.from_user.id}), adding to database")
set_user(message.from_user.id, None if not message.from_user.username else message.from_user.username)
- send_message(client, message.chat.id, GREET.format(mention=f"[{name}](tg://user?id={message.from_user.id})"),
- reply_markup=BUTTONS)
+ if GREET:
+ send_message(client,
+ message.chat.id,
+ GREET.format(mention=f"[{name}](tg://user?id={message.from_user.id})",
+ id=message.from_user.id,
+ username=message.from_user.username
+ ),
+ reply_markup=BUTTONS
+ )
@Client.on_callback_query(Filters.callback_data("info"))
diff --git a/BotBase/modules/start.py.save b/BotBase/modules/start.py.save
new file mode 100644
index 0000000..e50fca3
--- /dev/null
+++ b/BotBase/modules/start.py.save
@@ -0,0 +1,43 @@
+from pyrogram import Client, Filters, InlineKeyboardButton, InlineKeyboardMarkup
+from .antiflood import BANNED_USERS
+from ..config import GREET, BUTTONS, CREDITS
+from ..database.query import get_users, set_user
+import logging
+import itertools
+from ..methods.safe_send import send_message
+from ..methods.safe_edit import edit_message_text
+
+
+@Client.on_message(Filters.command("start") & ~BANNED_USERS & Filters.private)
+def start_handler(client, message):
+ """Simply handles the /start command sending a pre-defined greeting
+ and saving new users to the database"""
+
+ if message.from_user.first_name:
+ name = message.from_user.first_name
+ elif message.from_user.username:
+ name = message.from_user.username
+ else:
+ name = "Anonymous"
+ if message.from_user.id not in itertools.chain(*get_users()):
+ logging.warning(f"New user detected ({message.from_user.id}), adding to database")
+ set_user(message.from_user.id, None if not message.from_user.username else message.from_user.username)
+ if GREET:
+ send_message(client,
+ message.chat.id,
+ GREET.format(mention=f"[{name}](tg://user?id={message.from_user.id})",
+ id=message.from_user.id,
+ username=message.from_user.username
+ ),
+ reply_markup=BUTTONS
+ )
+
+
+@Client.on_callback_query(Filters.callback_data("info"))
+def bot_info(_, query):
+ buttons = InlineKeyboardMarkup([[InlineKeyboardButton("🔙 Back", "back_start")]])
+ edit_message_text(query, CREDITS, reply_markup=buttons)
+
+
+@Client.on_callback_query(Filters.callback_data("back_start"))
+def
diff --git a/DATABASE.md b/DATABASE.md
index 629b75f..bc6e884 100644
--- a/DATABASE.md
+++ b/DATABASE.md
@@ -30,9 +30,8 @@ the date and time the user was inserted in the database as a string
(formatted as d/m/Y H:M:S) and an integer (0 for `False` and 1 for `True`)
that represents the user's status (whether it is banned or not)
-- `get_users()` -> This acts similarly to the above `get_user`, but takes
-no parameters and returns a list of all the users in the database. The
-list contains tuples of the same structure of the ones returned by `get_user`
+- `get_users()` -> This method takes no parameter and returns a list
+of tuples. Each tuple contains a user ID as stored in the database
- `set_user()` -> Saves an ID/username pair (in this order)
to the database. The username parameter can be `None`
@@ -47,3 +46,13 @@ The API has been designed in a way that makes it easy to swap between different
database managers, so if you feel in the right mood make a PR to support a new
database and it'll be reviewed ASAP.
+
+# Adding more methods
+
+If you want to add custom methods to the API, we advise to follow the bot's convention:
+
+- Set the SQL query as a global variable whose name starts with `DB_` in `config.py`
+- Import it in the `BotBase.database.query` module
+- Create a new function that takes the required parameters whose name reflects the query name (without `DB_`)
+- Perform the query in a `with` context manager, close the cursor when you're done
+- Return `True` or the query result if the query was successful, or an `Exception` subclass if an error occurs