feat(bot): add proxy limit
This commit is contained in:
@@ -18,7 +18,7 @@ from aiogram.types import (
|
||||
)
|
||||
from convex import ConvexInt64
|
||||
|
||||
from bot.handlers.proxy.handler import get_proxy_config
|
||||
from bot.handlers.proxy.handler import get_proxy_config, increment_proxy_count
|
||||
from bot.modules.ai import (
|
||||
SUMMARIZE_PROMPT,
|
||||
AgentDeps,
|
||||
@@ -426,6 +426,8 @@ async def process_message( # noqa: C901, PLR0912, PLR0913, PLR0915
|
||||
await proxy_config.proxy_bot.send_message(
|
||||
proxy_config.target_chat_id, f"👤 {text}"
|
||||
)
|
||||
await increment_proxy_count(chat_id)
|
||||
proxy_config = get_proxy_config(chat_id)
|
||||
|
||||
rag_connections = await convex.query(
|
||||
"ragConnections:getActiveForUser", {"userId": convex_user_id}
|
||||
@@ -550,6 +552,7 @@ async def process_message( # noqa: C901, PLR0912, PLR0913, PLR0915
|
||||
await proxy_config.proxy_bot.send_message(
|
||||
proxy_config.target_chat_id, part
|
||||
)
|
||||
await increment_proxy_count(chat_id)
|
||||
|
||||
except Exception as e: # noqa: BLE001
|
||||
await state.stop_typing()
|
||||
@@ -663,6 +666,7 @@ async def on_album_message(
|
||||
await proxy_config.proxy_bot.send_media_group(
|
||||
proxy_config.target_chat_id, media
|
||||
)
|
||||
await increment_proxy_count(message.chat.id)
|
||||
|
||||
active_chat_id = user["activeChatId"]
|
||||
await convex.mutation(
|
||||
@@ -732,6 +736,7 @@ async def on_photo_message(message: types.Message, bot: Bot) -> None:
|
||||
BufferedInputFile(image_bytes, "photo.jpg"),
|
||||
caption=f"👤 {caption}",
|
||||
)
|
||||
await increment_proxy_count(message.chat.id)
|
||||
|
||||
active_chat_id = user["activeChatId"]
|
||||
await convex.mutation(
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
from dataclasses import dataclass
|
||||
import contextlib
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from aiogram import Bot, Router, types
|
||||
from aiogram.client.default import DefaultBotProperties
|
||||
@@ -13,6 +14,8 @@ class ProxyConfig:
|
||||
bot_token: str
|
||||
target_chat_id: int
|
||||
proxy_bot: Bot
|
||||
limit: int = 0
|
||||
message_count: int = field(default=0, init=False)
|
||||
|
||||
|
||||
proxy_states: dict[int, ProxyConfig] = {}
|
||||
@@ -28,25 +31,43 @@ async def cleanup_proxy(chat_id: int) -> None:
|
||||
await config.proxy_bot.session.close()
|
||||
|
||||
|
||||
async def increment_proxy_count(chat_id: int) -> bool:
|
||||
config = proxy_states.get(chat_id)
|
||||
if not config:
|
||||
return False
|
||||
|
||||
config.message_count += 1
|
||||
|
||||
if config.limit > 0 and config.message_count >= config.limit:
|
||||
await config.proxy_bot.send_message(
|
||||
config.target_chat_id, f"🔗 Proxy finished ({config.limit} messages)"
|
||||
)
|
||||
await cleanup_proxy(chat_id)
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
@router.message(Command("proxy"))
|
||||
async def on_proxy(message: types.Message) -> None: # noqa: C901
|
||||
async def on_proxy(message: types.Message) -> None: # noqa: C901, PLR0912
|
||||
if not message.from_user or not message.text:
|
||||
return
|
||||
|
||||
chat_id = message.chat.id
|
||||
args = message.text.split(maxsplit=2)
|
||||
args = message.text.split(maxsplit=3)
|
||||
|
||||
if len(args) == 1:
|
||||
config = get_proxy_config(chat_id)
|
||||
if config:
|
||||
cnt = f" ({config.message_count}/{config.limit})" if config.limit else ""
|
||||
await message.answer(
|
||||
f"Proxy active → chat {config.target_chat_id}\n\n"
|
||||
f"Proxy active → chat {config.target_chat_id}{cnt}\n\n"
|
||||
"Use /proxy deactivate to stop."
|
||||
)
|
||||
else:
|
||||
await message.answer(
|
||||
"Usage:\n"
|
||||
"/proxy BOT_TOKEN CHAT_ID - activate proxy\n"
|
||||
"/proxy BOT_TOKEN CHAT_ID [LIMIT] - activate proxy\n"
|
||||
"/proxy deactivate - stop proxy"
|
||||
)
|
||||
return
|
||||
@@ -60,7 +81,7 @@ async def on_proxy(message: types.Message) -> None: # noqa: C901
|
||||
return
|
||||
|
||||
if len(args) < 3: # noqa: PLR2004
|
||||
await message.answer("Usage: /proxy BOT_TOKEN CHAT_ID")
|
||||
await message.answer("Usage: /proxy BOT_TOKEN CHAT_ID [LIMIT]")
|
||||
return
|
||||
|
||||
bot_token = args[1]
|
||||
@@ -70,6 +91,11 @@ async def on_proxy(message: types.Message) -> None: # noqa: C901
|
||||
await message.answer("Invalid chat ID. Must be a number.")
|
||||
return
|
||||
|
||||
limit = 0
|
||||
if len(args) >= 4: # noqa: PLR2004
|
||||
with contextlib.suppress(ValueError):
|
||||
limit = max(int(args[3]), 0)
|
||||
|
||||
if chat_id in proxy_states:
|
||||
await cleanup_proxy(chat_id)
|
||||
|
||||
@@ -87,12 +113,16 @@ async def on_proxy(message: types.Message) -> None: # noqa: C901
|
||||
return
|
||||
|
||||
proxy_states[chat_id] = ProxyConfig(
|
||||
bot_token=bot_token, target_chat_id=target_chat_id, proxy_bot=proxy_bot
|
||||
bot_token=bot_token,
|
||||
target_chat_id=target_chat_id,
|
||||
proxy_bot=proxy_bot,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
limit_text = f"\nLimit: {limit} messages" if limit > 0 else ""
|
||||
await message.answer(
|
||||
f"✓ Proxy activated via @{bot_info.username}\n"
|
||||
f"Target: {target_chat_id}\n\n"
|
||||
f"Target: {target_chat_id}{limit_text}\n\n"
|
||||
"All messages will be forwarded."
|
||||
)
|
||||
await message.delete()
|
||||
|
||||
Reference in New Issue
Block a user