From caf5b88eef7e5dc2b97d3dabb84f3c21306a0bc0 Mon Sep 17 00:00:00 2001 From: 3252a8 <3252a8@proton.me> Date: Tue, 11 Nov 2025 22:32:30 +0300 Subject: [PATCH] Use random string in ref link instead of tg ID - Enable old links with tg id by default in .env with LEGACY_REFS=true --- .env.example | 1 + bot/handlers/inline_mode.py | 28 +++++++++--- bot/handlers/user/referral.py | 25 +++++++++-- bot/handlers/user/start.py | 22 ++++++++-- bot/services/referral_service.py | 32 ++++++++++++-- config/settings.py | 6 ++- db/dal/user_dal.py | 63 +++++++++++++++++++++++++++ db/migrator.py | 74 ++++++++++++++++++++++++++++++++ db/models.py | 1 + 9 files changed, 236 insertions(+), 16 deletions(-) diff --git a/.env.example b/.env.example index ebe2b6a..f1218a5 100644 --- a/.env.example +++ b/.env.example @@ -93,6 +93,7 @@ SUBSCRIPTION_NOTIFY_DAYS_BEFORE=3 # Da REFERRAL_ONE_BONUS_PER_REFEREE=False # Give a bonus only once per referee +LEGACY_REFS=true # Allow ref_ links. Leave unset/true unless you want to disable old links # Referral Bonus Days # Bonus for the inviting user REFERRAL_BONUS_DAYS_1_MONTH=3 diff --git a/bot/handlers/inline_mode.py b/bot/handlers/inline_mode.py index 762da06..145fbfb 100644 --- a/bot/handlers/inline_mode.py +++ b/bot/handlers/inline_mode.py @@ -38,7 +38,13 @@ async def inline_query_handler(inline_query: InlineQuery, # For all users: referral functionality if not query or "реф" in query or "ref" in query or "друг" in query or "friend" in query: referral_result = await create_referral_result( - inline_query, bot, referral_service, i18n, current_lang, settings + inline_query, + bot, + referral_service, + i18n, + current_lang, + settings, + session, ) if referral_result: results.append(referral_result) @@ -67,9 +73,15 @@ async def inline_query_handler(inline_query: InlineQuery, await inline_query.answer(results=[], cache_time=10) -async def create_referral_result(inline_query: InlineQuery, bot: Bot, - referral_service: ReferralService, - i18n_instance, lang: str, settings: Settings) -> Optional[InlineQueryResultArticle]: +async def create_referral_result( + inline_query: InlineQuery, + bot: Bot, + referral_service: ReferralService, + i18n_instance, + lang: str, + settings: Settings, + session: AsyncSession, +) -> Optional[InlineQueryResultArticle]: """Create referral link result for inline query""" _ = lambda key, **kwargs: i18n_instance.gettext(lang, key, **kwargs) @@ -80,7 +92,13 @@ async def create_referral_result(inline_query: InlineQuery, bot: Bot, return None user_id = inline_query.from_user.id - referral_link = referral_service.generate_referral_link(bot_username, user_id) + referral_link = await referral_service.generate_referral_link( + session, bot_username, user_id + ) + + if not referral_link: + logging.warning("Could not produce referral link for inline user %s", user_id) + return None # Create message content (use same text as friend message) message_text = _( diff --git a/bot/handlers/user/referral.py b/bot/handlers/user/referral.py index 7536370..06f3150 100644 --- a/bot/handlers/user/referral.py +++ b/bot/handlers/user/referral.py @@ -60,8 +60,18 @@ async def referral_command_handler(event: Union[types.Message, return inviter_user_id = event.from_user.id - referral_link = referral_service.generate_referral_link( - bot_username, inviter_user_id) + referral_link = await referral_service.generate_referral_link( + session, bot_username, inviter_user_id) + + if not referral_link: + logging.error( + "Failed to generate referral link for user %s (probably missing DB record).", + inviter_user_id, + ) + await target_message_obj.answer(_("error_generating_referral_link")) + if isinstance(event, types.CallbackQuery): + await event.answer() + return bonus_info_parts = [] if settings.subscription_options: @@ -132,7 +142,16 @@ async def referral_action_handler(callback: types.CallbackQuery, settings: Setti return inviter_user_id = callback.from_user.id - referral_link = referral_service.generate_referral_link(bot_username, inviter_user_id) + referral_link = await referral_service.generate_referral_link( + session, bot_username, inviter_user_id) + + if not referral_link: + logging.error( + "Failed to generate referral link for user %s via inline button.", + inviter_user_id, + ) + await callback.answer(_("error_generating_referral_link"), show_alert=True) + return friend_message = _("referral_friend_message", referral_link=referral_link) diff --git a/bot/handlers/user/start.py b/bot/handlers/user/start.py index ebfc2ed..1eb7551 100644 --- a/bot/handlers/user/start.py +++ b/bot/handlers/user/start.py @@ -302,7 +302,7 @@ async def ensure_required_channel_subscription( @router.message(CommandStart()) -@router.message(CommandStart(magic=F.args.regexp(r"^ref_(\d+)$").as_("ref_match"))) +@router.message(CommandStart(magic=F.args.regexp(r"^ref_((?:[uU][A-Za-z0-9]{9})|(?:[A-Za-z0-9]{9})|\d+)$").as_("ref_match"))) @router.message(CommandStart(magic=F.args.regexp(r"^promo_(\w+)$").as_("promo_match"))) @router.message(CommandStart(magic=F.args.regexp(r"^(?!ref_|promo_)([A-Za-z0-9_\-]{2,64})$").as_("ad_param_match"))) async def start_command_handler(message: types.Message, @@ -328,9 +328,23 @@ async def start_command_handler(message: types.Message, ad_start_param: Optional[str] = None if ref_match: - potential_referrer_id = int(ref_match.group(1)) - if await user_dal.get_user_by_id(session, potential_referrer_id): - referred_by_user_id = potential_referrer_id + raw_ref_value = ref_match.group(1) + if raw_ref_value.isdigit(): + if settings.LEGACY_REFS: + potential_referrer_id = int(raw_ref_value) + if potential_referrer_id != user_id and await user_dal.get_user_by_id( + session, potential_referrer_id): + referred_by_user_id = potential_referrer_id + else: + normalized_code = raw_ref_value.strip() + if normalized_code and normalized_code[0].lower() == "u": + normalized_code = normalized_code[1:] + ref_user = None + if normalized_code: + ref_user = await user_dal.get_user_by_referral_code( + session, normalized_code) + if ref_user and ref_user.user_id != user_id: + referred_by_user_id = ref_user.user_id elif promo_match: promo_code_to_apply = promo_match.group(1) logging.info(f"User {user_id} started with promo code: {promo_code_to_apply}") diff --git a/bot/services/referral_service.py b/bot/services/referral_service.py index 0fcaaf7..bbed41a 100644 --- a/bot/services/referral_service.py +++ b/bot/services/referral_service.py @@ -257,9 +257,35 @@ class ReferralService: raise - def generate_referral_link(self, bot_username: str, - inviter_user_id: int) -> str: - return f"https://t.me/{bot_username}?start=ref_{inviter_user_id}" + async def generate_referral_link(self, session: AsyncSession, + bot_username: str, + inviter_user_id: int) -> Optional[str]: + try: + user = await user_dal.get_user_by_id(session, inviter_user_id) + if not user: + logging.warning( + "Unable to generate referral link: user %s not found.", + inviter_user_id, + ) + return None + + referral_code = await user_dal.ensure_referral_code(session, user) + if not referral_code: + logging.warning( + "User %s has no referral code even after regeneration attempt.", + inviter_user_id, + ) + return None + + return f"https://t.me/{bot_username}?start=ref_u{referral_code}" + except Exception as exc: + logging.error( + "Failed to generate referral link for user %s: %s", + inviter_user_id, + exc, + exc_info=True, + ) + return None async def get_referral_stats(self, session: AsyncSession, user_id: int) -> dict: """Get referral statistics for a user""" diff --git a/config/settings.py b/config/settings.py index c296621..d86472e 100644 --- a/config/settings.py +++ b/config/settings.py @@ -118,7 +118,11 @@ class Settings(BaseSettings): # Referral program configuration REFERRAL_ONE_BONUS_PER_REFEREE: bool = Field( default=True, - description="When true, referral bonuses (for inviter and referee) are applied only once per invited user – on their first successful payment." + description="When true, referral bonuses (for inviter and referee) are applied only once per invited user - on their first successful payment." + ) + LEGACY_REFS: bool = Field( + default=True, + description="Allow legacy referral links like ref_ to continue working. Defaults to True when unset." ) PANEL_API_URL: Optional[str] = None diff --git a/db/dal/user_dal.py b/db/dal/user_dal.py index 4676a81..0622732 100644 --- a/db/dal/user_dal.py +++ b/db/dal/user_dal.py @@ -1,4 +1,6 @@ import logging +import secrets +import string from typing import Optional, List, Dict, Any, Tuple from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select @@ -18,6 +20,53 @@ from ..models import ( AdAttribution, ) +REFERRAL_CODE_ALPHABET = string.ascii_uppercase + string.digits +REFERRAL_CODE_LENGTH = 9 +MAX_REFERRAL_CODE_ATTEMPTS = 25 + + +def _generate_referral_code_candidate() -> str: + return "".join( + secrets.choice(REFERRAL_CODE_ALPHABET) for _ in range(REFERRAL_CODE_LENGTH) + ) + + +async def _referral_code_exists(session: AsyncSession, code: str) -> bool: + stmt = select(User.user_id).where(User.referral_code == code) + result = await session.execute(stmt) + return result.scalar_one_or_none() is not None + + +async def generate_unique_referral_code(session: AsyncSession) -> str: + """ + Generate a unique referral code consisting of uppercase alphanumeric characters. + Retries until a free code is found or raises RuntimeError after exceeding attempts. + """ + for _ in range(MAX_REFERRAL_CODE_ATTEMPTS): + candidate = _generate_referral_code_candidate() + if not await _referral_code_exists(session, candidate): + return candidate + raise RuntimeError("Failed to generate a unique referral code after several attempts.") + + +async def ensure_referral_code(session: AsyncSession, user: User) -> str: + """ + Ensure the provided user has a referral code, generating and persisting it if missing. + Returns the existing or newly generated code. + """ + if user.referral_code: + normalized = user.referral_code.strip().upper() + if normalized != user.referral_code: + user.referral_code = normalized + await session.flush() + await session.refresh(user) + return user.referral_code + + user.referral_code = await generate_unique_referral_code(session) + await session.flush() + await session.refresh(user) + return user.referral_code + async def get_user_by_id(session: AsyncSession, user_id: int) -> Optional[User]: stmt = select(User).where(User.user_id == user_id) @@ -52,6 +101,11 @@ async def create_user(session: AsyncSession, user_data: Dict[str, Any]) -> Tuple if "registration_date" not in user_data: user_data["registration_date"] = datetime.now(timezone.utc) + if not user_data.get("referral_code"): + user_data["referral_code"] = await generate_unique_referral_code(session) + else: + user_data["referral_code"] = user_data["referral_code"].strip().upper() + # Use PostgreSQL upsert to avoid IntegrityError on concurrent inserts stmt = ( pg_insert(User) @@ -80,6 +134,15 @@ async def create_user(session: AsyncSession, user_data: Dict[str, Any]) -> Tuple return user, created +async def get_user_by_referral_code(session: AsyncSession, referral_code: str) -> Optional[User]: + normalized = referral_code.strip().upper() + if not normalized: + return None + stmt = select(User).where(User.referral_code == normalized) + result = await session.execute(stmt) + return result.scalar_one_or_none() + + async def update_user( session: AsyncSession, user_id: int, update_data: Dict[str, Any] ) -> Optional[User]: diff --git a/db/migrator.py b/db/migrator.py index e91cb53..11c4ee0 100644 --- a/db/migrator.py +++ b/db/migrator.py @@ -48,12 +48,86 @@ def _migration_0001_add_channel_subscription_fields(connection: Connection) -> N connection.execute(text(stmt)) +def _migration_0002_add_referral_code(connection: Connection) -> None: + inspector = inspect(connection) + columns: Set[str] = {col["name"] for col in inspector.get_columns("users")} + + if "referral_code" not in columns: + connection.execute( + text("ALTER TABLE users ADD COLUMN referral_code VARCHAR(16)") + ) + + connection.execute( + text( + """ + WITH generated_codes AS ( + SELECT + user_id, + UPPER( + SUBSTRING( + md5( + user_id::text + || clock_timestamp()::text + || random()::text + ) + FROM 1 FOR 9 + ) + ) AS referral_code + FROM users + WHERE referral_code IS NULL OR referral_code = '' + ) + UPDATE users AS u + SET referral_code = g.referral_code + FROM generated_codes AS g + WHERE u.user_id = g.user_id + """ + ) + ) + + connection.execute( + text( + """ + CREATE UNIQUE INDEX IF NOT EXISTS uq_users_referral_code + ON users (referral_code) + WHERE referral_code IS NOT NULL + """ + ) + ) + + +def _migration_0003_normalize_referral_codes(connection: Connection) -> None: + inspector = inspect(connection) + columns: Set[str] = {col["name"] for col in inspector.get_columns("users")} + if "referral_code" not in columns: + return + + connection.execute( + text( + """ + UPDATE users + SET referral_code = UPPER(referral_code) + WHERE referral_code IS NOT NULL + AND referral_code <> UPPER(referral_code) + """ + ) + ) + MIGRATIONS: List[Migration] = [ Migration( id="0001_add_channel_subscription_fields", description="Add columns to track required channel subscription verification", upgrade=_migration_0001_add_channel_subscription_fields, ), + Migration( + id="0002_add_referral_code", + description="Store short referral codes for users and backfill existing rows", + upgrade=_migration_0002_add_referral_code, + ), + Migration( + id="0003_normalize_referral_codes", + description="Normalize referral codes to uppercase for consistent lookups", + upgrade=_migration_0003_normalize_referral_codes, + ), ] diff --git a/db/models.py b/db/models.py index 5e8d079..843300f 100644 --- a/db/models.py +++ b/db/models.py @@ -21,6 +21,7 @@ class User(Base): server_default=func.now()) is_banned = Column(Boolean, default=False) panel_user_uuid = Column(String, nullable=True, unique=True, index=True) + referral_code = Column(String(16), nullable=True, unique=True, index=True) referred_by_id = Column(BigInteger, ForeignKey("users.user_id"), nullable=True)