Add admin controls for selected user params and user note enrichment

Agent-Logs-Url: https://github.com/kr0sh512/tt-simple/sessions/34c3052e-7b55-4cb8-9fee-0ee1ba315691

Co-authored-by: kr0sh512 <46405586+kr0sh512@users.noreply.github.com>

README.md

@@ -13,9 +13,11 @@ Async Telegram bot for selling and extending Marzban VPN subscriptions.
   - `data_limit_reset_strategy = month`
   - `proxies = {"vless": {"flow": "xtls-rprx-vision"}}`
   - `inbounds = {"vless": ["VLESS TCP REALITY"]}`
+  - `note` containing Telegram first name, last name, and phone (if available)
 - If user has active subscription, a new purchase extends from current expiration.
 - Bot sends **subscription URL** after successful purchase.
 - Admin notifications for every creation/extension.
+- Per-user price multiplier support (`price_multiplier`) that affects invoice stars (final amount is always integer).
 - `/stats` command for admins:
   - payments count
   - stars earned
@@ -23,6 +25,13 @@ Async Telegram bot for selling and extending Marzban VPN subscriptions.
   - marzban users
   - traffic usage/total quota
 
+## Admin user management
+- `/select_user <username|telegram_id>` — select user for further updates and show current values.
+- `/selected_user` — show the currently selected user and current values.
+- `/set_expire <YYYY-MM-DD or ISO datetime>` — set selected user expiration (UTC).
+- `/set_traffic <GB>` — set selected user monthly traffic limit in GB.
+- `/set_multiplier <float>` — set selected user price multiplier.
+
 ## Setup
 ```bash
 python -m venv .venv

bot.py

@@ -4,7 +4,7 @@ import logging
 import os
 from dataclasses import dataclass
 from datetime import datetime, timezone
-from decimal import Decimal
+from decimal import Decimal, ROUND_HALF_UP
 from typing import Any
 
 import aiohttp
@@ -165,6 +165,7 @@ class VpnBot:
         )
         self.db: aiosqlite.Connection | None = None
         self.locks: dict[int, asyncio.Lock] = {}
+        self.admin_selected_users: dict[int, str] = {}
         self._register_handlers()
 
     async def setup(self) -> None:
@@ -193,10 +194,116 @@ class VpnBot:
     def _register_handlers(self) -> None:
         self.dp.message.register(self.cmd_start, Command("start"))
         self.dp.message.register(self.cmd_stats, Command("stats"))
+        self.dp.message.register(self.cmd_select_user, Command("select_user"))
+        self.dp.message.register(self.cmd_selected_user, Command("selected_user"))
+        self.dp.message.register(self.cmd_set_expire, Command("set_expire"))
+        self.dp.message.register(self.cmd_set_traffic, Command("set_traffic"))
+        self.dp.message.register(self.cmd_set_multiplier, Command("set_multiplier"))
         self.dp.callback_query.register(self.choose_plan, F.data.startswith("buy:"))
         self.dp.pre_checkout_query.register(self.pre_checkout)
         self.dp.message.register(self.successful_payment, F.successful_payment)
 
+    @staticmethod
+    def _extract_price_multiplier(note: str | None) -> float:
+        if not note:
+            return 1.0
+        for line in note.splitlines():
+            if line.startswith("price_multiplier="):
+                try:
+                    value = float(line.split("=", 1)[1].strip())
+                    return value if value > 0 else 1.0
+                except ValueError:
+                    return 1.0
+        return 1.0
+
+    @staticmethod
+    def _set_price_multiplier_note(note: str | None, multiplier: float) -> str:
+        lines = [line for line in (note or "").splitlines() if line and not line.startswith("price_multiplier=")]
+        lines.append(f"price_multiplier={multiplier}")
+        return "\n".join(lines)
+
+    @staticmethod
+    def _format_new_user_note(first_name: str | None, last_name: str | None, phone: str | None) -> str:
+        parts: list[str] = []
+        if first_name:
+            parts.append(f"First name: {first_name}")
+        if last_name:
+            parts.append(f"Last name: {last_name}")
+        if phone:
+            parts.append(f"Phone: {phone}")
+        return "\n".join(parts)
+
+    @staticmethod
+    def _stars_with_multiplier(base_stars: int, multiplier: float) -> int:
+        value = Decimal(base_stars) * Decimal(str(multiplier))
+        rounded = value.quantize(Decimal("1"), rounding=ROUND_HALF_UP)
+        return max(1, int(rounded))
+
+    @staticmethod
+    def _parse_command_arg(message: Message) -> str:
+        text = (message.text or "").strip()
+        parts = text.split(maxsplit=1)
+        return parts[1].strip() if len(parts) > 1 else ""
+
+    @staticmethod
+    def _parse_utc_datetime(value: str) -> datetime:
+        parsed = datetime.fromisoformat(value.strip())
+        if parsed.tzinfo is None:
+            return parsed.replace(tzinfo=timezone.utc)
+        return parsed.astimezone(timezone.utc)
+
+    @staticmethod
+    def _normalize_username(value: str) -> str:
+        cleaned = value.strip()
+        if cleaned.isdigit():
+            return f"tg{cleaned}"
+        return cleaned
+
+    async def _selected_user_or_reply(self, message: Message) -> str | None:
+        username = self.admin_selected_users.get(message.from_user.id)
+        if not username:
+            await message.answer("No selected user. Use /select_user <username> first.")
+            return None
+        return username
+
+    async def _get_admin_user(self, message: Message, username: str) -> dict[str, Any] | None:
+        user = await self.marzban.get_user(username)
+        if not user:
+            await message.answer(f"User <code>{username}</code> not found in Marzban.")
+            return None
+        return user
+
+    async def _update_admin_user(self, username: str, current: dict[str, Any], **changes: Any) -> dict[str, Any]:
+        payload = {
+            "username": current.get("username", username),
+            "status": current.get("status", "active"),
+            "data_limit": int(current.get("data_limit") or 0),
+            "data_limit_reset_strategy": current.get("data_limit_reset_strategy") or "month",
+            "expire": int(current.get("expire") or 0),
+            "proxies": current.get("proxies") or {},
+            "inbounds": current.get("inbounds") or {},
+            "note": current.get("note") or "",
+        }
+        payload.update(changes)
+        return await self.marzban.update_user(username, payload)
+
+    async def _send_selected_user_info(self, message: Message, username: str, user: dict[str, Any]) -> None:
+        expire_ts = int(user.get("expire") or 0)
+        expire_human = (
+            datetime.fromtimestamp(expire_ts, tz=timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
+            if expire_ts
+            else "not set"
+        )
+        traffic = int(user.get("data_limit") or 0)
+        multiplier = self._extract_price_multiplier(user.get("note"))
+        await message.answer(
+            "Selected user:\n"
+            f"Username: <code>{username}</code>\n"
+            f"Expire (UTC): <code>{expire_human}</code>\n"
+            f"Monthly traffic: <b>{bytes_to_gb(traffic)}</b>\n"
+            f"Price multiplier: <b>{multiplier}</b>"
+        )
+
     async def cmd_start(self, message: Message) -> None:
         kb = InlineKeyboardBuilder()
         for months, stars in PLANS.items():
@@ -216,6 +323,10 @@ class VpnBot:
             return
         months = int(query.data.split(":", 1)[1])
         stars = PLANS[months]
+        username = f"tg{query.from_user.id}"
+        current = await self.marzban.get_user(username)
+        multiplier = self._extract_price_multiplier((current or {}).get("note"))
+        stars = self._stars_with_multiplier(stars, multiplier)
         await self.bot.send_invoice(
             chat_id=query.message.chat.id,
             title=f"VPN {months} month(s)",
@@ -247,8 +358,10 @@ class VpnBot:
             return
 
         lock = self.locks.setdefault(message.from_user.id, asyncio.Lock())
+        order_info = getattr(payment, "order_info", None)
+        phone = getattr(order_info, "phone_number", None)
         async with lock:
-            info = await self.buy_or_extend(message.from_user.id, months)
+            info = await self.buy_or_extend(message.from_user, months, phone)
 
         await message.answer(
             "✅ Subscription updated!\n"
@@ -272,8 +385,8 @@ class VpnBot:
             f"Expires: <code>{info['expire_human']}</code>"
         )
 
-    async def buy_or_extend(self, user_id: int, months: int) -> dict[str, str]:
+    async def buy_or_extend(self, user: Any, months: int, phone: str | None = None) -> dict[str, str]:
-        username = f"tg{user_id}"
+        username = f"tg{user.id}"
         now = datetime.now(tz=timezone.utc)
         expire_from = now
         action = "created"
@@ -286,14 +399,22 @@ class VpnBot:
             action = "extended"
 
         expire_at = add_months(expire_from, months)
+        data_limit = int(current.get("data_limit") or 50 * GB) if current else 50 * GB
+        data_limit_reset_strategy = (
+            current.get("data_limit_reset_strategy") or "month"
+        ) if current else "month"
+        note = current.get("note") if current else self._format_new_user_note(
+            user.first_name, user.last_name, phone
+        )
         payload = {
             "username": username,
             "status": "active",
-            "data_limit": 50 * GB,
+            "data_limit": data_limit,
-            "data_limit_reset_strategy": "month",
+            "data_limit_reset_strategy": data_limit_reset_strategy,
             "expire": int(expire_at.timestamp()),
-            "proxies": {"vless": {"flow": "xtls-rprx-vision"}},
+            "proxies": (current.get("proxies") if current else {"vless": {"flow": "xtls-rprx-vision"}}),
-            "inbounds": {"vless": ["VLESS TCP REALITY"]},
+            "inbounds": (current.get("inbounds") if current else {"vless": ["VLESS TCP REALITY"]}),
+            "note": note,
         }
 
         if current:
@@ -369,6 +490,110 @@ class VpnBot:
             f"Traffic quota: <b>{bytes_to_gb(quota)}</b>"
         )
 
+    async def cmd_select_user(self, message: Message) -> None:
+        if message.from_user.id not in self.settings.admin_ids:
+            await message.answer("Admins only")
+            return
+        arg = self._parse_command_arg(message)
+        if not arg:
+            await message.answer("Usage: /select_user <username|telegram_id>")
+            return
+        username = self._normalize_username(arg)
+        user = await self._get_admin_user(message, username)
+        if not user:
+            return
+        self.admin_selected_users[message.from_user.id] = username
+        await self._send_selected_user_info(message, username, user)
+
+    async def cmd_selected_user(self, message: Message) -> None:
+        if message.from_user.id not in self.settings.admin_ids:
+            await message.answer("Admins only")
+            return
+        username = await self._selected_user_or_reply(message)
+        if not username:
+            return
+        user = await self._get_admin_user(message, username)
+        if not user:
+            return
+        await self._send_selected_user_info(message, username, user)
+
+    async def cmd_set_expire(self, message: Message) -> None:
+        if message.from_user.id not in self.settings.admin_ids:
+            await message.answer("Admins only")
+            return
+        arg = self._parse_command_arg(message)
+        if not arg:
+            await message.answer("Usage: /set_expire <YYYY-MM-DD or ISO datetime>")
+            return
+        username = await self._selected_user_or_reply(message)
+        if not username:
+            return
+        try:
+            expire_dt = self._parse_utc_datetime(arg)
+        except ValueError:
+            await message.answer("Invalid datetime format. Example: 2026-12-31 or 2026-12-31T15:30:00")
+            return
+        current = await self._get_admin_user(message, username)
+        if not current:
+            return
+        updated = await self._update_admin_user(username, current, expire=int(expire_dt.timestamp()))
+        await self._send_selected_user_info(message, username, updated)
+
+    async def cmd_set_traffic(self, message: Message) -> None:
+        if message.from_user.id not in self.settings.admin_ids:
+            await message.answer("Admins only")
+            return
+        arg = self._parse_command_arg(message)
+        if not arg:
+            await message.answer("Usage: /set_traffic <GB>")
+            return
+        username = await self._selected_user_or_reply(message)
+        if not username:
+            return
+        try:
+            gb_value = Decimal(arg)
+        except Exception:
+            await message.answer("Invalid value. Example: /set_traffic 75")
+            return
+        if gb_value <= 0:
+            await message.answer("Traffic must be greater than 0.")
+            return
+        data_limit = int((gb_value * Decimal(GB)).quantize(Decimal("1"), rounding=ROUND_HALF_UP))
+        current = await self._get_admin_user(message, username)
+        if not current:
+            return
+        updated = await self._update_admin_user(username, current, data_limit=data_limit, data_limit_reset_strategy="month")
+        await self._send_selected_user_info(message, username, updated)
+
+    async def cmd_set_multiplier(self, message: Message) -> None:
+        if message.from_user.id not in self.settings.admin_ids:
+            await message.answer("Admins only")
+            return
+        arg = self._parse_command_arg(message)
+        if not arg:
+            await message.answer("Usage: /set_multiplier <float>")
+            return
+        username = await self._selected_user_or_reply(message)
+        if not username:
+            return
+        try:
+            multiplier = float(arg)
+        except ValueError:
+            await message.answer("Invalid multiplier. Example: /set_multiplier 1.25")
+            return
+        if multiplier <= 0:
+            await message.answer("Multiplier must be greater than 0.")
+            return
+        current = await self._get_admin_user(message, username)
+        if not current:
+            return
+        updated = await self._update_admin_user(
+            username,
+            current,
+            note=self._set_price_multiplier_note(current.get("note"), multiplier),
+        )
+        await self._send_selected_user_info(message, username, updated)
+
 
 async def main() -> None:
     settings = Settings.from_env()