ASF-Control-Bot/bot/handlers/uploads.py

import asyncio
import hashlib
import io
import json
import posixpath
import re
import zipfile
from datetime import datetime, timezone
from pathlib import PurePosixPath

from aiogram import Router
from aiogram.fsm.context import FSMContext
from aiogram.types import CallbackQuery, Message

from bot.api.asf import get_bot, save_bot_config
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()
_groups: dict[str, list[Message]] = {}
_group_tasks: dict[str, asyncio.Task] = {}


@router.callback_query(lambda c: c.data == "upload")
async def upload_menu(callback: CallbackQuery, state: FSMContext) -> None:
    await callback.answer()
    await state.set_state(UploadFlow.waiting_upload)
    await callback.message.edit_text(
        "Отправьте JSON-конфиг, несколько JSON одним альбомом или ZIP-архив с JSON-файлами.",
        reply_markup=back_keyboard(callback_data="bots"),
    )


def _decode_config(raw: bytes) -> dict:
    value = json.loads(raw.decode("utf-8-sig"))
    if not isinstance(value, dict):
        raise ValueError("JSON должен содержать объект конфигурации")

    return value


def _zip_entries(raw: bytes) -> list[tuple[str, bytes]]:
    entries = []
    with zipfile.ZipFile(io.BytesIO(raw)) as archive:
        for item in archive.infolist():
            name = item.filename.replace("\\", "/")
            path = PurePosixPath(name)
            if path.is_absolute() or ".." in path.parts:
                raise ValueError("ZIP содержит небезопасный путь")
            if item.is_dir() or not name.lower().endswith(".json"):
                continue
            entries.append((name, archive.read(item)))
    return entries


async def _download(message: Message) -> bytes:
    buffer = io.BytesIO()
    await message.bot.download(message.document.file_id, destination=buffer)
    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()

    for filename, raw in files:
        try:
            config = _decode_config(raw)
            digest = hashlib.sha256(raw).hexdigest()
            name = config.get("SteamLogin")

            if name in seen_names or await bot_name_exists(name):
                raise ValueError("бот с таким именем уже существует")

            if digest in seen_hashes or await user_hash_exists(user_id, digest):
                raise ValueError("такой конфиг уже загружался")

            existing = get_bot(name)
            if existing.status_code == 200 and existing.json().get("Result"):
                raise ValueError("бот с таким именем уже существует в ASF")

            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})"
                    )

                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),
                )

            except Exception:
                await credit_user(user_id, slot_price)
                raise

            seen_names.add(name)
            seen_hashes.add(digest)
            results.append(f"✅ {filename}")

        except (
            InsufficientBalance,
            ValueError,
            json.JSONDecodeError,
            UnicodeDecodeError,
            zipfile.BadZipFile,
        ) as error:
            results.append(f"❌ {filename}: {error}")

        except Exception:
            results.append(f"❌ {filename}: внутренняя ошибка")

    return f"Обработано файлов: {len(files)}"


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:
        try:
            files.append(
                (message.document.file_name or "config.json", await _download(message))
            )
        except Exception:
            files.append((message.document.file_name or "config.json", b""))

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


@router.message(UploadFlow.waiting_upload, lambda message: bool(message.document))
async def upload_document(message: Message, state: FSMContext) -> None:
    document = message.document
    filename = document.file_name or "config.json"

    if message.media_group_id:
        key = f"{message.chat.id}:{message.media_group_id}"
        _groups.setdefault(key, []).append(message)
        if key not in _group_tasks:
            _group_tasks[key] = asyncio.create_task(_finish_group(key, state))
        return

    raw = await _download(message)
    if filename.lower().endswith(".zip"):
        try:
            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-файлов")
        else:
            await message.answer(await _process(message.from_user.id, files))

    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()