feat(bot): add bot account ownership and balance

2026-07-24 10:44:48 +05:00
parent 87385228ea
commit 8bd5c80693
22 changed files with 1008 additions and 393 deletions
@@ -0,0 +1,169 @@
+from sqlalchemy import func, select, update
+
+from bot.config import ADMIN_ID
+from bot.db.models import AppSetting, BotAccount, User
+from bot.db.session import async_session
+
+BOT_SLOT_PRICE_KEY = "bot_slot_price"
+DEFAULT_BOT_SLOT_PRICE = 1
+
+
+class BalanceError(ValueError):
+    pass
+
+
+class InsufficientBalance(BalanceError):
+    pass
+
+
+def _amount(value: int) -> int:
+    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+        raise BalanceError("Amount must be a positive integer")
+    return value
+
+
+async def _get_or_create_user(session, telegram_id: int) -> User:
+    user = await session.scalar(select(User).where(User.telegram_id == telegram_id))
+    if user is None:
+        user = User(telegram_id=telegram_id, balance_credits=0)
+        session.add(user)
+        await session.flush()
+    return user
+
+
+async def get_user_balance(telegram_id: int) -> int:
+    async with async_session() as session:
+        user = await session.scalar(select(User).where(User.telegram_id == telegram_id))
+        return int(user.balance_credits) if user else 0
+
+
+async def change_user_balance(
+    telegram_id: int, amount: int, *, actor_telegram_id: int | None = None
+) -> int:
+    amount = _amount(amount)
+    if (
+        actor_telegram_id is not None
+        and telegram_id != actor_telegram_id
+        and actor_telegram_id != ADMIN_ID
+    ):
+        raise PermissionError("Only an administrator can adjust another user")
+
+    async with async_session() as session:
+        user = await _get_or_create_user(session, telegram_id)
+        await session.execute(
+            update(User).where(User.id == user.id).values(balance_credits=amount)
+        )
+
+        await session.commit()
+        return int(
+            (
+                await session.scalar(
+                    select(User.balance_credits).where(User.id == user.id)
+                )
+            )
+            or 0
+        )
+
+
+async def credit_user(
+    telegram_id: int, amount: int, *, actor_telegram_id: int | None = None
+) -> int:
+    amount = _amount(amount)
+    if (
+        actor_telegram_id is not None
+        and telegram_id != actor_telegram_id
+        and actor_telegram_id != ADMIN_ID
+    ):
+        raise PermissionError("Only an administrator can adjust another user")
+
+    async with async_session() as session:
+        user = await _get_or_create_user(session, telegram_id)
+        await session.execute(
+            update(User)
+            .where(User.id == user.id)
+            .values(balance_credits=User.balance_credits + amount)
+        )
+        await session.commit()
+        return int(
+            (
+                await session.scalar(
+                    select(User.balance_credits).where(User.id == user.id)
+                )
+            )
+            or 0
+        )
+
+
+async def debit_user(
+    telegram_id: int, amount: int, *, actor_telegram_id: int | None = None
+) -> int:
+    amount = _amount(amount)
+    if (
+        actor_telegram_id is not None
+        and telegram_id != actor_telegram_id
+        and actor_telegram_id != ADMIN_ID
+    ):
+        raise PermissionError("Only an administrator can adjust another user")
+
+    async with async_session() as session:
+        user = await _get_or_create_user(session, telegram_id)
+        result = await session.execute(
+            update(User)
+            .where(User.id == user.id, User.balance_credits >= amount)
+            .values(balance_credits=User.balance_credits - amount)
+        )
+        if result.rowcount != 1:
+            await session.rollback()
+            raise InsufficientBalance("Insufficient balance")
+        await session.commit()
+        return int(
+            (
+                await session.scalar(
+                    select(User.balance_credits).where(User.id == user.id)
+                )
+            )
+            or 0
+        )
+
+
+async def get_bot_slot_price() -> int:
+    async with async_session() as session:
+        setting = await session.get(AppSetting, BOT_SLOT_PRICE_KEY)
+        if setting is None:
+            setting = AppSetting(
+                key=BOT_SLOT_PRICE_KEY, integer_value=DEFAULT_BOT_SLOT_PRICE
+            )
+            session.add(setting)
+            await session.commit()
+            return DEFAULT_BOT_SLOT_PRICE
+        if setting.integer_value <= 0:
+            return DEFAULT_BOT_SLOT_PRICE
+        return int(setting.integer_value)
+
+
+async def set_bot_slot_price(price: int, *, updated_by: int) -> int:
+    price = _amount(price)
+    async with async_session() as session:
+        setting = await session.get(AppSetting, BOT_SLOT_PRICE_KEY)
+        if setting is None:
+            setting = AppSetting(
+                key=BOT_SLOT_PRICE_KEY, integer_value=price, updated_by=updated_by
+            )
+            session.add(setting)
+        else:
+            setting.integer_value = price
+            setting.updated_by = updated_by
+        await session.commit()
+        return price
+
+
+async def get_user_profile(telegram_id: int) -> tuple[int, int, int]:
+    async with async_session() as session:
+        user = await session.scalar(select(User).where(User.telegram_id == telegram_id))
+        if user is None:
+            return telegram_id, 0, 0
+
+        count = await session.scalar(
+            select(func.count(BotAccount.id)).where(BotAccount.owner_id == user.id)
+        )
+        return telegram_id, int(user.balance_credits), int(count or 0)
@@ -0,0 +1,82 @@
+from sqlalchemy import select
+
+from bot.db.models import BotAccount, User
+from bot.db.session import async_session
+from bot.config import ADMIN_ID
+import hashlib
+
+
+async def list_user_bot_accounts(telegram_id: int) -> list[BotAccount]:
+    async with async_session() as session:
+        result = await session.execute(
+            select(BotAccount)
+            .join(User)
+            .where(User.telegram_id == telegram_id)
+            .order_by(BotAccount.bot_name)
+        )
+        return list(result.scalars())
+
+
+async def get_owned_bot(telegram_id: int, account_id: int) -> BotAccount | None:
+    async with async_session() as session:
+        result = await session.execute(
+            select(BotAccount)
+            .join(User)
+            .where(BotAccount.id == account_id, User.telegram_id == telegram_id)
+        )
+        return result.scalar_one_or_none()
+
+
+async def create_bot_account(telegram_id: int, **values) -> BotAccount:
+    async with async_session() as session:
+        user = await session.scalar(select(User).where(User.telegram_id == telegram_id))
+        if user is None:
+            user = User(telegram_id=telegram_id)
+            session.add(user)
+            await session.flush()
+        account = BotAccount(owner_id=user.id, **values)
+        session.add(account)
+        await session.commit()
+        await session.refresh(account)
+        return account
+
+
+async def bot_name_exists(bot_name: str) -> bool:
+    async with async_session() as session:
+        return (
+            await session.scalar(
+                select(BotAccount.id).where(BotAccount.bot_name == bot_name)
+            )
+            is not None
+        )
+
+
+async def user_hash_exists(telegram_id: int, config_sha256: str) -> bool:
+    async with async_session() as session:
+        return (
+            await session.scalar(
+                select(BotAccount.id)
+                .join(User)
+                .where(
+                    User.telegram_id == telegram_id,
+                    BotAccount.config_sha256 == config_sha256,
+                )
+            )
+            is not None
+        )
+
+
+async def ensure_admin_compatibility(bot_names: list[str]) -> None:
+    """Keep ASF bots created before ownership was introduced visible to the admin."""
+    for bot_name in bot_names:
+        if await bot_name_exists(bot_name):
+            continue
+
+        await create_bot_account(
+            ADMIN_ID,
+            bot_name=bot_name,
+            source_filename="existing-asf",
+            config_sha256=hashlib.sha256(bot_name.encode()).hexdigest(),
+            upload_state="attached",
+            is_attached=True,
+        )
@@ -7,7 +7,7 @@ from bot.ui.keyboards import inventory_menu_keyboard, inventory_page_keyboard
 logger = get_logger(__name__)


-async def build_inventory_menu(bot_name: str) -> tuple[str, object]:
+async def build_inventory_menu(bot_name: str, account_id: int) -> tuple[str, object]:
    logger.info("Формирование меню инвентаря: bot=%s", bot_name)
    cs2_inventory = await get_bot_inventory(bot_name, 730, 2)
    steam_inventory = await get_bot_inventory(bot_name, 753, 6)
@@ -25,7 +25,7 @@ async def build_inventory_menu(bot_name: str) -> tuple[str, object]:
        logger.info("Инвентарь пуст: bot=%s", bot_name)
        return (
            f"📦 Инвентарь {bot_name} пуст",
-            inventory_menu_keyboard(bot_name, 0, 0),
+            inventory_menu_keyboard(account_id, 0, 0),
        )
    logger.info(
        "Меню инвентаря сформировано: bot=%s cs2_items=%s steam_items=%s",
@@ -40,11 +40,11 @@ async def build_inventory_menu(bot_name: str) -> tuple[str, object]:
        "🔒 — нельзя продавать\n"
        "❌ — нельзя трейдить"
    )
-    return text, inventory_menu_keyboard(bot_name, len(cs2_assets), len(steam_assets))
+    return text, inventory_menu_keyboard(account_id, len(cs2_assets), len(steam_assets))


 async def render_inventory_page(
-    bot_name: str, inventory_type: str, appid: int, contextid: int, page: int
+    bot_name: str, inventory_type: str, appid: int, contextid: int, page: int, account_id: int
 ) -> tuple[str, object] | tuple[None, None]:
    logger.info(
        "Формирование страницы инвентаря: bot=%s type=%s appid=%s contextid=%s page=%s",
@@ -70,7 +70,7 @@ async def render_inventory_page(
        logger.info(
            "Инвентарь выбранного типа пуст: bot=%s type=%s", bot_name, inventory_type
        )
-        return f"Инвентарь {bot_name} пуст", inventory_menu_keyboard(bot_name, 0, 0)
+        return f"Инвентарь {bot_name} пуст", inventory_menu_keyboard(account_id, 0, 0)
    desc_map = {}
    for desc in descriptions:
        key = (str(desc.get("classid")), str(desc.get("instanceid")))
@@ -139,5 +139,5 @@ async def render_inventory_page(
    )

    return text[:4000], inventory_page_keyboard(
-        inventory_type, bot_name, page, total_pages
+        inventory_type, account_id, page, total_pages
    )
@@ -111,11 +111,12 @@ async def redeem_single_bot_keys(bot_name: str, keys: list[str]) -> str:
    return text[:4000] + ("\n\n...обрезано" if len(text) > 4000 else "")


-async def redeem_all_bots_keys(keys: list[str]) -> str:
+async def redeem_all_bots_keys(keys: list[str], allowed_bots: list[str] | None = None) -> str:
    logger.info("Начата активация ключей по всем ботам: count=%s", len(keys))
    response = get_all_bots()
    data = response.json()
-    bots = list(data.get("Result", {}).keys())
+    available = list(data.get("Result", {}).keys())
+    bots = [name for name in available if allowed_bots is None or name in allowed_bots]
    logger.info("Для массовой активации найдено ботов: count=%s", len(bots))
    success_count = 0
    failed_count = 0
@@ -27,7 +27,7 @@ async def stop_twofa_task(message_id: int) -> None:
            )


-async def auto_update_2fa(message, bot_name: str) -> None:
+async def auto_update_2fa(message, bot_name: str, account_id: int) -> None:
    last_code = None
    message_id = message.message_id
    logger.info(
@@ -45,7 +45,7 @@ async def auto_update_2fa(message, bot_name: str) -> None:
                try:
                    await message.edit_text(
                        text,
-                        reply_markup=twofa_keyboard(bot_name, bool(confirmations)),
+                        reply_markup=twofa_keyboard(account_id, bool(confirmations)),
                    )

                    last_code = code