Refactor and improve security with cryptographically secure random generation

- Replace `random` module with `secrets` for cryptographically secure random choices
- Update type hints for async API methods in client.py
- Optimize file reading with aiofiles in multiple handlers
- Improve error handling and logging in various modules
- Update Python target version to 3.12 in pyproject.toml
- Remove unused imports and improve code quality
This commit is contained in:
Zakhar Izmaylov
2025-01-30 22:35:09 +03:00
parent ee9e65d692
commit 899da825cb
14 changed files with 64 additions and 52 deletions
+2 -1
View File
@@ -97,9 +97,10 @@ async def create_backup_and_send_to_admins(xui) -> None:
async def _send_backup_to_admins(backup_file_path: str) -> None:
try:
from bot import bot
import aiofiles
from bot import bot
async with aiofiles.open(backup_file_path, "rb") as backup_file:
backup_data = await backup_file.read()
backup_input_file = BufferedInputFile(file=backup_data, filename=os.path.basename(backup_file_path))
+3 -3
View File
@@ -23,7 +23,7 @@ class ClientConfig:
sub_id: str
async def add_client(xui: py3xui.API, config: ClientConfig) -> dict[str, Any]:
async def add_client(xui: py3xui.AsyncApi, config: ClientConfig) -> dict[str, Any]:
"""
Добавляет клиента на сервер через 3x-ui.
@@ -66,7 +66,7 @@ async def add_client(xui: py3xui.API, config: ClientConfig) -> dict[str, Any]:
async def extend_client_key(
xui: py3xui.API, inbound_id: int, email: str, new_expiry_time: int, client_id: str, total_gb: int, sub_id: str
xui: py3xui.AsyncApi, inbound_id: int, email: str, new_expiry_time: int, client_id: str, total_gb: int, sub_id: str
) -> bool | None:
"""
Обновляет срок действия ключа клиента.
@@ -117,7 +117,7 @@ async def extend_client_key(
async def delete_client(
xui: py3xui.API,
xui: py3xui.AsyncApi,
inbound_id: int,
email: str,
client_id: str,
+2 -2
View File
@@ -85,8 +85,8 @@ async def handle_message_input(message: types.Message, state: FSMContext, sessio
try:
await message.bot.send_message(chat_id=tg_id, text=text_message)
success_count += 1
except Exception:
pass
except Exception as e:
logger.error(e)
text = (
f"📤 Рассылка завершена!"
+12 -8
View File
@@ -1,3 +1,5 @@
from typing import Any
import asyncpg
from aiogram import F, Router, types
from aiogram.fsm.context import FSMContext
@@ -33,7 +35,7 @@ class AdminServersEditor(StatesGroup):
IsAdminFilter(),
)
async def handle_servers(callback_query: types.CallbackQuery):
servers = await get_servers_from_db()
servers = await get_servers()
text = (
"<b>🔧 Управление кластерами</b>\n\n"
@@ -227,8 +229,7 @@ async def handle_inbound_id_input(message: types.Message, state: FSMContext):
@router.callback_query(AdminServerEditorCallback.filter(F.action == "clusters_manage"), IsAdminFilter())
async def handle_clusters_manage(
callback_query: types.CallbackQuery,
callback_data: AdminServerEditorCallback,
callback_query: types.CallbackQuery, callback_data: AdminServerEditorCallback, session: Any
):
cluster_name = callback_data.data
@@ -242,7 +243,9 @@ async def handle_clusters_manage(
@router.callback_query(AdminServerEditorCallback.filter(F.action == "servers_availability"), IsAdminFilter())
async def handle_servers_availability(callback_query: types.CallbackQuery, callback_data: AdminServerEditorCallback):
async def handle_servers_availability(
callback_query: types.CallbackQuery, callback_data: AdminServerEditorCallback, session: Any
):
cluster_name = callback_data.data
servers = await get_servers(session)
@@ -279,7 +282,7 @@ async def handle_servers_availability(callback_query: types.CallbackQuery, callb
@router.callback_query(AdminServerEditorCallback.filter(F.action == "servers_manage"), IsAdminFilter())
async def handle_servers_manage(callback_query: types.CallbackQuery, callback_data: AdminServerEditorCallback):
server_name = callback_data.data
servers = await get_servers_from_db()
servers = await get_servers()
cluster_name, server = next(
((c, s) for c, cs in servers.items() for s in cs if s["server_name"] == server_name), (None, None)
@@ -316,7 +319,9 @@ async def handle_servers_delete(callback_query: types.CallbackQuery, callback_da
@router.callback_query(AdminServerEditorCallback.filter(F.action == "servers_delete_confirm"), IsAdminFilter())
async def handle_servers_delete_confirm(callback_query: types.CallbackQuery, callback_data: AdminServerEditorCallback):
async def handle_servers_delete_confirm(
callback_query: types.CallbackQuery, callback_data: AdminServerEditorCallback, session: Any
):
server_name = callback_data.data
await delete_server(server_name, session)
@@ -350,8 +355,7 @@ async def handle_servers_add(
@router.callback_query(AdminServerEditorCallback.filter(F.action == "clusters_backup"), IsAdminFilter())
async def handle_clusters_backup(
callback_query: types.CallbackQuery,
callback_data: AdminServerEditorCallback,
callback_query: types.CallbackQuery, callback_data: AdminServerEditorCallback, session: Any
):
cluster_name = callback_data.data
+5 -5
View File
@@ -11,6 +11,7 @@ from aiogram.utils.keyboard import InlineKeyboardBuilder
from config import TOTAL_GB
from database import (
delete_key,
delete_user_data,
get_client_id_by_email,
get_key_details,
@@ -22,7 +23,6 @@ from database import (
from filters.admin import IsAdminFilter
from handlers.keys.key_utils import (
delete_key_from_cluster,
delete_key_from_db,
renew_key_in_cluster,
)
from handlers.utils import sanitize_key_name
@@ -415,7 +415,7 @@ async def handle_expiry_time_input(message: types.Message, state: FSMContext, se
async def update_key_on_all_servers():
tasks = []
for cluster_name, cluster_servers in clusters.items():
for server in cluster_servers:
for _server in cluster_servers:
tasks.append(
asyncio.create_task(
renew_key_in_cluster(
@@ -493,12 +493,12 @@ async def process_callback_confirm_delete(callback_query: types.CallbackQuery, s
async def delete_key_from_servers(email, client_id):
tasks = []
for cluster_name, cluster_servers in clusters.items():
for server in cluster_servers:
for _server in cluster_servers:
tasks.append(delete_key_from_cluster(cluster_name, email, client_id))
await asyncio.gather(*tasks)
await delete_key_from_servers(email, client_id)
await delete_key_from_db(client_id, session)
await delete_key(client_id, session)
await callback_query.message.answer(response_message, reply_markup=builder.as_markup())
else:
@@ -570,7 +570,7 @@ async def delete_user(callback_query: types.CallbackQuery, session: Any):
tasks = []
for email, client_id in key_records:
servers = await get_servers(session)
for cluster_id, cluster in servers.items():
for cluster_id, _cluster in servers.items():
tasks.append(delete_key_from_cluster(cluster_id, email, client_id))
await asyncio.gather(*tasks)
except Exception as e:
+7 -8
View File
@@ -10,11 +10,10 @@ from aiogram.fsm.state import State, StatesGroup
from aiogram.types import CallbackQuery
from config import TOTAL_GB
from database import delete_user_data, get_client_id_by_email, get_servers_from_db, restore_trial, update_key_expiry
from database import delete_key, delete_user_data, get_client_id_by_email, get_servers, update_key_expiry, update_trial
from filters.admin import IsAdminFilter
from handlers.keys.key_utils import (
delete_key_from_cluster,
delete_key_from_db,
renew_key_in_cluster,
update_subscription,
)
@@ -168,7 +167,7 @@ async def handle_trial_restore(
):
tg_id = callback_data.tg_id
await restore_trial(tg_id, session)
await update_trial(tg_id, 0, session)
await callback_query.message.edit_text(text="✅ Триал успешно восстановлен!", reply_markup=build_editor_kb(tg_id))
@@ -479,7 +478,7 @@ async def handle_delete_key_confirm(
if record:
client_id = record["client_id"]
clusters = await get_servers_from_db()
clusters = await get_servers()
async def delete_key_from_servers():
tasks = []
@@ -489,7 +488,7 @@ async def handle_delete_key_confirm(
await asyncio.gather(*tasks)
await delete_key_from_servers()
await delete_key_from_db(client_id, session)
await delete_key(client_id, session)
await callback_query.message.edit_text(text="✅ Ключ успешно удален.", reply_markup=kb)
else:
@@ -515,8 +514,8 @@ async def handle_delete_user_confirm(
try:
tasks = []
for email, client_id in key_records:
servers = await get_servers_from_db()
for cluster_id, cluster in servers.items():
servers = await get_servers()
for cluster_id, _cluster in servers.items():
tasks.append(delete_key_from_cluster(cluster_id, email, client_id))
await asyncio.gather(*tasks)
except Exception as e:
@@ -619,7 +618,7 @@ async def change_expiry_time(expiry_time: int, email: str, session: Any) -> Exce
if not server_id:
return ValueError(f"User with client_id {server_id} was not found")
clusters = await get_servers_from_db()
clusters = await get_servers()
async def update_key_on_all_servers():
tasks = [
+4 -3
View File
@@ -1,4 +1,5 @@
import random
import secrets
from typing import Any
from aiogram import F, Router
@@ -7,7 +8,6 @@ from aiogram.types import CallbackQuery, Message
from aiogram.utils.keyboard import InlineKeyboardBuilder
from config import CAPTCHA_EMOJIS
from handlers.start import start_command
from logger import logger
router = Router()
@@ -27,7 +27,7 @@ async def generate_captcha(message: Message, state: FSMContext):
}
"""
# Выбираем случайный эмодзи и его описание из конфига
correct_emoji, correct_text = random.choice(list(CAPTCHA_EMOJIS.items()))
correct_emoji, correct_text = secrets.choice(list(CAPTCHA_EMOJIS.items()))
# Получаем 3 случайных неправильных эмодзи
wrong_emojis = random.sample([e for e in CAPTCHA_EMOJIS.keys() if e != correct_emoji], 3)
@@ -54,7 +54,8 @@ async def generate_captcha(message: Message, state: FSMContext):
@router.callback_query(F.data.startswith("captcha_"))
async def check_captcha(callback: CallbackQuery, state: FSMContext, session: Any, admin: bool):
"""Проверяет ответ пользователя на капчу"""
from handlers.start import start_command
selected_emoji = callback.data.split("captcha_")[1]
state_data = await state.get_data()
correct_emoji = state_data.get("correct_emoji")
-1
View File
@@ -57,7 +57,6 @@ async def send_instructions(
@router.callback_query(F.data.startswith("connect_pc|"))
async def process_connect_pc(callback_query: types.CallbackQuery, session: Any):
tg_id = callback_query.message.chat.id
key_name = callback_query.data.split("|")[1]
record = await get_key_details(key_name, session)
+2 -3
View File
@@ -48,7 +48,6 @@ from handlers.buttons.add_subscribe import (
)
from handlers.keys.key_utils import (
delete_key_from_cluster,
delete_key_from_db,
renew_key_in_cluster,
update_subscription,
)
@@ -354,7 +353,7 @@ async def process_callback_confirm_delete(callback_query: types.CallbackQuery, s
async def delete_key_from_servers():
try:
tasks = []
for cluster_id, cluster in servers.items():
for cluster_id, _cluster in servers.items():
tasks.append(delete_key_from_cluster(cluster_id, email, client_id))
await asyncio.gather(*tasks)
@@ -364,7 +363,7 @@ async def process_callback_confirm_delete(callback_query: types.CallbackQuery, s
asyncio.create_task(delete_key_from_servers())
await delete_key_from_db(client_id, session)
await delete_key(client_id, session)
else:
response_message = "Ключ не найден или уже удален."
+18 -12
View File
@@ -2,6 +2,7 @@ import asyncio
import os
from datetime import datetime, timedelta
import aiofiles
import asyncpg
import pytz
from aiogram import Bot, Router, types
@@ -169,10 +170,11 @@ async def process_10h_record(record, bot, conn):
)
if os.path.isfile(image_path):
with open(image_path, "rb") as image_file:
async with aiofiles.open(image_path, "rb") as image_file:
image_data = await image_file.read()
await bot.send_photo(
tg_id,
photo=BufferedInputFile(image_file.read(), filename="notify_10h.jpg"),
photo=BufferedInputFile(image_data, filename="notify_10h.jpg"),
caption=KEY_RENEWED.format(email=email),
reply_markup=keyboard,
)
@@ -258,10 +260,11 @@ async def process_24h_record(record, bot, conn):
)
if os.path.isfile(image_path):
with open(image_path, "rb") as image_file:
async with aiofiles.open(image_path, "rb") as image_file:
image_data = await image_file.read()
await bot.send_photo(
tg_id,
photo=BufferedInputFile(image_file.read(), filename="notify_24h.jpg"),
photo=BufferedInputFile(image_data, filename="notify_24h.jpg"),
caption=KEY_RENEWED.format(email=email),
reply_markup=keyboard,
)
@@ -286,10 +289,11 @@ async def send_renewal_notification(bot, tg_id, email, message, conn, client_id,
image_path = os.path.join("img", "notify_24h.jpg")
if os.path.isfile(image_path):
with open(image_path, "rb") as image_file:
async with aiofiles.open(image_path, "rb") as image_file:
image_data = await image_file.read()
await bot.send_photo(
tg_id,
photo=BufferedInputFile(image_file.read(), filename="notify_24h.jpg"),
photo=BufferedInputFile(image_data, filename="notify_24h.jpg"),
caption=message,
reply_markup=keyboard.as_markup(),
)
@@ -298,7 +302,7 @@ async def send_renewal_notification(bot, tg_id, email, message, conn, client_id,
logger.info(f"Уведомление отправлено пользователю {tg_id}.")
await conn.execute(f"UPDATE keys SET {flag} = TRUE WHERE client_id = $1", client_id)
await conn.execute("UPDATE keys SET notified_24h = $1 WHERE client_id = $2", flag, client_id)
except Exception as e:
logger.error(f"Ошибка при отправке уведомления пользователю {tg_id}: {e}")
@@ -456,10 +460,11 @@ async def process_key(record, bot, conn):
try:
if os.path.isfile(image_path):
with open(image_path, "rb") as image_file:
async with aiofiles.open(image_path, "rb") as image_file:
image_data = await image_file.read()
await bot.send_photo(
tg_id,
photo=BufferedInputFile(image_file.read(), filename="notify_expired.jpg"),
photo=BufferedInputFile(image_data, filename="notify_expired.jpg"),
caption=KEY_RENEWED.format(email=email),
reply_markup=keyboard,
)
@@ -473,10 +478,11 @@ async def process_key(record, bot, conn):
message_expired = f"Ваша подписка {email} истекла. Пополните баланс для продления."
try:
if os.path.isfile(image_path):
with open(image_path, "rb") as image_file:
async with aiofiles.open(image_path, "rb") as image_file:
image_data = await image_file.read()
await bot.send_photo(
tg_id,
photo=BufferedInputFile(image_file.read(), filename="notify_expired.jpg"),
photo=BufferedInputFile(image_data, filename="notify_expired.jpg"),
caption=message_expired,
reply_markup=keyboard,
)
@@ -511,7 +517,7 @@ async def process_key(record, bot, conn):
async def check_online_users():
servers = await get_servers()
for cluster_id, cluster in servers.items():
for _cluster_id, cluster in servers.items():
for server_id, server in enumerate(cluster):
xui = AsyncApi(server["api_url"], username=ADMIN_USERNAME, password=ADMIN_PASSWORD)
await xui.login()
+3 -1
View File
@@ -1,6 +1,8 @@
import json
import random
import re
import secrets
import string
import aiohttp
import asyncpg
@@ -49,7 +51,7 @@ def generate_random_email(length: int = 6) -> str:
Returns:
str: Сгенерированная случайная строка.
"""
return "".join(random.choices("abcdefghijklmnopqrstuvwxyz0123456789", k=length))
return "".join(secrets.choice(string.ascii_lowercase + string.digits) for _ in range(length)) if length > 0 else ""
async def get_least_loaded_cluster() -> str:
+3 -2
View File
@@ -9,6 +9,7 @@ from keyboards.admin.panel_kb import AdminPanelCallback
from keyboards.admin.sender_kb import AdminSenderCallback
from keyboards.admin.servers_kb import AdminServerEditorCallback
from keyboards.admin.users_kb import AdminUserEditorCallback, AdminUserKeyEditorCallback
from logger import logger
pass_callbacks = [
AdminPanelCallback,
@@ -31,8 +32,8 @@ class DeleteMessageMiddleware(BaseMiddleware):
if not event.text or not event.text.startswith("/start"):
try:
await event.bot.delete_message(event.chat.id, event.message_id - 1)
except Exception:
pass
except Exception as e:
logger.error(e)
await event.delete()
if isinstance(event, CallbackQuery):
+2 -2
View File
@@ -1,10 +1,10 @@
[tool.ruff]
line-length = 120
target-version = "py310"
target-version = "py312"
[tool.ruff.lint]
select = ["E", "F", "W", "I", "N", "UP", "ANN", "ASYNC", "S", "BLE", "FBT", "B", "A", "C4", "DTZ", "T10", "ISC", "ICN", "G", "PIE"]
ignore = ["ANN101", "ANN102", "S101",'ANN201','ANN001','BLE001']
ignore = ["ANN101", "ANN102", "S101",'ANN201','ANN001','BLE001','W291','ANN401','DTZ003','DTZ005','F401','FBT002','FBT001','FBT003','A005','E501','UP017','DTZ004','W293','ANN202','DTZ007']
exclude = [
".git",
"venv",
+1 -1
View File
@@ -33,7 +33,7 @@ async def sync_servers_with_db():
logger.info("Подключение к базе данных для синхронизации серверов успешно.")
for cluster_name, servers in CLUSTERS.items():
for server_key, server_info in servers.items():
for _server_key, server_info in servers.items():
exists = await check_unique_server_name(server_info["name"], conn, cluster_name)
if not exists: