feat(bot): add proxy limit

This commit is contained in:
h
2026-01-29 16:28:49 +01:00
parent 0e14e945bd
commit 504282cb0f
2 changed files with 44 additions and 9 deletions

View File

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