Files
remnawave-minishop/bot/services/tariff_worker.py
T

669 lines
27 KiB
Python

import asyncio
import logging
from datetime import datetime, timezone
from typing import Optional
from aiogram import Bot
from aiogram.types import InlineKeyboardButton, InlineKeyboardMarkup, WebAppInfo
from aiogram.utils.text_decorations import html_decoration as hd
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import sessionmaker
from bot.middlewares.i18n import JsonI18n
from bot.services.panel_api_service import PanelApiService
from bot.services.subscription_service import SubscriptionService
from bot.utils.date_utils import month_start
from bot.utils.mini_app_url import subscription_mini_app_topup_url
from config.settings import Settings
from db.dal import subscription_dal, tariff_dal, user_dal
from db.models import Subscription
PREMIUM_WARNING_LEVEL_OFFSET = 1000
# Single warning per premium billing period when usage reached or exceeded the quota.
PREMIUM_WARNING_DEPLETED_LEVEL = PREMIUM_WARNING_LEVEL_OFFSET + 100
class TariffTrafficWorker:
def __init__(
self,
settings: Settings,
session_factory: sessionmaker,
panel_service: PanelApiService,
subscription_service: SubscriptionService,
bot: Optional[Bot] = None,
i18n: Optional[JsonI18n] = None,
):
self.settings = settings
self.session_factory = session_factory
self.panel_service = panel_service
self.subscription_service = subscription_service
self.bot = bot
self.i18n = i18n
self._stopped = asyncio.Event()
self._premium_nodes_cache = {}
async def _user_lang(self, session: AsyncSession, user_id: int) -> str:
try:
row = await user_dal.get_user_by_id(session, user_id)
if row and getattr(row, "language_code", None):
code = str(row.language_code or "").strip()
if code:
return code
except Exception:
logging.exception("TariffTrafficWorker: failed to load user language for %s", user_id)
return self.settings.DEFAULT_LANGUAGE
def _usage_placeholders(self, used_bytes: int, limit_bytes: int) -> dict:
"""Formatted traffic stats for warning messages (HTML-safe quoted)."""
used_b = max(0, int(used_bytes or 0))
lim_b = max(0, int(limit_bytes or 0))
remaining_b = max(0, lim_b - used_b)
return {
"used": hd.quote(self._fmt_bytes(used_b)),
"remaining": hd.quote(self._fmt_bytes(remaining_b)),
"limit_total": hd.quote(self._fmt_bytes(lim_b)),
}
def _traffic_topup_markup(self, user_lang: str, kind: str) -> Optional[InlineKeyboardMarkup]:
if not self.bot:
return None
_ = lambda k, **kw: (
self.i18n.gettext(user_lang, k, **kw) if self.i18n else (lambda key, **_: key)
)
normalized = "premium" if str(kind or "").lower() == "premium" else "regular"
url = subscription_mini_app_topup_url(self.settings, normalized)
if normalized == "premium":
label_key = "traffic_warn_btn_topup_webapp_premium"
fallback_key = "traffic_warn_btn_topup_premium"
else:
label_key = "traffic_warn_btn_topup_webapp_regular"
fallback_key = "traffic_warn_btn_topup_regular"
# Mini App inside Telegram when SUBSCRIPTION_MINI_APP_URL is configured.
if url:
button = InlineKeyboardButton(text=_(label_key), web_app=WebAppInfo(url=url))
else:
button = InlineKeyboardButton(text=_(fallback_key), callback_data="tariff_topup:list")
return InlineKeyboardMarkup(inline_keyboard=[[button]])
async def run(self) -> None:
if not self.settings.tariffs_config:
return
while not self._stopped.is_set():
try:
async with self.session_factory() as session:
await self.traffic_period_tick(session)
await session.commit()
async with self.session_factory() as session:
await self.legacy_throttle_recovery_tick(session)
await session.commit()
except Exception:
logging.exception("TariffTrafficWorker tick failed")
try:
await asyncio.wait_for(self._stopped.wait(), timeout=300)
except asyncio.TimeoutError:
pass
def stop(self) -> None:
self._stopped.set()
async def traffic_period_tick(self, session: AsyncSession) -> None:
now = datetime.now(timezone.utc)
warning_period_start = month_start(now)
result = await session.execute(
select(Subscription).where(
Subscription.is_active == True,
Subscription.end_date > now,
Subscription.tariff_key.is_not(None),
)
)
for sub in result.scalars().all():
try:
tariff = self.settings.tariffs_config.require(sub.tariff_key)
except Exception:
continue
panel_data = (
await self.panel_service.get_user_by_uuid(sub.panel_user_uuid, log_response=False)
or {}
)
used, limit, panel_strategy = self.subscription_service._extract_panel_traffic_details(
panel_data
)
panel_status = str(panel_data.get("status") or "").upper()
panel_username = panel_data.get("username") if isinstance(panel_data, dict) else None
if used is not None and used != sub.traffic_used_bytes:
sub.traffic_used_bytes = used
if limit is not None and limit != sub.traffic_limit_bytes:
sub.traffic_limit_bytes = limit
if panel_status and panel_status != (sub.status_from_panel or "").upper():
sub.status_from_panel = panel_status
if tariff.billing_model == "period":
await self._ensure_period_reset_strategy(sub, tariff, limit, panel_strategy)
await self._maybe_warn_or_throttle(
session,
sub,
tariff,
used,
limit,
warning_period_start=warning_period_start
if tariff.billing_model == "period"
else None,
)
await self._sync_premium_squad_limit(
session,
sub,
tariff,
now,
panel_username=panel_username,
panel_user_dict=panel_data,
)
async def _ensure_period_reset_strategy(
self,
sub: Subscription,
tariff,
limit: Optional[int],
panel_strategy: Optional[str],
) -> None:
if str(panel_strategy or "").upper() == "MONTH":
return
rb = int(getattr(sub, "regular_bonus_bytes", 0) or 0)
if bool(getattr(sub, "regular_unlimited_override", False)):
baseline = int(sub.tier_baseline_bytes or (tariff.monthly_bytes if tariff else 0) or 0)
traffic_limit_bytes = self.subscription_service._compute_main_traffic_limit_bytes(
tier_baseline_bytes=baseline,
topup_balance_bytes=int(sub.topup_balance_bytes or 0),
regular_bonus_bytes=rb,
regular_unlimited_override=True,
traffic_used_bytes=int(sub.traffic_used_bytes or 0),
)
else:
traffic_limit_bytes = int(
limit
or sub.traffic_limit_bytes
or (tariff.monthly_bytes + int(sub.topup_balance_bytes or 0) + rb)
)
payload = self.subscription_service._build_panel_update_payload(
panel_user_uuid=sub.panel_user_uuid,
expire_at=sub.end_date,
traffic_limit_bytes=traffic_limit_bytes,
traffic_limit_strategy="MONTH",
)
payload["activeInternalSquads"] = self.subscription_service._panel_squads_for_tariff(
tariff,
include_premium=not bool(getattr(sub, "premium_is_limited", False)),
)
await self.panel_service.update_user_details_on_panel(
sub.panel_user_uuid, payload, log_response=False
)
async def _maybe_warn_or_throttle(
self,
session: AsyncSession,
sub: Subscription,
tariff,
used: Optional[int],
limit: Optional[int],
*,
warning_period_start: Optional[datetime] = None,
) -> None:
if bool(getattr(sub, "regular_unlimited_override", False)):
return
used_val = int(used or sub.traffic_used_bytes or 0)
limit_val = int(limit or sub.traffic_limit_bytes or 0)
if limit_val <= 0:
return
ratio = used_val / limit_val
levels = list(getattr(self.settings, "tariff_traffic_warning_levels", [85, 90, 95]))
for level in levels:
threshold = level / 100
if ratio < threshold:
continue
warning = await tariff_dal.get_warning(
session,
subscription_id=sub.subscription_id,
period_start_at=warning_period_start if tariff.billing_model == "period" else None,
level=level,
traffic_limit_bytes=limit_val if tariff.billing_model == "traffic" else None,
)
if warning:
continue
await tariff_dal.create_warning(
session,
subscription_id=sub.subscription_id,
period_start_at=warning_period_start if tariff.billing_model == "period" else None,
level=level,
traffic_limit_bytes=limit_val if tariff.billing_model == "traffic" else None,
)
if self.bot:
try:
user_lang = await self._user_lang(session, sub.user_id)
_ = (
(lambda k, **kw: self.i18n.gettext(user_lang, k, **kw))
if self.i18n
else (lambda k, **kw: k)
)
left_pct = max(0, 100 - level)
tariff_name = hd.quote(str(tariff.name(user_lang)))
usage = self._usage_placeholders(used_val, limit_val)
if level < 100:
text = _(
"traffic_warning_regular_almost",
tariff_name=tariff_name,
left_pct=left_pct,
**usage,
)
else:
text = _(
"traffic_warning_regular_depleted",
tariff_name=tariff_name,
**usage,
)
markup = self._traffic_topup_markup(user_lang, "regular")
await self.bot.send_message(
sub.user_id,
text,
reply_markup=markup,
parse_mode="HTML",
)
except Exception:
logging.exception("Failed to send traffic warning to user %s", sub.user_id)
if ratio >= 1.0 and not sub.is_throttled:
logging.info(
"Tariff traffic limit reached for user %s subscription %s. "
"Leaving access control to Remnawave status handling.",
sub.user_id,
sub.subscription_id,
)
async def _sync_premium_squad_limit(
self,
session: AsyncSession,
sub: Subscription,
tariff,
now: datetime,
*,
panel_username: Optional[str] = None,
panel_user_dict: Optional[dict] = None,
) -> None:
if not getattr(tariff, "premium_squad_uuids", None):
if (
any(
int(value or 0) > 0
for value in (
sub.premium_baseline_bytes,
sub.premium_topup_balance_bytes,
sub.premium_used_bytes,
)
)
or sub.premium_is_limited
):
sub.premium_baseline_bytes = 0
sub.premium_topup_balance_bytes = 0
sub.premium_used_bytes = 0
sub.premium_is_limited = False
return
premium_period_start = month_start(now)
same_period = bool(getattr(sub, "premium_period_start_at", None) == premium_period_start)
premium_baseline = int(tariff.premium_monthly_bytes or 0)
premium_topup_balance = int(sub.premium_topup_balance_bytes or 0)
premium_topup_used = (
int(getattr(sub, "premium_topup_used_bytes", 0) or 0) if same_period else 0
)
# Admin-side overrides for free gifted premium traffic.
premium_unlimited_override = bool(getattr(sub, "premium_unlimited_override", False))
premium_bonus = max(0, int(getattr(sub, "premium_bonus_bytes", 0) or 0))
premium_limit = (
premium_baseline + premium_topup_balance + premium_topup_used + premium_bonus
)
if premium_limit <= 0 and not premium_unlimited_override:
return
node_uuids = await self._premium_node_uuids_for_tariff(tariff)
if not node_uuids:
logging.warning("Premium squads for tariff %s have no accessible nodes", tariff.key)
return
start_date = now.date().replace(day=1).isoformat()
end_date = now.date().isoformat()
premium_used = await self._premium_usage_for_user(
sub.panel_user_uuid,
node_uuids,
start_date,
end_date,
panel_username=panel_username,
)
if premium_used is None:
return
# Consume paid top-up balance only for overflow beyond baseline+bonus.
# Admin-granted bonus is "spent" against usage first along with baseline,
# so the user's paid top-up survives longer.
free_quota = premium_baseline + premium_bonus
overflow = max(0, int(premium_used) - free_quota)
delta_overflow = max(0, overflow - premium_topup_used)
consume_from_topup = min(premium_topup_balance, delta_overflow)
if consume_from_topup > 0:
premium_topup_balance -= consume_from_topup
premium_topup_used += consume_from_topup
premium_limit = (
premium_baseline + premium_topup_balance + premium_topup_used + premium_bonus
)
if premium_unlimited_override:
should_limit = False
else:
should_limit = premium_used >= premium_limit
panel_needs_update = bool(sub.premium_is_limited) != should_limit
desired_squads = self.subscription_service._panel_squads_for_tariff(
tariff,
include_premium=not should_limit,
)
desired_set = self._internal_squad_uuid_set(desired_squads)
if isinstance(panel_user_dict, dict):
current_known = False
current_raw = None
for key in ("activeInternalSquads", "active_internal_squads"):
if key in panel_user_dict:
current_raw = panel_user_dict.get(key)
current_known = True
break
if current_known and desired_set != self._internal_squad_uuid_set(current_raw):
panel_needs_update = True
sub.premium_baseline_bytes = premium_baseline
sub.premium_topup_balance_bytes = premium_topup_balance
sub.premium_topup_used_bytes = premium_topup_used
sub.premium_used_bytes = int(premium_used)
sub.premium_is_limited = bool(should_limit)
sub.premium_period_start_at = premium_period_start
if not premium_unlimited_override:
await self._maybe_warn_premium_squad_limit(
session,
sub,
tariff,
premium_used,
premium_limit,
premium_period_start,
)
if not panel_needs_update:
return
squads = desired_squads
await self.panel_service.update_user_details_on_panel(
sub.panel_user_uuid,
{"uuid": sub.panel_user_uuid, "activeInternalSquads": squads},
log_response=False,
)
logging.info(
"Premium squad access %s for user %s tariff %s: %s/%s bytes",
"limited" if should_limit else "restored",
sub.user_id,
tariff.key,
premium_used,
premium_limit,
)
@staticmethod
def _internal_squad_uuid_set(raw) -> set[str]:
if not isinstance(raw, list):
return set()
out: set[str] = set()
for item in raw:
if isinstance(item, dict):
u = item.get("uuid") or item.get("internalSquadUuid") or item.get("squadUuid")
if u:
out.add(str(u))
elif item:
out.add(str(item))
return out
@staticmethod
def _fmt_bytes(value: int) -> str:
size = float(max(0, int(value or 0)))
for unit in ("B", "KB", "MB", "GB", "TB"):
if size < 1024 or unit == "TB":
return f"{size:.1f} {unit}" if unit != "B" else f"{int(size)} B"
size /= 1024
return f"{size:.1f} TB"
async def _maybe_warn_premium_squad_limit(
self,
session: AsyncSession,
sub: Subscription,
tariff,
used: int,
limit: int,
period_start_at: datetime,
) -> None:
if limit <= 0:
return
used_val = int(used or 0)
limit_val = int(limit)
ratio = used_val / limit_val
levels = list(getattr(self.settings, "tariff_traffic_warning_levels", [85, 90, 95]))
# Fully exhausted or over quota — one message per period (same idea as regular traffic at 100%).
if ratio >= 1.0:
depleted_existing = await tariff_dal.get_warning(
session,
subscription_id=sub.subscription_id,
period_start_at=period_start_at,
level=PREMIUM_WARNING_DEPLETED_LEVEL,
)
if depleted_existing:
return
await tariff_dal.create_warning(
session,
subscription_id=sub.subscription_id,
period_start_at=period_start_at,
level=PREMIUM_WARNING_DEPLETED_LEVEL,
traffic_limit_bytes=None,
)
if self.bot:
try:
user_lang = await self._user_lang(session, sub.user_id)
_ = (
(lambda k, **kw: self.i18n.gettext(user_lang, k, **kw))
if self.i18n
else (lambda k, **kw: k)
)
access = await self.subscription_service.premium_access_for_tariff(tariff)
labels = access.get("node_labels") or access.get("squad_labels") or []
if labels:
visible = [hd.quote(str(x)) for x in labels[:8]]
servers = "\n".join(f"• {label}" for label in visible)
if len(labels) > len(visible):
more = len(labels) - len(visible)
servers += "\n" + _("traffic_warning_premium_servers_more", count=more)
else:
servers = _("traffic_warning_premium_generic_servers")
usage = self._usage_placeholders(used_val, limit_val)
text = _(
"traffic_warning_premium_depleted",
tariff_name=hd.quote(str(tariff.name(user_lang))),
servers=servers,
**usage,
)
markup = self._traffic_topup_markup(user_lang, "premium")
await self.bot.send_message(
sub.user_id,
text,
reply_markup=markup,
parse_mode="HTML",
)
except Exception:
logging.exception(
"Failed to send premium traffic depleted warning to user %s", sub.user_id
)
return
for level in levels:
if level >= 100:
continue
if ratio < level / 100:
continue
storage_level = PREMIUM_WARNING_LEVEL_OFFSET + int(level)
warning = await tariff_dal.get_warning(
session,
subscription_id=sub.subscription_id,
period_start_at=period_start_at,
level=storage_level,
)
if warning:
continue
await tariff_dal.create_warning(
session,
subscription_id=sub.subscription_id,
period_start_at=period_start_at,
level=storage_level,
traffic_limit_bytes=None,
)
if not self.bot:
continue
try:
user_lang = await self._user_lang(session, sub.user_id)
_ = (
(lambda k, **kw: self.i18n.gettext(user_lang, k, **kw))
if self.i18n
else (lambda k, **kw: k)
)
access = await self.subscription_service.premium_access_for_tariff(tariff)
labels = access.get("node_labels") or access.get("squad_labels") or []
if labels:
visible = [hd.quote(str(x)) for x in labels[:8]]
servers = "\n".join(f"• {label}" for label in visible)
if len(labels) > len(visible):
more = len(labels) - len(visible)
servers += "\n" + _("traffic_warning_premium_servers_more", count=more)
else:
servers = _("traffic_warning_premium_generic_servers")
left_pct = max(0, 100 - int(level))
usage = self._usage_placeholders(used_val, limit_val)
text = _(
"traffic_warning_premium_almost",
tariff_name=hd.quote(str(tariff.name(user_lang))),
left_pct=left_pct,
servers=servers,
**usage,
)
markup = self._traffic_topup_markup(user_lang, "premium")
await self.bot.send_message(
sub.user_id,
text,
reply_markup=markup,
parse_mode="HTML",
)
except Exception:
logging.exception("Failed to send premium traffic warning to user %s", sub.user_id)
async def _premium_node_uuids_for_tariff(self, tariff) -> list[str]:
cache_key = tuple(sorted(tariff.premium_squad_uuids or []))
cached = self._premium_nodes_cache.get(cache_key)
now_ts = datetime.now(timezone.utc).timestamp()
if cached and now_ts - cached["ts"] < 600:
return list(cached["nodes"])
nodes: list[str] = []
for squad_uuid in tariff.premium_squad_uuids or []:
accessible = (
await self.panel_service.get_internal_squad_accessible_nodes(squad_uuid) or []
)
for node in accessible:
if not isinstance(node, dict):
continue
node_uuid = node.get("uuid") or node.get("nodeUuid") or node.get("node_uuid")
if node_uuid:
nodes.append(str(node_uuid))
deduped = list(dict.fromkeys(nodes))
self._premium_nodes_cache[cache_key] = {"ts": now_ts, "nodes": deduped}
return deduped
async def _premium_usage_for_user(
self,
user_uuid: str,
node_uuids: list[str],
start_date: str,
end_date: str,
*,
panel_username: Optional[str] = None,
) -> Optional[int]:
total = 0
found = False
username = (panel_username or "").strip() or None
for node_uuid in node_uuids:
stats = await self.panel_service.get_node_users_bandwidth_stats(
node_uuid,
start=start_date,
end=end_date,
)
if not stats:
continue
entries = stats.get("topUsers") or stats.get("usersStats") or stats.get("users") or []
if not isinstance(entries, list):
continue
for entry in entries:
if not isinstance(entry, dict):
continue
user_obj = entry.get("user") if isinstance(entry.get("user"), dict) else {}
entry_uuid = (
user_obj.get("uuid")
or entry.get("userUuid")
or entry.get("uuid")
or entry.get("user_uuid")
)
entry_username = (
user_obj.get("username") or entry.get("username") or entry.get("userUsername")
)
# Remnawave's /bandwidth-stats/nodes/{uuid}/users response
# currently exposes only {color, username, total}; match by
# username first, fall back to UUID if a future version
# adds it back.
matched = False
if entry_uuid and entry_uuid == user_uuid:
matched = True
elif username and entry_username and entry_username == username:
matched = True
if not matched:
continue
value = entry.get("total")
if value is None:
value = int(entry.get("download", 0) or 0) + int(entry.get("upload", 0) or 0)
total += int(value or 0)
found = True
if len(node_uuids) > 1:
await asyncio.sleep(0.1)
return total if found else 0
async def legacy_throttle_recovery_tick(self, session: AsyncSession) -> None:
"""Recover subscriptions throttled by older bot versions.
Current Remnawave versions enforce exhausted user traffic limits by
switching the user status to LIMITED, so new ticks must not remove users
from Internal Squads.
"""
result = await session.execute(
select(Subscription).where(
Subscription.is_active == True,
Subscription.is_throttled == True,
)
)
for sub in result.scalars().all():
try:
tariff = self.settings.tariffs_config.require(sub.tariff_key)
except Exception:
continue
if int(sub.traffic_limit_bytes or 0) <= int(sub.traffic_used_bytes or 0):
continue
for squad_uuid in tariff.squad_uuids:
await self.panel_service.add_users_to_internal_squad(
squad_uuid, [sub.panel_user_uuid]
)
await subscription_dal.update_subscription(
session,
sub.subscription_id,
{"is_throttled": False, "status_from_panel": "ACTIVE"},
)