281 lines
12 KiB
Python
281 lines
12 KiB
Python
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass
|
||
from typing import Any
|
||
|
||
from aiogram import Bot, Dispatcher, F, Router
|
||
from aiogram.filters import CommandStart
|
||
from aiogram.types import CallbackQuery, Message
|
||
|
||
from .config import Config
|
||
from .db import Database
|
||
from .keyboards import admin_menu, extend_keyboard, main_menu, reset_confirm_keyboard
|
||
from .marzban import MarzbanClient, MarzbanError
|
||
from .messages import format_vpn_info, generate_invite_token, payment_amount, user_display_name
|
||
from .utils import build_marzban_note, format_marzban_username, format_timestamp, now_ts
|
||
|
||
|
||
@dataclass
|
||
class PendingPayment:
|
||
months: int
|
||
amount: int
|
||
|
||
|
||
class BotApp:
|
||
def __init__(self, *, config: Config, db: Database, marzban: MarzbanClient):
|
||
self.config = config
|
||
self.db = db
|
||
self.marzban = marzban
|
||
self.router = Router()
|
||
self.pending_payments: dict[int, PendingPayment] = {}
|
||
self._register_handlers()
|
||
|
||
def is_admin(self, user_id: int | None) -> bool:
|
||
return user_id is not None and user_id in self.config.admin_ids
|
||
|
||
def _register_handlers(self) -> None:
|
||
self.router.message(CommandStart())(self.start)
|
||
self.router.message(F.text == "➕ New invite")(self.new_invite)
|
||
self.router.message(F.text == "/new_user")(self.new_invite)
|
||
self.router.message(F.text == "📄 VPN info")(self.vpn_info)
|
||
self.router.message(F.text == "💳 Extend VPN")(self.extend_menu)
|
||
self.router.message(F.text == "🔄 Reset traffic")(self.reset_explain)
|
||
self.router.message(F.contact)(self.save_contact)
|
||
self.router.callback_query(F.data.startswith("extend:"))(self.select_extend)
|
||
self.router.callback_query(F.data == "reset:confirm")(self.confirm_reset)
|
||
self.router.callback_query(F.data == "reset:cancel")(self.cancel_reset)
|
||
self.router.message(F.photo | F.document)(self.payment_proof)
|
||
|
||
async def start(self, message: Message) -> None:
|
||
user_id = message.from_user.id if message.from_user else None
|
||
args = message.text.split(maxsplit=1)[1].strip() if message.text and " " in message.text else ""
|
||
|
||
if args:
|
||
await self.activate_invite(message, args)
|
||
return
|
||
|
||
if self.is_admin(user_id):
|
||
await message.answer("Admin menu", reply_markup=admin_menu())
|
||
return
|
||
|
||
if user_id and self.db.get_by_tg_user_id(user_id):
|
||
await message.answer("Welcome back.", reply_markup=main_menu())
|
||
return
|
||
|
||
await message.answer("Ask an admin for an invite link.")
|
||
|
||
async def new_invite(self, message: Message) -> None:
|
||
if not self.is_admin(message.from_user.id if message.from_user else None):
|
||
await message.answer("Only admins can create invite links.")
|
||
return
|
||
bot_user = await message.bot.get_me()
|
||
token = generate_invite_token()
|
||
self.db.create_invite(token=token, admin_id=message.from_user.id, created_at=now_ts())
|
||
link = f"https://t.me/{bot_user.username}?start={token}"
|
||
text = f"Invite link created:\n{link}"
|
||
await message.answer(text)
|
||
await self.notify_admins(message.bot, f"➕ New invite created by admin {message.from_user.id}\n{link}")
|
||
|
||
async def activate_invite(self, message: Message, token: str) -> None:
|
||
tg = message.from_user
|
||
if tg is None:
|
||
return
|
||
existing = self.db.get_by_tg_user_id(tg.id)
|
||
if existing:
|
||
await message.answer("You are already activated.", reply_markup=main_menu())
|
||
return
|
||
invite = self.db.get_by_token(token)
|
||
if invite is None:
|
||
await message.answer("Invalid invite link.")
|
||
return
|
||
if invite["tg_user_id"] is not None:
|
||
await message.answer("This invite link was already used.")
|
||
return
|
||
|
||
marzban_username = format_marzban_username(tg.id)
|
||
try:
|
||
await self.marzban.create_trial_user(
|
||
tg_user_id=tg.id,
|
||
first_name=tg.first_name,
|
||
last_name=tg.last_name,
|
||
tg_username=tg.username,
|
||
)
|
||
self.db.activate_invite(
|
||
token=token,
|
||
tg_user_id=tg.id,
|
||
tg_username=tg.username,
|
||
tg_first_name=tg.first_name,
|
||
tg_last_name=tg.last_name,
|
||
tg_phone=None,
|
||
marzban_username=marzban_username,
|
||
activated_at=now_ts(),
|
||
)
|
||
user = await self.marzban.get_user(marzban_username)
|
||
except (MarzbanError, ValueError) as exc:
|
||
await message.answer(f"Activation failed: {exc}")
|
||
await self.notify_admins(message.bot, f"⚠️ Activation failed for {tg.id}: {exc}")
|
||
return
|
||
|
||
await message.answer("VPN trial activated for 3 days / 10 GB.", reply_markup=main_menu())
|
||
await message.answer(format_vpn_info(user))
|
||
await self.notify_admins(
|
||
message.bot,
|
||
f"✅ User activated invite\nTelegram: {tg.first_name or ''} {tg.last_name or ''} @{tg.username or '-'}\nMarzban: {marzban_username}",
|
||
)
|
||
|
||
async def get_current_local_user(self, message_or_query: Message | CallbackQuery) -> dict[str, Any] | None:
|
||
from_user = message_or_query.from_user
|
||
if from_user is None:
|
||
return None
|
||
return self.db.get_by_tg_user_id(from_user.id)
|
||
|
||
async def vpn_info(self, message: Message) -> None:
|
||
row = await self.get_current_local_user(message)
|
||
if not row or not row.get("marzban_username"):
|
||
await message.answer("You are not activated. Use an invite link first.")
|
||
return
|
||
try:
|
||
user = await self.marzban.get_user(row["marzban_username"])
|
||
except MarzbanError as exc:
|
||
await message.answer(f"Could not fetch VPN info: {exc}")
|
||
return
|
||
await message.answer(format_vpn_info(user))
|
||
|
||
async def extend_menu(self, message: Message) -> None:
|
||
row = await self.get_current_local_user(message)
|
||
if not row:
|
||
await message.answer("You are not activated. Use an invite link first.")
|
||
return
|
||
await message.answer("Choose extension period:", reply_markup=extend_keyboard())
|
||
|
||
async def select_extend(self, query: CallbackQuery) -> None:
|
||
if query.from_user is None:
|
||
return
|
||
row = self.db.get_by_tg_user_id(query.from_user.id)
|
||
if not row:
|
||
await query.message.answer("You are not activated. Use an invite link first.")
|
||
await query.answer()
|
||
return
|
||
months = int(query.data.split(":", 1)[1])
|
||
amount = payment_amount(months)
|
||
self.pending_payments[query.from_user.id] = PendingPayment(months=months, amount=amount)
|
||
await query.message.answer(
|
||
f"Extension selected: {months} month(s), {amount} ₽.\n\n"
|
||
f"{self.config.payment_text}\n\n"
|
||
"After payment, send a screenshot/photo or file here."
|
||
)
|
||
await query.answer()
|
||
|
||
async def payment_proof(self, message: Message) -> None:
|
||
if message.from_user is None:
|
||
return
|
||
pending = self.pending_payments.get(message.from_user.id)
|
||
if pending is None:
|
||
return
|
||
row = self.db.get_by_tg_user_id(message.from_user.id)
|
||
if not row or not row.get("marzban_username"):
|
||
await message.answer("You are not activated. Use an invite link first.")
|
||
return
|
||
try:
|
||
user = await self.marzban.extend_user(row["marzban_username"], pending.months)
|
||
except MarzbanError as exc:
|
||
await message.answer(f"Could not extend VPN: {exc}")
|
||
await self.notify_admins(message.bot, f"⚠️ Extension failed for {row['marzban_username']}: {exc}")
|
||
return
|
||
self.pending_payments.pop(message.from_user.id, None)
|
||
await message.answer("Payment proof received. VPN was extended automatically.")
|
||
await message.answer(format_vpn_info(user))
|
||
caption = (
|
||
"💳 Payment proof received and VPN extended\n"
|
||
f"User: {user_display_name(row)}\n"
|
||
f"Marzban: {row['marzban_username']}\n"
|
||
f"Period: {pending.months} month(s)\n"
|
||
f"Amount: {pending.amount} ₽\n"
|
||
f"New expiration: {format_timestamp(user.get('expire'))}"
|
||
)
|
||
await self.forward_payment_to_admins(message, caption)
|
||
|
||
async def reset_explain(self, message: Message) -> None:
|
||
row = await self.get_current_local_user(message)
|
||
if not row:
|
||
await message.answer("You are not activated. Use an invite link first.")
|
||
return
|
||
await message.answer(
|
||
"Reset traffic will restore your traffic limit to 50 GB.\n\n"
|
||
"Important: all days of your current partial paid month will be removed. "
|
||
"Your remaining VPN time will be rounded down to full 30-day months.\n\n"
|
||
"Example: if you have 1 month and 12 days left, after reset you will have exactly 1 month left.\n\n"
|
||
"Do you want to continue?",
|
||
reply_markup=reset_confirm_keyboard(),
|
||
)
|
||
|
||
async def confirm_reset(self, query: CallbackQuery) -> None:
|
||
row = self.db.get_by_tg_user_id(query.from_user.id)
|
||
if not row or not row.get("marzban_username"):
|
||
await query.message.answer("You are not activated. Use an invite link first.")
|
||
await query.answer()
|
||
return
|
||
try:
|
||
user, full_months, _new_expire = await self.marzban.reset_paid_traffic_with_time_penalty(row["marzban_username"])
|
||
except MarzbanError as exc:
|
||
await query.message.answer(f"Could not reset traffic: {exc}")
|
||
await query.answer()
|
||
return
|
||
await query.message.answer(
|
||
f"Traffic reset complete. Remaining time was rounded down to {full_months} full month(s)."
|
||
)
|
||
await query.message.answer(format_vpn_info(user))
|
||
await query.answer()
|
||
|
||
async def cancel_reset(self, query: CallbackQuery) -> None:
|
||
await query.message.answer("Reset cancelled.")
|
||
await query.answer()
|
||
|
||
async def save_contact(self, message: Message) -> None:
|
||
if not message.from_user or not message.contact:
|
||
return
|
||
if message.contact.user_id and message.contact.user_id != message.from_user.id:
|
||
await message.answer("Please share your own contact, not another contact.")
|
||
return
|
||
phone = message.contact.phone_number
|
||
self.db.update_phone(tg_user_id=message.from_user.id, phone=phone)
|
||
row = self.db.get_by_tg_user_id(message.from_user.id)
|
||
if row and row.get("marzban_username"):
|
||
note = build_marzban_note(
|
||
first_name=row.get("tg_first_name"),
|
||
last_name=row.get("tg_last_name"),
|
||
username=row.get("tg_username"),
|
||
phone=phone,
|
||
)
|
||
try:
|
||
await self.marzban.modify_user(row["marzban_username"], {"note": note})
|
||
except MarzbanError:
|
||
await message.answer("Phone saved locally, but I could not update Marzban note now.")
|
||
return
|
||
await message.answer("Phone saved.", reply_markup=main_menu())
|
||
|
||
async def notify_admins(self, bot: Bot, text: str) -> None:
|
||
for admin_id in self.config.admin_ids:
|
||
try:
|
||
await bot.send_message(admin_id, text)
|
||
except Exception:
|
||
pass
|
||
|
||
async def forward_payment_to_admins(self, message: Message, caption: str) -> None:
|
||
for admin_id in self.config.admin_ids:
|
||
try:
|
||
if message.photo:
|
||
await message.bot.send_photo(admin_id, message.photo[-1].file_id, caption=caption)
|
||
elif message.document:
|
||
await message.bot.send_document(admin_id, message.document.file_id, caption=caption)
|
||
else:
|
||
await message.bot.send_message(admin_id, caption)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def build_dispatcher(app: BotApp) -> Dispatcher:
|
||
dp = Dispatcher()
|
||
dp.include_router(app.router)
|
||
return dp
|