From 9deb98bbe72895a29a71e5536dc7755fd095a48e Mon Sep 17 00:00:00 2001 From: 3252a8 <3252a8@proton.me> Date: Tue, 28 Apr 2026 14:34:00 +0300 Subject: [PATCH] feat: add traffic periods topups and throttling --- bot/main_bot.py | 12 +++ bot/services/tariff_worker.py | 195 ++++++++++++++++++++++++++++++++++ 2 files changed, 207 insertions(+) create mode 100644 bot/services/tariff_worker.py diff --git a/bot/main_bot.py b/bot/main_bot.py index 9416b78..b0f42cf 100644 --- a/bot/main_bot.py +++ b/bot/main_bot.py @@ -30,6 +30,7 @@ from bot.routers import build_root_router from bot.services.yookassa_service import YooKassaService from bot.services.panel_api_service import PanelApiService from bot.services.subscription_service import SubscriptionService +from bot.services.tariff_worker import TariffTrafficWorker from bot.services.referral_service import ReferralService from bot.services.promo_code_service import PromoCodeService from bot.services.stars_service import StarsService @@ -273,6 +274,15 @@ async def run_bot(settings_param: Settings): dp[key] = service dp["panel_service"] = services["panel_service"] dp["async_session_factory"] = local_async_session_factory + tariff_worker = TariffTrafficWorker( + settings_param, + local_async_session_factory, + services["panel_service"], + services["subscription_service"], + bot, + i18n_instance, + ) + dp["tariff_worker"] = tariff_worker # Wrap startup/shutdown handlers to satisfy aiogram event signature (no args passed) async def _on_startup_wrapper(): @@ -306,6 +316,8 @@ async def run_bot(settings_param: Settings): await build_and_start_web_app(dp, bot, settings_param, local_async_session_factory) main_tasks.append(asyncio.create_task(web_server_task(), name="AIOHTTPServerTask")) + if settings_param.tariffs_config: + main_tasks.append(asyncio.create_task(tariff_worker.run(), name="TariffTrafficWorker")) # Recurring billing moved to panel webhook (24h before expiry). No periodic task needed here. diff --git a/bot/services/tariff_worker.py b/bot/services/tariff_worker.py new file mode 100644 index 0000000..71e3899 --- /dev/null +++ b/bot/services/tariff_worker.py @@ -0,0 +1,195 @@ +import asyncio +import logging +from datetime import datetime, timedelta, timezone +from typing import Optional + +from aiogram import Bot +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 config.settings import Settings +from db.dal import subscription_dal, tariff_dal +from db.models import Subscription + + +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() + + 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.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) + 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, _ = self.subscription_service._extract_panel_traffic_details(panel_data) + 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 tariff.billing_model == "period": + await self._maybe_reset_period(session, sub, tariff, used) + await self._maybe_warn_or_throttle(session, sub, tariff, used, limit) + + async def _maybe_reset_period(self, session: AsyncSession, sub: Subscription, tariff, used: Optional[int]) -> None: + if not sub.period_start_at: + sub.period_start_at = datetime.now(timezone.utc) + return + now = datetime.now(timezone.utc) + next_reset = sub.period_start_at + timedelta(days=30) + if now < next_reset or sub.end_date <= now: + return + + used_now = int(used or sub.traffic_used_bytes or 0) + baseline = int(sub.tier_baseline_bytes or tariff.monthly_bytes) + topup_used = max(0, used_now - baseline) + new_topup = max(0, int(sub.topup_balance_bytes or 0) - topup_used) + await self.panel_service.reset_user_traffic(sub.panel_user_uuid) + new_limit = baseline + new_topup + payload = self.subscription_service._build_panel_update_payload( + panel_user_uuid=sub.panel_user_uuid, + expire_at=sub.end_date, + status="ACTIVE", + traffic_limit_bytes=new_limit, + ) + payload["activeInternalSquads"] = tariff.squad_uuids + await self.panel_service.update_user_details_on_panel(sub.panel_user_uuid, payload, log_response=False) + await subscription_dal.update_subscription( + session, + sub.subscription_id, + { + "period_start_at": next_reset, + "traffic_used_bytes": 0, + "traffic_limit_bytes": new_limit, + "topup_balance_bytes": new_topup, + "is_throttled": False, + }, + ) + await tariff_dal.clear_period_warnings(session, sub.subscription_id) + + async def _maybe_warn_or_throttle( + self, + session: AsyncSession, + sub: Subscription, + tariff, + used: Optional[int], + limit: Optional[int], + ) -> None: + 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 + for level, threshold in ((80, 0.8), (95, 0.95), (100, 1.0)): + if ratio < threshold: + continue + warning = await tariff_dal.get_warning( + session, + subscription_id=sub.subscription_id, + period_start_at=sub.period_start_at 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=sub.period_start_at 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: + left_pct = max(0, 100 - level) + if level < 100: + text = f"Трафик тарифа {tariff.name(self.settings.DEFAULT_LANGUAGE)} почти закончился. Осталось около {left_pct}%." + else: + text = "Трафик закончился. Доступ временно ограничен до сброса или докупки пакета." + await self.bot.send_message(sub.user_id, text) + except Exception: + logging.exception("Failed to send traffic warning to user %s", sub.user_id) + if level == 100: + await self._throttle(session, sub, tariff) + + async def _throttle(self, session: AsyncSession, sub: Subscription, tariff) -> None: + if sub.is_throttled: + return + for squad_uuid in tariff.squad_uuids: + await self.panel_service.remove_users_from_internal_squad(squad_uuid, [sub.panel_user_uuid]) + await subscription_dal.update_subscription( + session, + sub.subscription_id, + {"is_throttled": True, "status_from_panel": "THROTTLED_BY_BOT"}, + ) + + async def throttle_recovery_tick(self, session: AsyncSession) -> None: + 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"}, + )