Code cleanup and bug fixes

Author: snoups

README.md

@@ -47,17 +47,18 @@ The bot enables efficient subscription sales with advanced features:
     - Format text using HTML
     - Preview notifications before sending
     - System notifications for the developer and administrators
-- **Referral Program**
+- **Two-Level Referral Program** (by [@Heimlet](https://github.com/Heimlet))
     - View referral statistics
     - Reward users for inviting new members
-- **Trial Period**
+    - Support for two-tier referral rewards
+- **Trial Period** (by [@Heimlet](https://github.com/Heimlet))
     - Provide free trial subscription
+    - Extend trial period for referred users
     - Configure and disable the trial period
 - **Flexible Payment System**
     - Change the default currency
     - Easily extendable architecture for adding new payment gateways
     - ~~Add, edit, and delete subscription plans at any time~~
-    - ~~Enable or disable payment methods~~
     - ~~Change the display order of payment options~~
 - **~~User Editor~~**
     - ~~View user information~~
@@ -83,11 +84,13 @@ Administrators do not have access to server management.
 
 
 ### 🚧 Current Tasks
-- [ ] Trial period
-- [ ] Referral system
+- [x] Trial period
+- [x] Referral system
 - [ ] Statistics
 - [ ] User editor
 - [ ] Plans editor
+- [ ] Flexible server pool
+- [ ] Custom promocodes
 
 <a id="installation-guide"></a>
 
@@ -140,17 +143,14 @@ Before starting the installation, make sure you have the installed [**Docker**](
 | | | |
 | SHOP_EMAIL | ⭕ | [email protected] | Email for receipts |
 | SHOP_CURRENCY | ⭕ | RUB | Currency for buttons (e.g., RUB, USD, XTR) |
-| SHOP_TRIAL_ENABLED | ⭕ | True | Enable trial subscription for new users. |
-| SHOP_TRIAL_PERIOD | ⭕ | 3 | Duration of the trial subscription in days. |
-| SHOP_REFERRED_TRIAL_ENABLED | ⭕ | False | Enable specific trial subscription for referral users. |
-| SHOP_REFERRED_TRIAL_PERIOD | ⭕ | 7 | Trial duration in days for the referred (invited) user |
+| SHOP_TRIAL_ENABLED | ⭕ | True | Enable trial subscription for new users |
+| SHOP_TRIAL_PERIOD | ⭕ | 3 | Duration of the trial subscription in days |
+| SHOP_REFERRED_TRIAL_ENABLED | ⭕ | False | Enable extended trial period for referred users |
+| SHOP_REFERRED_TRIAL_PERIOD | ⭕ | 7 | Duration of the extended trial for referred users (in days) |
 | SHOP_REFERRER_REWARD_ENABLED | ⭕ | True | Enable the two-level referral reward system |
-| ~~SHOP_REFERRER_REWARD_TYPE~~ | ⭕ | DAYS | Type of referrer reward. 'days' only now. Awaits user balance implementation. (e.g. days, money) |
-| SHOP_REFERRER_LEVEL_ONE_PERIOD | ⭕ | 10 | Days reward for the first-level referrer (user who invited) |
-| SHOP_REFERRER_LEVEL_TWO_PERIOD | ⭕ | 3 | Days reward for the second-level referrer (user invited by the invited) |
-| ~~SHOP_REFERRER_LEVEL_ONE_RATE~~ | ⭕ | 50 | Percentage reward for the first-level referrer (user who invited) |
-| ~~SHOP_REFERRER_LEVEL_TWO_RATE~~ | ⭕ | 5 | Percentage reward for the second-level referrer (user invited by the invited) |
-| SHOP_BONUS_DEVICES_COUNT | ⭕ | 1 | Number of devices by default for trial and referral users (mirrors tariff plans settings) |
+| SHOP_REFERRER_LEVEL_ONE_PERIOD | ⭕ | 10 | Reward in days for the first-level referrer (inviter) |
+| SHOP_REFERRER_LEVEL_TWO_PERIOD | ⭕ | 3 | Reward in days for the second-level referrer (inviter of the inviter). |
+| SHOP_BONUS_DEVICES_COUNT | ⭕ | 1 | Default Device Limit for Promocode, Trial, and Referral Users (Based on Plan Settings) |
 | SHOP_PAYMENT_STARS_ENABLED | ⭕ | True | Enable Telegram stars payment |
 | SHOP_PAYMENT_CRYPTOMUS_ENABLED | ⭕ | False | Enable Cryptomus payment |
 | SHOP_PAYMENT_YOOKASSA_ENABLED | ⭕ | False | Enable Yookassa payment |
@@ -171,12 +171,6 @@ Before starting the installation, make sure you have the installed [**Docker**](
 | YOOMONEY_WALLET_ID | ⭕ | - | Wallet ID for Yoomoney payment |
 | YOOMONEY_NOTIFICATION_SECRET | ⭕ | - | Notification secret key for Yoomoney payment |
 | | | |
-| REDIS_HOST | ⭕ | 3xui-shop-redis | Host of the Redis server |
-| REDIS_PORT | ⭕ | 6379 | Port of the Redis server |
-| REDIS_DB_NAME | ⭕ | 0 | Name of the Redis database |
-| REDIS_USERNAME | ⭕ | - | Username for authentication in the Redis server |
-| REDIS_PASSWORD | ⭕ | - | Password for authentication in the Redis server |
-| | | |
 | LOG_LEVEL | ⭕ | DEBUG | Log level (e.g., INFO, DEBUG) |
 | LOG_FORMAT | ⭕ | %(asctime)s \| %(name)s \| %(levelname)s \| %(message)s | Log format |
 | LOG_ARCHIVE_FORMAT | ⭕ | zip | Log archive format (e.g., zip, gz) |
@@ -262,14 +256,15 @@ To ensure the bot functions correctly, you must configure the 3X-UI panel:
 <a id="bugs-and-feature-requests"></a>
 
 ### Referral and Trial Rewards Configuration
-Shop now supports **trial subscriptions** and a **two-level referral reward system**. Here’s how it works:
-All configuration is available via .env (see it above).
-
-| Type of reward                   | How it works                                                                                                                                                                                                                        |
-|----------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| Trial period                     | A trial subscription is available by 'TRY FOR FREE' button at start menu to any user who opens the bot and does not have an active subscription.                                                                                    |
-| Referred user  Trial period      | This option is just like previous 'trial period', but allows bot admin to configure **extended trial period** for an invited user.                                                                                                  |
-| Referral 2-level payment rewards | When a referred user pays for a subscription, the referrer and the second-level referrer (the user who invited the referrer) receive ~~a percentage of the payment as a reward~~ fixed count of days at the moment fore each level. |
+
+Bot now supports **trial subscriptions** and a **two-level referral reward system**. Here’s how it works:
+All configuration is available via `.env` [(see it above)](#environment-variables-configuration).
+
+| Type of reward | How it works |
+| - | - |
+| Trial period | A trial subscription is available by 'TRY FOR FREE' button at start menu to any user who opens the bot and does not have an active subscription. |
+| Extended Trial period | This option is just like previous 'trial period', but allows to configure **extended trial period** for an invited user. |
+| Two-Level Referral Payment Rewards | When a referred user pays for a subscription, the referrer and the second-level referrer (the user who invited the referrer) receive fixed count of days at the moment fore each level. |
 
 ## 🐛 Bugs and Feature Requests
 

app/__main__.py

@@ -51,7 +51,9 @@ async def on_startup(config: Config, bot: Bot, services: ServicesContainer, db:
 
     tasks.transactions.start_scheduler(db.session)
     if config.shop.REFERRER_REWARD_ENABLED:
-        tasks.referral.start_referral_scheduler(session_factory=db.session, referral_service=services.referral)
+        tasks.referral.start_scheduler(
+            session_factory=db.session, referral_service=services.referral
+        )
 
 
 async def main() -> None:
@@ -66,6 +68,7 @@ async def main() -> None:
 
     # Initialize database
     db = Database(config.database)
+    await db.initialize()
 
     # Set up storage for FSM (Finite State Machine)
     storage = RedisStorage.from_url(url=config.redis.url())
@@ -114,7 +117,7 @@ async def main() -> None:
     dispatcher.shutdown.register(on_shutdown)
 
     # Enable Maintenance mode for developing # WARNING: remove before production
-    MaintenanceMiddleware.set_mode(True)
+    MaintenanceMiddleware.set_mode(False)
 
     # Register middlewares
     middlewares.register(dispatcher=dispatcher, i18n=i18n, session=db.session)
@@ -129,9 +132,6 @@ async def main() -> None:
     # Include bot routers
     routers.include(app=app, dispatcher=dispatcher)
 
-    # Initialize database
-    await db.initialize()
-
     # Set up bot commands
     await commands.setup(bot)
 

app/bot/middlewares/maintenance.py

@@ -33,7 +33,11 @@ async def __call__(
 
                 if self.active and not is_admin and user.id != event.bot.id:
                     logger.info(f"User {user.id} tried to use bot in maintenance")
-                    message = event.message or event.callback_query.message
+                    
+                    if event.message:
+                        message = event.message
+                    elif event.callback_query and event.callback_query.message:
+                        message = event.callback_query.message
 
                     if message:
                         await NotificationService.notify_by_message(

app/bot/payment_gateways/_gateway.py

@@ -1,7 +1,6 @@
 import logging
 from abc import ABC, abstractmethod
 
-import requests
 from aiogram import Bot
 from aiogram.fsm.storage.redis import RedisStorage
 from aiogram.utils.i18n import I18n
@@ -43,7 +42,7 @@ def __init__(
         bot: Bot,
         i18n: I18n,
         services: ServicesContainer,
-    ):
+    ) -> None:
         self.app = app
         self.config = config
         self.session = session
@@ -82,7 +81,7 @@ async def _on_payment_succeeded(self, payment_id: str) -> None:
         if self.config.shop.REFERRER_REWARD_ENABLED:
             await self.services.referral.add_referrers_rewards_on_payment(
                 referred_tg_id=data.user_id,
-                payment_amount=data.price,  # todo: (!) add currency unified processing
+                payment_amount=data.price,  # TODO: (!) add currency unified processing
                 payment_id=payment_id,
             )
 

app/bot/payment_gateways/telegram_stars.py

@@ -34,7 +34,7 @@ def __init__(
         bot: Bot,
         i18n: I18n,
         services: ServicesContainer,
-    ):
+    ) -> None:
         self.name = __("payment:gateway:telegram_stars")
         self.app = app
         self.config = config

app/bot/payment_gateways/yookassa.py

@@ -42,7 +42,7 @@ def __init__(
         bot: Bot,
         i18n: I18n,
         services: ServicesContainer,
-    ):
+    ) -> None:
         self.name = __("payment:gateway:yookassa")
         self.app = app
         self.config = config

app/bot/payment_gateways/yoomoney.py

@@ -36,7 +36,7 @@ def __init__(
         bot: Bot,
         i18n: I18n,
         services: ServicesContainer,
-    ):
+    ) -> None:
         self.name = __("payment:gateway:yoomoney")
         self.app = app
         self.config = config

app/bot/routers/main_menu/handler.py

@@ -14,36 +14,36 @@
 from app.bot.utils.constants import MAIN_MESSAGE_ID_KEY
 from app.bot.utils.navigation import NavMain
 from app.config import Config
-from app.db.models import User, Referral
+from app.db.models import Referral, User
 
 from .keyboard import main_menu_keyboard
 
 logger = logging.getLogger(__name__)
 router = Router(name=__name__)
 
 
-async def process_creating_referral(
-    session: AsyncSession,
-    user: User,
-    referrer_id: int
-) -> bool:
+async def process_creating_referral(session: AsyncSession, user: User, referrer_id: int) -> bool:
     logger.info(f"Assigning user {user.tg_id} as a referred to a referrer user {referrer_id}")
     try:
         referrer = await User.get(session=session, tg_id=referrer_id)
         if not referrer or referrer.tg_id == user.tg_id:
-            logger.info(f"Failed to assign user {user.tg_id} as a referred to a referrer user {referrer_id}."
-                        f"Invalid string received.")
+            logger.info(
+                f"Failed to assign user {user.tg_id} as a referred to a referrer user {referrer_id}."
+                f"Invalid string received."
+            )
             return False
 
         await Referral.create(
-            session=session,
-            referrer_tg_id=referrer.tg_id,
-            referred_tg_id=user.tg_id
+            session=session, referrer_tg_id=referrer.tg_id, referred_tg_id=user.tg_id
+        )
+        logger.info(
+            f"User {user.tg_id} assigned as referred to a referrer with tg id {referrer.tg_id}"
         )
-        logger.info(f"User {user.tg_id} assigned as referred to a referrer with tg id {referrer.tg_id}")
         return True
-    except Exception as e:
-        logger.critical(f"Error creating Referral to a referred {user.tg_id} (start arg: {referrer_id}): {e}")
+    except Exception as exception:
+        logger.critical(
+            f"Referral creation error for {user.tg_id} (arg: {referrer_id}): {exception}"
+        )
         return False
 
 
@@ -56,7 +56,7 @@ async def command_main_menu(
     config: Config,
     session: AsyncSession,
     command: CommandObject,
-    is_new_user: bool
+    is_new_user: bool,
 ) -> None:
     logger.info(f"User {user.tg_id} opened main menu page.")
     previous_message_id = await state.get_value(MAIN_MESSAGE_ID_KEY)
@@ -73,9 +73,7 @@ async def command_main_menu(
     received_referrer_id = int(command.args) if command.args and command.args.isdigit() else None
     if received_referrer_id and is_new_user:
         await process_creating_referral(
-            session=session,
-            user=user,
-            referrer_id=received_referrer_id
+            session=session, user=user, referrer_id=received_referrer_id
         )
 
     is_admin = await IsAdmin()(user_id=user.tg_id)
@@ -93,11 +91,11 @@ async def command_main_menu(
 
 @router.callback_query(F.data == NavMain.MAIN_MENU)
 async def callback_main_menu(
-        callback: CallbackQuery,
-        user: User,
-        services: ServicesContainer,
-        state: FSMContext,
-        config: Config,
+    callback: CallbackQuery,
+    user: User,
+    services: ServicesContainer,
+    state: FSMContext,
+    config: Config,
 ) -> None:
     logger.info(f"User {user.tg_id} returned to main menu page.")
     await state.clear()
@@ -142,8 +140,10 @@ async def redirect_to_main_menu(
                 is_admin,
                 is_referral_available=config.shop.REFERRER_REWARD_ENABLED,
                 is_trial_available=await services.subscription.is_trial_available(user),
-                is_referred_trial_available=await services.referral.is_referred_trial_available(user),
+                is_referred_trial_available=await services.referral.is_referred_trial_available(
+                    user
+                ),
             ),
         )
-    except Exception as e:
-        logger.critical(f"Error redirecting to main menu page: {e}")
+    except Exception as exception:
+        logger.error(f"Error redirecting to main menu page: {exception}")

app/bot/routers/main_menu/keyboard.py

@@ -12,10 +12,10 @@
 
 
 def main_menu_keyboard(
-        is_admin: bool = False,
-        is_referral_available: bool = False,
-        is_trial_available: bool = False,
-        is_referred_trial_available: bool = False,
+    is_admin: bool = False,
+    is_referral_available: bool = False,
+    is_trial_available: bool = False,
+    is_referred_trial_available: bool = False,
 ) -> InlineKeyboardMarkup:
     builder = InlineKeyboardBuilder()
 
@@ -29,8 +29,7 @@ def main_menu_keyboard(
     elif is_trial_available:
         builder.row(
             InlineKeyboardButton(
-                text=_("subscription:button:get_trial"),
-                callback_data=NavSubscription.GET_TRIAL
+                text=_("subscription:button:get_trial"), callback_data=NavSubscription.GET_TRIAL
             )
         )
 
@@ -45,12 +44,16 @@ def main_menu_keyboard(
         ),
     )
     builder.row(
-        *([
-              InlineKeyboardButton(
-                  text=_("main_menu:button:referral"),
-                  callback_data=NavReferral.MAIN,
-              )
-          ] if is_referral_available else []),
+        *(
+            [
+                InlineKeyboardButton(
+                    text=_("main_menu:button:referral"),
+                    callback_data=NavReferral.MAIN,
+                )
+            ]
+            if is_referral_available
+            else []
+        ),
         InlineKeyboardButton(
             text=_("main_menu:button:support"),
             callback_data=NavSupport.MAIN,