From 94612e11f2dbf1c91d8e882120734960f5d34da4 Mon Sep 17 00:00:00 2001 From: 3252a8 <3252a8@proton.me> Date: Tue, 12 May 2026 22:01:11 +0300 Subject: [PATCH] fix: fix unlimited premium squads applying via admin panel --- bot/app/web/admin_api.py | 6 +++ bot/handlers/admin/user_management.py | 6 +++ bot/services/subscription_service.py | 68 ++++++++++++++++++++++++++- bot/services/tariff_worker.py | 38 +++++++++++++-- tests/test_admin_traffic_grants.py | 52 ++++++++++++++++++++ 5 files changed, 163 insertions(+), 7 deletions(-) diff --git a/bot/app/web/admin_api.py b/bot/app/web/admin_api.py index e52dfa1..4b6a201 100644 --- a/bot/app/web/admin_api.py +++ b/bot/app/web/admin_api.py @@ -1241,6 +1241,7 @@ async def admin_user_premium_override_route(request: web.Request) -> web.Respons actor_id = _require_admin_user_id(request) target_id = int(request.match_info["user_id"]) payload = await _read_json(request) + subscription_service = request.app.get("subscription_service") unlimited = bool(payload.get("unlimited")) bonus_bytes_raw = payload.get("bonus_bytes") @@ -1285,6 +1286,11 @@ async def admin_user_premium_override_route(request: web.Request) -> web.Respons await session.commit() await session.refresh(active) + if subscription_service is not None: + await subscription_service.sync_premium_squad_access_to_panel(session, target_id) + await session.commit() + await session.refresh(active) + return _ok({"subscription": _serialize_subscription(active)}) diff --git a/bot/handlers/admin/user_management.py b/bot/handlers/admin/user_management.py index ef7f99b..f05fcb2 100644 --- a/bot/handlers/admin/user_management.py +++ b/bot/handlers/admin/user_management.py @@ -759,6 +759,9 @@ async def handle_premium_override_apply( ) await session.commit() + await subscription_service.sync_premium_squad_access_to_panel(session, user.user_id) + await session.commit() + await callback.answer(_("admin_premium_override_saved"), show_alert=False) await handle_refresh_user_card( callback, user, subscription_service, session, i18n_instance, lang @@ -1757,6 +1760,9 @@ async def process_premium_override_bonus_handler( ) await session.commit() + await subscription_service.sync_premium_squad_access_to_panel(session, target_user_id) + await session.commit() + await message.answer( _("admin_premium_override_bonus_set", gb=f"{gb:.2f}", user_id=target_user_id) ) diff --git a/bot/services/subscription_service.py b/bot/services/subscription_service.py index 9088658..9dfb295 100644 --- a/bot/services/subscription_service.py +++ b/bot/services/subscription_service.py @@ -1121,8 +1121,11 @@ class SubscriptionService: previous_topup_used = int(sub.premium_topup_used_bytes or 0) if same_period else 0 premium_used = int(sub.premium_used_bytes or 0) if same_period else 0 premium_baseline = int(tariff.premium_monthly_bytes or sub.premium_baseline_bytes or 0) + premium_bonus = max(0, int(getattr(sub, "premium_bonus_bytes", 0) or 0)) premium_topup_balance = int(sub.premium_topup_balance_bytes or 0) + purchase_bytes - overflow_to_cover = max(0, premium_used - premium_baseline - previous_topup_used) + overflow_to_cover = max( + 0, premium_used - premium_baseline - previous_topup_used - premium_bonus + ) consume_now = min(premium_topup_balance, overflow_to_cover) premium_topup_balance -= consume_now premium_topup_used = previous_topup_used + consume_now @@ -1130,8 +1133,12 @@ class SubscriptionService: premium_baseline, premium_topup_balance, premium_topup_used, + premium_bonus, + ) + premium_unlimited = bool(getattr(sub, "premium_unlimited_override", False)) + premium_is_limited = ( + not premium_unlimited and premium_limit > 0 and premium_used >= premium_limit ) - premium_is_limited = premium_limit > 0 and premium_used >= premium_limit updated_sub = await subscription_dal.update_subscription( session, @@ -1173,6 +1180,63 @@ class SubscriptionService: "tariff_key": tariff.key, } + async def sync_premium_squad_access_to_panel( + self, + session: AsyncSession, + user_id: int, + ) -> None: + """Recompute premium quota flags from DB and push internal squads to Remnawave. + + Used when admin overrides change without going through the traffic worker + (Telegram/Web admin premium bonus / unlimited). + """ + db_user = await user_dal.get_user_by_id(session, user_id) + if not db_user or not db_user.panel_user_uuid: + return + sub = await subscription_dal.get_active_subscription_by_user_id( + session, user_id, db_user.panel_user_uuid + ) + if not sub: + return + tariff = self._resolve_tariff(sub.tariff_key) if sub.tariff_key else None + if not tariff or not getattr(tariff, "premium_squad_uuids", None): + return + + premium_baseline = int(tariff.premium_monthly_bytes or sub.premium_baseline_bytes or 0) + premium_bonus = max(0, int(getattr(sub, "premium_bonus_bytes", 0) 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) + premium_used = int(sub.premium_used_bytes or 0) + premium_limit = self._premium_effective_limit_bytes( + premium_baseline, + premium_topup_balance, + premium_topup_used, + premium_bonus, + ) + premium_unlimited = bool(getattr(sub, "premium_unlimited_override", False)) + premium_is_limited = ( + not premium_unlimited and premium_limit > 0 and premium_used >= premium_limit + ) + + if bool(getattr(sub, "premium_is_limited", False)) != premium_is_limited: + await subscription_dal.update_subscription( + session, + sub.subscription_id, + {"premium_is_limited": premium_is_limited}, + ) + + squads = self._panel_squads_for_tariff(tariff, include_premium=not premium_is_limited) + try: + await self.panel_service.update_user_details_on_panel( + db_user.panel_user_uuid, + {"uuid": db_user.panel_user_uuid, "activeInternalSquads": squads}, + log_response=False, + ) + except Exception: + logging.exception( + "sync_premium_squad_access_to_panel: failed to push squads for user %s", user_id + ) + async def admin_grant_topup( self, session: AsyncSession, diff --git a/bot/services/tariff_worker.py b/bot/services/tariff_worker.py index 8b2f245..3fc2c4d 100644 --- a/bot/services/tariff_worker.py +++ b/bot/services/tariff_worker.py @@ -152,7 +152,12 @@ class TariffTrafficWorker: ) await self._sync_premium_squad_limit( - session, sub, tariff, now, panel_username=panel_username + session, + sub, + tariff, + now, + panel_username=panel_username, + panel_user_dict=panel_data, ) async def _ensure_period_reset_strategy( @@ -281,6 +286,7 @@ class TariffTrafficWorker: now: datetime, *, panel_username: Optional[str] = None, + panel_user_dict: Optional[dict] = None, ) -> None: if not getattr(tariff, "premium_squad_uuids", None): if ( @@ -359,6 +365,17 @@ class TariffTrafficWorker: or bool(sub.premium_is_limited) != should_limit or getattr(sub, "premium_period_start_at", None) != premium_period_start ) + 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_raw = panel_user_dict.get("activeInternalSquads") or panel_user_dict.get( + "active_internal_squads" + ) + if desired_set != self._internal_squad_uuid_set(current_raw): + changed = True sub.premium_baseline_bytes = premium_baseline sub.premium_topup_balance_bytes = premium_topup_balance sub.premium_topup_used_bytes = premium_topup_used @@ -377,10 +394,7 @@ class TariffTrafficWorker: if not changed: return - squads = self.subscription_service._panel_squads_for_tariff( - tariff, - include_premium=not should_limit, - ) + squads = desired_squads await self.panel_service.update_user_details_on_panel( sub.panel_user_uuid, {"uuid": sub.panel_user_uuid, "activeInternalSquads": squads}, @@ -395,6 +409,20 @@ class TariffTrafficWorker: 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))) diff --git a/tests/test_admin_traffic_grants.py b/tests/test_admin_traffic_grants.py index 1f82caf..f4ebf38 100644 --- a/tests/test_admin_traffic_grants.py +++ b/tests/test_admin_traffic_grants.py @@ -234,6 +234,58 @@ class AdminGrantTopupTests(unittest.IsolatedAsyncioTestCase): ): self.assertIsNone(await service.admin_grant_premium_topup(AsyncMock(), 11, 10.0)) + async def test_sync_premium_squad_access_updates_limited_flag_and_panel(self): + with tempfile.TemporaryDirectory() as tmpdir: + settings = _make_settings(_tariffs_config_payload(premium=True), tmpdir) + panel_service = AsyncMock(spec=PanelApiService) + panel_service.update_user_details_on_panel = AsyncMock(return_value={"response": {}}) + service = SubscriptionService(settings, panel_service) + + db_user = SimpleNamespace( + user_id=88, + panel_user_uuid="panel-uuid", + ) + sub = SimpleNamespace( + subscription_id=3, + user_id=88, + panel_user_uuid="panel-uuid", + tariff_key="standard", + premium_baseline_bytes=25 * (1024**3), + premium_topup_balance_bytes=0, + premium_topup_used_bytes=0, + premium_used_bytes=30 * (1024**3), + premium_is_limited=True, + premium_period_start_at=datetime.now(timezone.utc).replace( + day=1, hour=0, minute=0, second=0, microsecond=0 + ), + premium_unlimited_override=False, + premium_bonus_bytes=10 * (1024**3), + ) + + with ( + patch( + "bot.services.subscription_service.user_dal.get_user_by_id", + new=AsyncMock(return_value=db_user), + ), + patch( + "bot.services.subscription_service.subscription_dal.get_active_subscription_by_user_id", + new=AsyncMock(return_value=sub), + ), + patch( + "bot.services.subscription_service.subscription_dal.update_subscription", + new=AsyncMock(return_value=sub), + ) as upd, + ): + await service.sync_premium_squad_access_to_panel(AsyncMock(), 88) + + upd.assert_awaited_once() + self.assertFalse(upd.await_args.args[2]["premium_is_limited"]) + panel_service.update_user_details_on_panel.assert_awaited_once() + squads = panel_service.update_user_details_on_panel.await_args.args[1][ + "activeInternalSquads" + ] + self.assertIn("premium-squad", squads) + if __name__ == "__main__": unittest.main()