Files
stealth-ai-relay/backend/src/bot/handlers/proxy/handler.py
2026-01-29 16:28:49 +01:00

132 lines
3.7 KiB
Python

import contextlib
from dataclasses import dataclass, field
from aiogram import Bot, Router, types
from aiogram.client.default import DefaultBotProperties
from aiogram.enums import ParseMode
from aiogram.filters import Command
router = Router()
@dataclass
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] = {}
def get_proxy_config(chat_id: int) -> ProxyConfig | None:
return proxy_states.get(chat_id)
async def cleanup_proxy(chat_id: int) -> None:
config = proxy_states.pop(chat_id, None)
if config:
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, PLR0912
if not message.from_user or not message.text:
return
chat_id = message.chat.id
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}{cnt}\n\n"
"Use /proxy deactivate to stop."
)
else:
await message.answer(
"Usage:\n"
"/proxy BOT_TOKEN CHAT_ID [LIMIT] - activate proxy\n"
"/proxy deactivate - stop proxy"
)
return
if args[1].lower() == "deactivate":
if chat_id in proxy_states:
await cleanup_proxy(chat_id)
await message.answer("✓ Proxy deactivated.")
else:
await message.answer("No active proxy.")
return
if len(args) < 3: # noqa: PLR2004
await message.answer("Usage: /proxy BOT_TOKEN CHAT_ID [LIMIT]")
return
bot_token = args[1]
try:
target_chat_id = int(args[2])
except ValueError:
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)
try:
proxy_bot = Bot(
token=bot_token, default=DefaultBotProperties(parse_mode=ParseMode.HTML)
)
bot_info = await proxy_bot.get_me()
try:
await proxy_bot.send_message(target_chat_id, "🔗 Proxy connected")
except Exception as e: # noqa: BLE001
await proxy_bot.session.close()
await message.answer(f"Cannot send to chat {target_chat_id}: {e}")
return
proxy_states[chat_id] = ProxyConfig(
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}{limit_text}\n\n"
"All messages will be forwarded."
)
await message.delete()
except Exception as e: # noqa: BLE001
await message.answer(f"Invalid bot token: {e}")