houseassassin/Solo_bot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update Middleware

Browse Source
This commit is contained in:
Zakhar Izmaylov
2024-11-17 23:42:09 +03:00
parent 853840475c
commit 39db3e2e64
8 changed files with 149 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.flake8
+1 -1
View File
@@ -1,6 +1,6 @@
[flake8]
max-line-length = 250
ignore = E203, E266, E501, W503, F541, E704, W293, W291, E126, E121, E123, E128, E302, E131, E231, W292, E402, E261, E305
ignore = E203, E266, E501, W503, F541, E704, W293, W291, E126, E121, E123, E128, E302, E131, E231, W292, E402, E261, E305, E701
max-complexity = 25
select = B, C, E, F, W, T4, B9
exclude = .venv,.git,.tox,dist,doc,*lib/python*,*egg,build,.txt
backup.py
+26 -10
View File
@@ -1,6 +1,7 @@
import os
import subprocess
from datetime import datetime
from typing import Union
from aiogram.types import BufferedInputFile
@@ -11,6 +12,15 @@ from logger import logger
async def backup_database():
from bot import bot
try:
if backup_file_path := _create_database_backup():
await _send_backup_to_admin(bot, backup_file_path)
_cleanup_old_backups()
except Exception as e:
logger.error(f"Ошибка при создании или отправке бэкапа: {e}")
def _create_database_backup():
USER = DB_USER
HOST = "localhost"
BACKUP_DIR = BACK_DIR
@@ -25,29 +35,37 @@ async def backup_database():
check=True,
)
logger.info(f"Бэкап базы данных создан: {BACKUP_FILE}")
return BACKUP_FILE
except subprocess.CalledProcessError as e:
logger.error(f"Ошибка при создании бэкапа базы данных: {e}")
return
return None
finally:
del os.environ["PGPASSWORD"]
async def _send_backup_to_admin(bot, backup_file_path):
try:
with open(BACKUP_FILE, "rb") as backup_file:
with open(backup_file_path, "rb") as backup_file:
backup_input_file = BufferedInputFile(
backup_file.read(), filename=os.path.basename(BACKUP_FILE)
backup_file.read(), filename=os.path.basename(backup_file_path)
)
if isinstance(ADMIN_ID, list):
for id in ADMIN_ID:
admin_ids: Union[int, list[int]] = ADMIN_ID
if isinstance(admin_ids, list):
for id in admin_ids:
await bot.send_document(id, backup_input_file)
elif isinstance(ADMIN_ID, int):
await bot.send_document(ADMIN_ID, backup_input_file)
else:
await bot.send_document(admin_ids, backup_input_file)
logger.info(f"Бэкап базы данных отправлен админу: {ADMIN_ID}")
except Exception as e:
logger.error(f"Ошибка при отправке бэкапа в Telegram: {e}")
def _cleanup_old_backups():
try:
subprocess.run(
[
"find",
BACKUP_DIR,
BACK_DIR,
"-type",
"f",
"-name",
@@ -64,5 +82,3 @@ async def backup_database():
logger.info("Старые бэкапы удалены.")
except subprocess.CalledProcessError as e:
logger.error(f"Ошибка при удалении старых бэкапов: {e}")
del os.environ["PGPASSWORD"]
filters/admin.py
+11 -5
View File
@@ -1,3 +1,5 @@
from typing import Union
from aiogram.filters import BaseFilter
from aiogram.types import Message
@@ -6,9 +8,13 @@ from config import ADMIN_ID
class IsAdminFilter(BaseFilter):
async def __call__(self, message: Message) -> bool:
if isinstance(ADMIN_ID, list):
return message.from_user.id in ADMIN_ID
elif isinstance(ADMIN_ID, int):
return message.from_user.id == ADMIN_ID
else:
try:
admin_ids: Union[int, list[int]] = ADMIN_ID
if isinstance(admin_ids, list):
return message.from_user.id in admin_ids
return message.from_user.id == admin_ids
except Exception:
return False
handlers/utils.py
+60 -37
View File
@@ -1,5 +1,6 @@
import random
import re
from typing import Optional
import asyncpg
@@ -9,35 +10,48 @@ from logger import logger
def sanitize_key_name(key_name: str) -> str:
"""
Очищает название ключа, оставляя только допустимые символы.
Args:
key_name (str): Исходное название ключа.
Returns:
str: Очищенное название ключа в нижнем регистре.
"""
return re.sub(r"[^a-z0-9@._-]", "", key_name.lower())
def generate_random_email():
"""Генерирует случайный набор символов."""
random_string = "".join(random.choices("abcdefghijklmnopqrstuvwxyz0123456789", k=6))
return random_string
async def get_least_loaded_cluster():
def generate_random_email(length: int = 6) -> str:
"""
Функция для получения кластера с наименьшей загрузкой (по количеству ключей).
Возвращает идентификатор кластера с наименьшей загрузкой или первый кластер из конфигурации,
если загруженность не определяется. В случае отсутствия кластеров с номером, возвращает 'cluster1'.
Генерирует случайный email с заданной длиной.
Args:
length (int, optional): Длина случайной строки. По умолчанию 6.
Returns:
str: Сгенерированная случайная строка.
"""
cluster_loads = {}
return "".join(random.choices("abcdefghijklmnopqrstuvwxyz0123456789", k=length))
conn = await asyncpg.connect(DATABASE_URL)
try:
keys = await conn.fetch("SELECT * FROM keys")
for key in keys:
cluster_id = key["server_id"]
if re.match(r"^cluster\d+$", cluster_id):
if cluster_id not in cluster_loads:
cluster_loads[cluster_id] = 0
cluster_loads[cluster_id] += 1
finally:
await conn.close()
async def get_least_loaded_cluster() -> str:
"""
Определяет кластер с наименьшей загрузкой.
Returns:
str: Идентификатор наименее загруженного кластера.
"""
cluster_loads: dict[str, int] = {}
async with asyncpg.create_pool(DATABASE_URL) as pool:
async with pool.acquire() as conn:
keys = await conn.fetch("SELECT * FROM keys")
for key in keys:
cluster_id = key["server_id"]
if re.match(r"^cluster\d+$", cluster_id):
cluster_loads[cluster_id] = cluster_loads.get(cluster_id, 0) + 1
logger.info(f"Cluster loads: {cluster_loads}")
@@ -51,29 +65,38 @@ async def get_least_loaded_cluster():
logger.info(f"Available clusters from config: {available_clusters}")
if available_clusters:
logger.info(
f"Returning the first available cluster: {available_clusters[0]}"
)
return available_clusters[0]
else:
logger.warning("No valid clusters found in config, returning 'cluster1'.")
return "cluster1"
selected_cluster = available_clusters[0]
logger.info(f"Returning the first available cluster: {selected_cluster}")
return selected_cluster
logger.warning("No valid clusters found in config, returning 'cluster1'.")
return "cluster1"
least_loaded_cluster = min(cluster_loads, key=cluster_loads.get)
logger.info(f"Least loaded cluster selected: {least_loaded_cluster}")
return least_loaded_cluster
async def handle_error(tg_id, callback_query, message):
async def handle_error(
tg_id: int, callback_query: Optional[object] = None, message: str = ""
) -> None:
"""
Обрабатывает ошибку, отправляя сообщение пользователю.
Args:
tg_id (int): Идентификатор пользователя в Telegram.
callback_query (Optional[object], optional): Объект запроса обратного вызова. По умолчанию None.
message (str, optional): Текст сообщения об ошибке. По умолчанию пустая строка.
"""
try:
try:
await bot.delete_message(
chat_id=tg_id, message_id=callback_query.message.message_id
)
except Exception:
pass
if callback_query and hasattr(callback_query, "message"):
try:
await bot.delete_message(
chat_id=tg_id, message_id=callback_query.message.message_id
)
except Exception as delete_error:
logger.warning(f"Не удалось удалить сообщение: {delete_error}")
await bot.send_message(tg_id, message, parse_mode="HTML")
middlewares/admin.py
+13 -14
View File
@@ -1,4 +1,4 @@
from typing import Any, Awaitable, Callable, Dict
from typing import Any, Awaitable, Callable, Dict,Union
from aiogram import BaseMiddleware
from aiogram.types import TelegramObject
@@ -14,17 +14,16 @@ class AdminMiddleware(BaseMiddleware):
event: TelegramObject,
data: Dict[str, Any],
) -> Any:
data["admin"] = False
try:
if isinstance(ADMIN_ID, list):
if event.from_user.id in ADMIN_ID:
data["admin"] = True
elif isinstance(ADMIN_ID, int):
if event.from_user.id == ADMIN_ID:
data["admin"] = True
else:
data["admin"] = False
except Exception as e:
logger.error(e)
data["admin"] = False
data["admin"] = self._check_admin_access(event)
return await handler(event, data)
def _check_admin_access(self, event: TelegramObject) -> bool:
try:
admin_ids: Union[int, list[int]] = ADMIN_ID
if isinstance(admin_ids, list):
return event.from_user.id in admin_ids
return event.from_user.id == admin_ids
except Exception as e:
logger.error(f"Ошибка проверки администратора: {e}")
return False
middlewares/database.py
+7 -6
View File
@@ -14,9 +14,10 @@ class DatabaseMiddleware(BaseMiddleware):
event: TelegramObject,
data: Dict[str, Any],
) -> Any:
session = await asyncpg.connect(DATABASE_URL)
data["session"] = session
try:
return await handler(event, data)
finally:
await session.close()
async with await asyncpg.create_pool(DATABASE_URL) as pool:
async with pool.acquire() as session:
data["session"] = session
try:
return await handler(event, data)
finally:
await pool.release(session)
middlewares/logging.py
+19 -12
View File
@@ -1,4 +1,4 @@
from typing import Any, Awaitable, Callable, Dict
from typing import Any, Awaitable, Callable, Dict, Optional
from aiogram import BaseMiddleware
from aiogram.types import CallbackQuery, Message, TelegramObject
@@ -13,23 +13,30 @@ class LoggingMiddleware(BaseMiddleware):
event: TelegramObject,
data: Dict[str, Any],
) -> Any:
user_info = self._extract_user_info(event)
logger.info(
f"Активность пользователя - "
f"ID пользователя: {user_info['user_id']}, "
f"Имя пользователя: {user_info['username']}, "
f"Действие: {user_info['action']}"
)
return await handler(event, data)
def _extract_user_info(self, event: TelegramObject) -> Dict[str, Optional[str]]:
user_id = None
username = None
action = None
if isinstance(event, Message):
user_id = event.from_user.id
username = event.from_user.username
user = event.from_user
user_id = user.id
username = user.username
action = f"Сообщение: {event.text}"
elif isinstance(event, CallbackQuery):
user_id = event.from_user.id
username = event.from_user.username
user = event.from_user
user_id = user.id
username = user.username
action = f"Обратный вызов: {event.data}"
logger.info(
f"Активность пользователя - "
f"ID пользователя: {user_id}, "
f"Имя пользователя: {username}, "
f"Действие: {action}"
)
return await handler(event, data)
return {"user_id": user_id, "username": username, "action": action}
middlewares/user.py
+12 -12
View File
@@ -13,16 +13,16 @@ class UserMiddleware(BaseMiddleware):
event: TelegramObject,
data: Dict[str, Any],
) -> Any:
user: User = data.get("event_from_user")
if user:
await upsert_user(
tg_id=user.id,
username=user.username,
first_name=user.first_name,
last_name=user.last_name,
language_code=user.language_code,
is_bot=user.is_bot,
)
if user := data.get("event_from_user"):
await self._process_user(user)
return await handler(event, data)
async def _process_user(self, user: User) -> None:
await upsert_user(
tg_id=user.id,
username=user.username,
first_name=user.first_name,
last_name=user.last_name,
language_code=user.language_code,
is_bot=user.is_bot,
)