refactor: balance handling and upload processing

2026-07-24 16:48:32 +05:00
parent 8bd5c80693
commit 9009147b52
9 changed files with 341 additions and 122 deletions
@@ -1,5 +1,5 @@
 from aiogram import Router
-from aiogram.types import CallbackQuery
+from aiogram.types import CallbackQuery, Message

 from bot.api.asf import get_all_bots
 from bot.logging_utils import get_logger
@@ -27,6 +27,28 @@ async def owned_callback_bot(callback: CallbackQuery, prefix: str):
    return await get_owned_bot(callback.from_user.id, account_id)


+async def send_bots_menu(message: Message, user_id: int) -> None:
+    """Отправляет отдельное сообщение со списком ботов пользователя."""
+    response = get_all_bots()
+    if response.status_code != 200 or not response.json().get("Success"):
+        await message.answer("Не удалось загрузить список ботов")
+        return
+
+    asf_bots = response.json().get("Result", {})
+    if user_id == ADMIN_ID:
+        await ensure_admin_compatibility(list(asf_bots))
+
+    accounts = await list_user_bot_accounts(user_id)
+    rows = [
+        (account.id, account.bot_name, asf_bots.get(account.bot_name, {}))
+        for account in accounts
+    ]
+    await message.answer(
+        f"<b>🤖 Ваших ботов:</b> {len(rows)}\n\n{get_farm_summary({name: bot for _, name, bot in rows})}",
+        reply_markup=bots_keyboard(rows),
+    )
+
+
@router.callback_query(lambda c: c.data == "bot_loading")
 async def bot_loading_handler(callback: CallbackQuery) -> None:
    await callback.answer(
@@ -37,10 +59,8 @@ async def bot_loading_handler(callback: CallbackQuery) -> None:
@router.callback_query(lambda c: c.data == "bots")
 async def bots_handler(callback: CallbackQuery) -> None:
    await stop_twofa_task(callback.message.message_id)
-
-    response = get_all_bots()
    await callback.answer()
-
+    response = get_all_bots()
    if response.status_code != 200 or not response.json().get("Success"):
        await callback.message.edit_text("Не удалось загрузить список ботов")
        return
@@ -54,7 +74,6 @@ async def bots_handler(callback: CallbackQuery) -> None:
        (account.id, account.bot_name, asf_bots.get(account.bot_name, {}))
        for account in accounts
    ]
-
    await callback.message.edit_text(
        f"<b>🤖 Ваших ботов:</b> {len(rows)}\n\n{get_farm_summary({name: bot for _, name, bot in rows})}",
        reply_markup=bots_keyboard(rows),
@@ -11,7 +11,10 @@ from bot.services.redeem import (
    redeem_all_bots_keys,
    redeem_single_bot_keys,
 )
-from bot.services.bot_accounts import list_user_bot_accounts
+from bot.services.bot_accounts import (
+    get_owned_bot,
+    list_user_activation_bot_accounts,
+)

 from bot.states import ConsoleFlow, RedeemFlow

@@ -103,12 +106,19 @@ async def handle_single_bot_redeem(message: Message, state: FSMContext):

    bot_name = state_data.get("bot_name")
    account_id = state_data.get("account_id")
-    if not isinstance(bot_name, str) or not bot_name:
+    if not isinstance(account_id, int):
        await message.answer("Контекст потерян, попробуйте еще раз.")
        await state.clear()

        return

+    account = await get_owned_bot(message.from_user.id, account_id)
+    if account is None:
+        await message.answer("Бот больше недоступен для активации ключей.")
+        await state.clear()
+        return
+    bot_name = account.bot_name
+
    chat_id = message.chat.id
    menu_message_id = state_data.get("menu_message_id")

@@ -257,7 +267,7 @@ async def handle_all_bots_redeem(message: Message, state: FSMContext):
    )

    try:
-        accounts = await list_user_bot_accounts(message.from_user.id)
+        accounts = await list_user_activation_bot_accounts(message.from_user.id)
        await checking.edit_text(
            await redeem_all_bots_keys(keys, [account.bot_name for account in accounts])
        )
@@ -105,9 +105,7 @@ async def restart_asf(callback: CallbackQuery) -> None:
            )
            return

-        # ! Меняется подход к перезапуску ASF: вместо того чтобы ждать
-        # ! 6 секунд, лучше проверять статус ASF до тех пор он не запустится
-        #  await asyncio.sleep(6)
+        # Проверяем статус ASF до его запуска, не ожидая фиксированные 6 секунд.

        status = get_asf_status()
        await asyncio.sleep(1)
@@ -2,9 +2,8 @@ import asyncio
 import hashlib
 import io
 import json
-import posixpath
-import re
 import zipfile
+from dataclasses import dataclass
 from datetime import datetime, timezone
 from pathlib import PurePosixPath

@@ -13,21 +12,26 @@ from aiogram.fsm.context import FSMContext
 from aiogram.types import CallbackQuery, Message

 from bot.api.asf import get_bot, save_bot_config
+from bot.config import ADMIN_ID
+from bot.handlers.bots import send_bots_menu
+from bot.logging_utils import get_logger
+from bot.services.balance import (
+    InsufficientBalance,
+    increase_balance,
+    debit_user,
+    get_bot_slot_price,
+    get_user_balance,
+)
 from bot.services.bot_accounts import (
    bot_name_exists,
    create_bot_account,
    user_hash_exists,
 )
-from bot.services.balance import (
-    InsufficientBalance,
-    credit_user,
-    debit_user,
-    get_bot_slot_price,
-)
 from bot.states import UploadFlow
 from bot.ui.keyboards import back_keyboard

 router = Router()
+logger = get_logger(__name__)
 _groups: dict[str, list[Message]] = {}
 _group_tasks: dict[str, asyncio.Task] = {}

@@ -36,8 +40,14 @@ _group_tasks: dict[str, asyncio.Task] = {}
 async def upload_menu(callback: CallbackQuery, state: FSMContext) -> None:
    await callback.answer()
    await state.set_state(UploadFlow.waiting_upload)
+
+    balance = await get_user_balance(callback.from_user.id)
+    slot_price = await get_bot_slot_price()
+
    await callback.message.edit_text(
-        "Отправьте JSON-конфиг, несколько JSON одним альбомом или ZIP-архив с JSON-файлами.",
+        "Отправьте документы .json или .zip с JSON-конфигурациями.\n\n"
+        f"💰 Ваш баланс: {balance} ₽\n"
+        f"💳 Цена одного слота: {slot_price} ₽",
        reply_markup=back_keyboard(callback_data="bots"),
    )

@@ -46,7 +56,9 @@ def _decode_config(raw: bytes) -> dict:
    value = json.loads(raw.decode("utf-8-sig"))
    if not isinstance(value, dict):
        raise ValueError("JSON должен содержать объект конфигурации")
-
+    name = value.get("SteamLogin")
+    if not isinstance(name, str) or not name.strip():
+        raise ValueError("отсутствует непустой SteamLogin")
    return value


@@ -58,8 +70,10 @@ def _zip_entries(raw: bytes) -> list[tuple[str, bytes]]:
            path = PurePosixPath(name)
            if path.is_absolute() or ".." in path.parts:
                raise ValueError("ZIP содержит небезопасный путь")
-            if item.is_dir() or not name.lower().endswith(".json"):
+            if item.is_dir():
                continue
+            if not name.lower().endswith(".json"):
+                raise ValueError(f"ZIP содержит неподдерживаемый файл: {name}")
            entries.append((name, archive.read(item)))
    return entries

@@ -70,91 +84,235 @@ async def _download(message: Message) -> bytes:
    return buffer.getvalue()


-async def _process(user_id: int, files: list[tuple[str, bytes]]) -> str:
-    results = []
-    seen_hashes: set[str] = set()
-    seen_names: set[str] = set()
+@dataclass
+class _Candidate:
+    filename: str
+    config: dict
+    digest: str
+
+
+def _digest(config: dict) -> str:
+    canonical = json.dumps(
+        config, ensure_ascii=False, sort_keys=True, separators=(",", ":")
+    )
+    return hashlib.sha256(canonical.encode("utf-8")).hexdigest()
+
+
+async def _preflight(
+    user_id: int, files: list[tuple[str, bytes]]
+) -> tuple[list[_Candidate], list[str]]:
+    candidates: list[_Candidate] = []
+    errors: list[str] = []
+    parsed: list[tuple[str, dict, str]] = []

    for filename, raw in files:
        try:
            config = _decode_config(raw)
-            digest = hashlib.sha256(raw).hexdigest()
-            name = config.get("SteamLogin")
+            name = config["SteamLogin"].strip()
+            digest = _digest(config)
+            parsed.append((filename, config, digest))

-            if name in seen_names or await bot_name_exists(name):
-                raise ValueError("бот с таким именем уже существует")
+        except (ValueError, json.JSONDecodeError, UnicodeDecodeError) as error:
+            logger.warning(
+                "Предварительная проверка отклонена: пользователь=%s файл=%s ошибка=%s",
+                user_id,
+                filename,
+                error,
+            )
+            errors.append(f"❌ {filename}: {error}")

-            if digest in seen_hashes or await user_hash_exists(user_id, digest):
-                raise ValueError("такой конфиг уже загружался")
+        except Exception as error:
+            logger.exception(
+                "Ошибка предварительной проверки: пользователь=%s файл=%s",
+                user_id,
+                filename,
+            )
+            errors.append(f"❌ {filename}: ошибка проверки ({error})")

-            existing = get_bot(name)
-            if existing.status_code == 200 and existing.json().get("Result"):
-                raise ValueError("бот с таким именем уже существует в ASF")
+    name_counts: dict[str, int] = {}
+    hash_counts: dict[str, int] = {}
+    for _, config, digest in parsed:
+        name_counts[config["SteamLogin"].strip()] = (
+            name_counts.get(config["SteamLogin"].strip(), 0) + 1
+        )
+        hash_counts[digest] = hash_counts.get(digest, 0) + 1

-            slot_price = await get_bot_slot_price()
-            await debit_user(user_id, slot_price)
-            try:
-                response = save_bot_config(name, config)
-                if response.status_code != 200 or not response.json().get("Success"):
-                    raise ValueError(
-                        f"ASF отклонил конфиг (HTTP {response.status_code})"
-                    )
+    for filename, config, digest in parsed:
+        name = config["SteamLogin"].strip()
+        try:
+            if name_counts[name] > 1:
+                raise ValueError("дублируется имя бота в этой загрузке")

-                await create_bot_account(
-                    user_id,
-                    bot_name=name,
-                    steam_id=(
-                        str(config.get("SteamID")) if config.get("SteamID") else None
-                    ),
-                    source_filename=filename[:255],
-                    config_sha256=digest,
-                    upload_state="attached",
-                    is_attached=True,
-                    attached_at=datetime.now(timezone.utc),
+            if hash_counts[digest] > 1 or await user_hash_exists(user_id, digest):
+                raise ValueError(
+                    "такая конфигурация уже загружалась или дублируется в загрузке"
                )

-            except Exception:
-                await credit_user(user_id, slot_price)
-                raise
+            if await bot_name_exists(name):
+                raise ValueError("бот с таким именем уже есть в базе данных")
+            existing = get_bot(name)

-            seen_names.add(name)
-            seen_hashes.add(digest)
-            results.append(f"✅ {filename}")
+            if existing.status_code == 200 and existing.json().get("Result"):
+                raise ValueError("бот с таким именем уже есть в ASF")

-        except (
-            InsufficientBalance,
-            ValueError,
-            json.JSONDecodeError,
-            UnicodeDecodeError,
-            zipfile.BadZipFile,
-        ) as error:
-            results.append(f"❌ {filename}: {error}")
+            candidates.append(_Candidate(filename, config, digest))
+            logger.info(
+                "Предварительная проверка пройдена: пользователь=%s файл=%s бот=%s",
+                user_id,
+                filename,
+                name,
+            )

-        except Exception:
-            results.append(f"❌ {filename}: внутренняя ошибка")
+        except (ValueError, json.JSONDecodeError, UnicodeDecodeError) as error:
+            logger.warning(
+                "Предварительная проверка отклонена: пользователь=%s файл=%s ошибка=%s",
+                user_id,
+                filename,
+                error,
+            )
+            errors.append(f"❌ {filename}: {error}")

-    return f"Обработано файлов: {len(files)}"
+        except Exception as error:
+            logger.exception(
+                "Ошибка предварительной проверки: пользователь=%s файл=%s",
+                user_id,
+                filename,
+            )
+            errors.append(f"❌ {filename}: ошибка проверки ({error})")
+
+    return candidates, errors
+
+
+async def _process(user_id: int, files: list[tuple[str, bytes]]) -> str:
+    candidates, results = await _preflight(user_id, files)
+    if not candidates:
+        return "\n".join(results) or "❌ Нет конфигураций для загрузки"
+
+    is_admin = user_id == ADMIN_ID
+    slot_price = await get_bot_slot_price()
+    total = slot_price * len(candidates)
+    reserved = False
+    if not is_admin:
+        try:
+            await debit_user(user_id, total)
+            reserved = True
+            logger.info(
+                "Средства зарезервированы: пользователь=%s сумма=%s ₽ количество=%s",
+                user_id,
+                total,
+                len(candidates),
+            )
+
+        except InsufficientBalance:
+            balance = await get_user_balance(user_id)
+            logger.warning(
+                "Недостаточно средств для загрузки: пользователь=%s баланс=%s ₽ требуется=%s ₽",
+                user_id,
+                balance,
+                total,
+            )
+
+            results.append(
+                f"❌ Недостаточно средств: нужно {total} ₽ для {len(candidates)} конф. "
+                f"(доступно {balance} ₽)\n"
+            )
+
+            results.extend(
+                f"❌ {candidate.filename}: загрузка не начата"
+                for candidate in candidates
+            )
+            return "\n".join(results)
+
+    for candidate in candidates:
+        name = candidate.config["SteamLogin"].strip()
+        try:
+            response = save_bot_config(name, candidate.config)
+            logger.info(
+                "Ответ ASF на загрузку: пользователь=%s бот=%s статус=%s",
+                user_id,
+                name,
+                response.status_code,
+            )
+
+            payload = response.json()
+            if response.status_code != 200 or not payload.get("Success"):
+                detail = (
+                    payload.get("Message")
+                    or payload.get("Result")
+                    or f"HTTP {response.status_code}"
+                )
+                raise ValueError(f"ASF отклонил конфигурацию: {detail}")
+
+            await create_bot_account(
+                user_id,
+                bot_name=name,
+                steam_id=(
+                    str(candidate.config.get("SteamID"))
+                    if candidate.config.get("SteamID")
+                    else None
+                ),
+                source_filename=candidate.filename[:255],
+                config_sha256=candidate.digest,
+                upload_state="attached",
+                is_attached=True,
+                attached_at=datetime.now(timezone.utc),
+            )
+            logger.info(
+                "Регистрация в базе завершена: пользователь=%s бот=%s", user_id, name
+            )
+            results.append(
+                f"✅ {candidate.filename}: ASF принял, аккаунт зарегистрирован ({name})"
+            )
+
+        except Exception as error:
+            logger.exception(
+                "Ошибка загрузки: пользователь=%s файл=%s бот=%s",
+                user_id,
+                candidate.filename,
+                name,
+            )
+            if reserved:
+                await increase_balance(user_id, slot_price)
+                logger.info(
+                    "Средства возвращены: пользователь=%s сумма=%s ₽ файл=%s",
+                    user_id,
+                    slot_price,
+                    candidate.filename,
+                )
+            results.append(f"❌ {candidate.filename}: {error}")
+
+    return "\n".join(results)


 async def _finish_group(key: str, state: FSMContext) -> None:
    await asyncio.sleep(0.8)
    messages = _groups.pop(key, [])
    _group_tasks.pop(key, None)
+
    if not messages:
        return
+
    files = []
    for message in messages:
+        filename = message.document.file_name or "config.json"
        try:
-            files.append(
-                (message.document.file_name or "config.json", await _download(message))
-            )
+            files.append((filename, await _download(message)))
        except Exception:
-            files.append((message.document.file_name or "config.json", b""))
+            files.append((filename, b""))

-    await messages[0].answer(await _process(messages[0].from_user.id, files))
+    await _send_result_and_menu(messages[0], messages[0].from_user.id, files)
    await state.clear()


+async def _send_result_and_menu(
+    message: Message, user_id: int, files: list[tuple[str, bytes]]
+) -> None:
+    result = await _process(user_id, files)
+    await message.answer(result)
+    if "✅" in result:
+        await send_bots_menu(message, user_id)
+
+
@router.message(UploadFlow.waiting_upload, lambda message: bool(message.document))
 async def upload_document(message: Message, state: FSMContext) -> None:
    document = message.document
@@ -167,24 +325,28 @@ async def upload_document(message: Message, state: FSMContext) -> None:
            _group_tasks[key] = asyncio.create_task(_finish_group(key, state))
        return

-    raw = await _download(message)
-    if filename.lower().endswith(".zip"):
-        try:
+    try:
+        raw = await _download(message)
+        if filename.lower().endswith(".zip"):
            files = _zip_entries(raw)
-
-        except Exception as error:
-            await message.answer(f"❌ {filename}: {error}")
-            await state.clear()
-            return
-
-        if not files:
-            await message.answer("❌ В архиве нет JSON-файлов")
+            if not files:
+                await message.answer("❌ В архиве нет JSON-файлов")
+            else:
+                await _send_result_and_menu(message, message.from_user.id, files)
+        elif filename.lower().endswith(".json"):
+            await _send_result_and_menu(
+                message, message.from_user.id, [(filename, raw)]
+            )
        else:
-            await message.answer(await _process(message.from_user.id, files))
+            await message.answer("❌ Поддерживаются только документы .json и .zip")
+    except (zipfile.BadZipFile, ValueError) as error:
+        await message.answer(f"❌ {filename}: {error}")
+    finally:
+        await state.clear()

-    elif filename.lower().endswith(".json"):
-        await message.answer(await _process(message.from_user.id, [(filename, raw)]))
-    else:
-        await message.answer("❌ Поддерживаются только .json и .zip")

-    await state.clear()
+@router.message(UploadFlow.waiting_upload)
+async def upload_non_document(message: Message, state: FSMContext) -> None:
+    await message.answer(
+        "❌ Отправьте конфигурацию именно документом .json или архивом .zip"
+    )