Archived
feat(bot): add bot account ownership and balance
This commit is contained in:
@@ -0,0 +1,190 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user