rate configurator/ purchase of conditions/ fixed payments/ refactoring/ major improvements and fixes

database/hot_leads.py

@@ -1,8 +1,8 @@
 from sqlalchemy import func, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from database.models import Key, Payment, User
 from core.constants import PAYMENT_SYSTEMS_EXCLUDED
+from database.models import Key, Payment, User
 
 
 async def get_hot_leads(session: AsyncSession):

database/keys.py

@@ -17,29 +17,49 @@ async def store_key(
     key: str,
     server_id: str,
     remnawave_link: str = None,
-    tariff_id: int = None,
-    alias: str = None,
+    tariff_id: int | None = None,
+    alias: str | None = None,
+    selected_device_limit: int | None = None,
+    selected_traffic_limit: int | None = None,
+    selected_price_rub: int | None = None,
+    current_device_limit: int | None = None,
+    current_traffic_limit: int | None = None,
 ):
+    """Сохраняет или обновляет ключ подписки."""
     try:
         exists = await session.execute(select(Key).where(Key.tg_id == tg_id, Key.client_id == client_id))
         existing_key = exists.scalar_one_or_none()
 
         if existing_key:
-            await session.execute(
-                update(Key)
-                .where(Key.tg_id == tg_id, Key.client_id == client_id)
-                .values(
-                    email=email,
-                    expiry_time=expiry_time,
-                    key=key,
-                    server_id=server_id,
-                    remnawave_link=remnawave_link,
-                    tariff_id=tariff_id,
-                    alias=alias,
-                )
-            )
+            values: dict = {
+                "email": email,
+                "expiry_time": expiry_time,
+                "key": key,
+                "server_id": server_id,
+                "remnawave_link": remnawave_link,
+                "tariff_id": tariff_id,
+                "alias": alias,
+            }
+
+            if selected_device_limit is not None:
+                values["selected_device_limit"] = selected_device_limit
+            if selected_traffic_limit is not None:
+                values["selected_traffic_limit"] = selected_traffic_limit
+            if selected_price_rub is not None:
+                values["selected_price_rub"] = selected_price_rub
+            if current_device_limit is not None:
+                values["current_device_limit"] = current_device_limit
+            if current_traffic_limit is not None:
+                values["current_traffic_limit"] = current_traffic_limit
+
+            await session.execute(update(Key).where(Key.tg_id == tg_id, Key.client_id == client_id).values(**values))
             logger.info(f"[Store Key] Ключ обновлён: tg_id={tg_id}, client_id={client_id}, server_id={server_id}")
         else:
+            if current_device_limit is None:
+                current_device_limit = selected_device_limit
+            if current_traffic_limit is None:
+                current_traffic_limit = selected_traffic_limit
+
             new_key = Key(
                 tg_id=tg_id,
                 client_id=client_id,
@@ -51,6 +71,11 @@ async def store_key(
                 remnawave_link=remnawave_link,
                 tariff_id=tariff_id,
                 alias=alias,
+                selected_device_limit=selected_device_limit,
+                selected_traffic_limit=selected_traffic_limit,
+                selected_price_rub=selected_price_rub,
+                current_device_limit=current_device_limit,
+                current_traffic_limit=current_traffic_limit,
             )
             session.add(new_key)
             logger.info(f"[Store Key] Ключ создан: tg_id={tg_id}, client_id={client_id}, server_id={server_id}")
@@ -79,6 +104,7 @@ async def get_key_by_server(session: AsyncSession, tg_id: int, client_id: str):
 
 
 async def get_key_details(session: AsyncSession, email: str) -> dict | None:
+    """Возвращает подробную информацию о ключе по email."""
     stmt = select(Key, User).join(User, Key.tg_id == User.tg_id).where(Key.email == email)
     result = await session.execute(stmt)
     row = result.first()
@@ -116,6 +142,11 @@ async def get_key_details(session: AsyncSession, email: str) -> dict | None:
         "cluster_name": key.server_id,
         "location_name": key.server_id,
         "tariff_id": key.tariff_id,
+        "selected_device_limit": key.selected_device_limit,
+        "selected_traffic_limit": key.selected_traffic_limit,
+        "selected_price_rub": key.selected_price_rub,
+        "current_device_limit": key.current_device_limit,
+        "current_traffic_limit": key.current_traffic_limit,
     }
 
 
@@ -156,13 +187,18 @@ async def mark_key_as_frozen(session: AsyncSession, tg_id: int, client_id: str,
                 is_frozen = TRUE
             WHERE tg_id = :tg_id
               AND client_id = :client_id
-        """
+            """
         ),
         {"expiry": time_left, "tg_id": tg_id, "client_id": client_id},
     )
 
 
-async def mark_key_as_unfrozen(session: AsyncSession, tg_id: int, client_id: str, new_expiry_time: int):
+async def mark_key_as_unfrozen(
+    session: AsyncSession,
+    tg_id: int,
+    client_id: str,
+    new_expiry_time: int,
+):
     await session.execute(
         text(
             """
@@ -171,7 +207,7 @@ async def mark_key_as_unfrozen(session: AsyncSession, tg_id: int, client_id: str
                 is_frozen = FALSE
             WHERE tg_id = :tg_id
               AND client_id = :client_id
-        """
+            """
         ),
         {"expiry": new_expiry_time, "tg_id": tg_id, "client_id": client_id},
     )
@@ -199,3 +235,53 @@ async def update_key_link(session: AsyncSession, email: str, link: str) -> bool:
     res = await session.execute(q)
     await session.commit()
     return res.scalar_one_or_none() is not None
+
+
+async def save_key_config_with_mode(
+    session: AsyncSession,
+    email: str,
+    selected_devices: int | None,
+    selected_traffic_gb: int | None,
+    total_price: int,
+    has_device_choice: bool,
+    has_traffic_choice: bool,
+    config_mode: str,
+) -> None:
+    values: dict = {}
+
+    if config_mode == "pack":
+        if has_device_choice and selected_devices is not None:
+            values["current_device_limit"] = int(selected_devices)
+        if has_traffic_choice and selected_traffic_gb is not None:
+            values["current_traffic_limit"] = int(selected_traffic_gb)
+    else:
+        device_val = int(selected_devices) if selected_devices is not None and has_device_choice else None
+        traffic_val = int(selected_traffic_gb) if selected_traffic_gb is not None and has_traffic_choice else None
+
+        values["selected_device_limit"] = device_val
+        values["selected_traffic_limit"] = traffic_val
+        values["selected_price_rub"] = int(total_price)
+        values["current_device_limit"] = device_val
+        values["current_traffic_limit"] = traffic_val
+
+    if not values:
+        return
+
+    await session.execute(update(Key).where(Key.email == email).values(**values))
+
+
+async def reset_key_current_limits_to_selected(session: AsyncSession, client_id: str):
+    """Сбрасывает текущие лимиты к выбранным для ключа."""
+    await session.execute(
+        text(
+            """
+            UPDATE keys
+            SET current_device_limit = selected_device_limit,
+                current_traffic_limit = selected_traffic_limit
+            WHERE client_id = :client_id
+            """
+        ),
+        {"client_id": client_id},
+    )
+    await session.commit()
+    logger.info(f"Текущие лимиты ключа {client_id} сброшены к выбранным")

database/models.py

@@ -77,6 +77,13 @@ class Key(DictLikeMixin, Base):
     notified = Column(Boolean, default=False)
     notified_24h = Column(Boolean, default=False)
 
+    selected_device_limit = Column(Integer, nullable=True)
+    selected_traffic_limit = Column(BigInteger, nullable=True)
+    selected_price_rub = Column(Integer, nullable=True)
+
+    current_device_limit = Column(Integer, nullable=True)
+    current_traffic_limit = Column(BigInteger, nullable=True)
+
 
 class Tariff(DictLikeMixin, Base):
     __tablename__ = "tariffs"
@@ -95,6 +102,17 @@ class Tariff(DictLikeMixin, Base):
     sort_order = Column(Integer, nullable=True)
     vless = Column(Boolean, default=False)
 
+    configurable = Column(Boolean, nullable=False, server_default="false")
+
+    device_options = Column(JSONB, nullable=True)
+    traffic_options_gb = Column(JSONB, nullable=True)
+
+    device_step_rub = Column(Integer, nullable=True)
+    device_overrides = Column(JSONB, nullable=True)
+
+    traffic_step_rub = Column(Integer, nullable=True)
+    traffic_overrides = Column(JSONB, nullable=True)
+
 
 class Server(DictLikeMixin, Base):
     __tablename__ = "servers"

database/referrals.py

@@ -3,6 +3,7 @@ from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from config import CHECK_REFERRAL_REWARD_ISSUED, REFERRAL_BONUS_PERCENTAGES
+from core.bootstrap import BUTTONS_CONFIG
 from database.models import Referral
 from logger import logger
 
@@ -57,6 +58,11 @@ async def mark_referral_reward_issued(session: AsyncSession, referred_tg_id: int
 
 
 async def get_total_referral_bonus(session: AsyncSession, referrer_tg_id: int, max_levels: int) -> float:
+    referral_enabled = bool(BUTTONS_CONFIG.get("REFERRAL_BUTTON_ENABLED", True))
+    if not referral_enabled:
+        logger.debug("Реферальная программа отключена, бонусы не начисляются")
+        return 0.0
+
     if CHECK_REFERRAL_REWARD_ISSUED:
         bonus_cte = """
             WITH RECURSIVE

database/statistics.py

@@ -3,8 +3,8 @@ from datetime import date, datetime
 from sqlalchemy import and_, exists, func, not_, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from database.models import Key, Payment, Referral, Tariff, User
 from core.constants import PAYMENT_SYSTEMS_EXCLUDED
+from database.models import Key, Payment, Referral, Tariff, User
 
 
 async def count_total_users(session: AsyncSession) -> int:

database/tracking_sources.py

@@ -2,9 +2,9 @@ from sqlalchemy import and_, func, insert, not_, select
 from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from core.constants import PAYMENT_SYSTEMS_EXCLUDED
 from database.models import Payment, TrackingSource, User
 from logger import logger
-from core.constants import PAYMENT_SYSTEMS_EXCLUDED
 
 
 async def create_tracking_source(session: AsyncSession, name: str, code: str, type_: str, created_by: int):