# - *- coding: utf- 8 - *- from math import isfinite import re from urllib.parse import urlparse from uuid import uuid4 from sqlalchemy import BigInteger, Float, Integer, String, UniqueConstraint, text from sqlalchemy.orm import Mapped, mapped_column from sqlalchemy.ext.asyncio import AsyncConnection from typing import Any, Dict, Optional, Tuple, Union from tgbot.database.core import Base, session_factory from tgbot.database.entities import Referral, ReferralTransaction, ReferralWithdrawal from tgbot.database.repository import BaseRepository from tgbot.utils.const_functions import get_unix def is_valid_referral_withdrawal_recipient( withdrawal_method: str, recipient: str ) -> bool: if withdrawal_method == "Cryptobot": return bool(re.fullmatch(r"@[A-Za-z][A-Za-z0-9_]{4,31}", recipient)) if withdrawal_method == "Lolzteam": parsed_url = urlparse(recipient) return parsed_url.scheme in ("http", "https") and parsed_url.hostname in ( "zelenka.guru", "lolz.team", "lolz.live", "lolz.guru", ) return False async def ensure_referrals_schema(conn: AsyncConnection) -> None: user_columns_result = await conn.exec_driver_sql("PRAGMA table_info(storage_users)") user_columns = {column[1] for column in user_columns_result.fetchall()} # Обновление полей пользователя if user_columns: if "user_referrer_id" not in user_columns: await conn.exec_driver_sql( "ALTER TABLE storage_users ADD COLUMN user_referrer_id BIGINT" ) if "user_referral_balance" not in user_columns: await conn.exec_driver_sql( "ALTER TABLE storage_users " "ADD COLUMN user_referral_balance FLOAT NOT NULL DEFAULT 0" ) if "user_referral_hold" not in user_columns: await conn.exec_driver_sql( "ALTER TABLE storage_users " "ADD COLUMN user_referral_hold FLOAT NOT NULL DEFAULT 0" ) await conn.exec_driver_sql( "CREATE INDEX IF NOT EXISTS ix_storage_users_user_referrer_id " "ON storage_users (user_referrer_id)" ) settings_columns_result = await conn.exec_driver_sql( "PRAGMA table_info(storage_settings)" ) settings_columns = {column[1] for column in settings_columns_result.fetchall()} # Обновление полей настроек if settings_columns: if "status_referral" not in settings_columns: await conn.exec_driver_sql( "ALTER TABLE storage_settings " "ADD COLUMN status_referral VARCHAR(16) NOT NULL DEFAULT 'False'" ) if "referral_bonus_rub" not in settings_columns: await conn.exec_driver_sql( "ALTER TABLE storage_settings " "ADD COLUMN referral_bonus_rub FLOAT NOT NULL DEFAULT 0" ) if "referral_refill_percent" not in settings_columns: await conn.exec_driver_sql( "ALTER TABLE storage_settings " "ADD COLUMN referral_refill_percent FLOAT NOT NULL DEFAULT 0" ) class ReferralModel(Base): __tablename__ = "storage_referrals" increment: Mapped[int] = mapped_column( Integer, primary_key=True, autoincrement=True ) referrer_id: Mapped[int] = mapped_column(BigInteger, nullable=False, index=True) referral_id: Mapped[int] = mapped_column( BigInteger, nullable=False, unique=True, index=True ) referral_unix: Mapped[int] = mapped_column( Integer, nullable=False, default=get_unix ) class ReferralTransactionModel(Base): __tablename__ = "storage_referral_transactions" __table_args__ = ( UniqueConstraint( "user_id", "transaction_type", "source_id", name="uq_referral_transaction_source", ), ) increment: Mapped[int] = mapped_column( Integer, primary_key=True, autoincrement=True ) user_id: Mapped[int] = mapped_column(BigInteger, nullable=False, index=True) related_user_id: Mapped[Optional[int]] = mapped_column(BigInteger, nullable=True) transaction_type: Mapped[str] = mapped_column(String(64), nullable=False) source_id: Mapped[str] = mapped_column(String(128), nullable=False) amount: Mapped[float] = mapped_column(Float, nullable=False, default=0) transaction_unix: Mapped[int] = mapped_column( Integer, nullable=False, default=get_unix ) class ReferralWithdrawalModel(Base): __tablename__ = "storage_referral_withdrawals" increment: Mapped[int] = mapped_column( Integer, primary_key=True, autoincrement=True ) user_id: Mapped[int] = mapped_column(BigInteger, nullable=False, index=True) withdrawal_amount: Mapped[float] = mapped_column(Float, nullable=False, default=0) withdrawal_method: Mapped[str] = mapped_column(String(32), nullable=False) withdrawal_recipient: Mapped[str] = mapped_column(String(255), nullable=False) withdrawal_status: Mapped[str] = mapped_column( String(32), nullable=False, default="pending" ) withdrawal_unix: Mapped[int] = mapped_column( Integer, nullable=False, default=get_unix ) processed_unix: Mapped[Optional[int]] = mapped_column(Integer, nullable=True) processed_admin_id: Mapped[Optional[int]] = mapped_column(BigInteger, nullable=True) class Referralx(BaseRepository[ReferralModel, Referral]): def __init__(self): super().__init__() self.table_model = ReferralModel self.entity_model = Referral self.storage_name = ReferralModel.__tablename__ async def add(self, referrer_id: int, referral_id: int) -> Referral: return await self._insert( referrer_id=referrer_id, referral_id=referral_id, referral_unix=get_unix(), ) async def register_referral(self, referral_id: int, referrer_id: int) -> str: if referral_id == referrer_id: return "SELF_REFERRAL" async with session_factory() as session: try: await session.execute(text("BEGIN IMMEDIATE")) settings_result = await session.execute( text( "SELECT status_referral, referral_bonus_rub " "FROM storage_settings WHERE id = 1" ) ) settings = settings_result.mappings().first() if settings is None or settings["status_referral"] != "True": await session.rollback() return "DISABLED" referrer_result = await session.execute( text( "SELECT user_referral_balance FROM storage_users " "WHERE user_id = :user_id" ), {"user_id": referrer_id}, ) referrer = referrer_result.mappings().first() if referrer is None: await session.rollback() return "REFERRER_NOT_FOUND" referral_result = await session.execute( text( "SELECT user_balance, user_referrer_id " "FROM storage_users WHERE user_id = :user_id" ), {"user_id": referral_id}, ) referral = referral_result.mappings().first() if referral is None: await session.rollback() return "REFERRAL_NOT_FOUND" existing_referral_result = await session.execute( text( "SELECT increment FROM storage_referrals " "WHERE referral_id = :user_id LIMIT 1" ), {"user_id": referral_id}, ) if ( referral["user_referrer_id"] is not None or existing_referral_result.mappings().first() is not None ): await session.rollback() return "ALREADY" bonus = round(float(settings["referral_bonus_rub"]), 2) source_id = f"registration:{referral_id}" now_unix = get_unix() # Приглашённому — бонус на основной баланс. await session.execute( text( "INSERT INTO storage_referrals " "(referrer_id, referral_id, referral_unix) " "VALUES (:referrer_id, :referral_id, :referral_unix)" ), { "referrer_id": referrer_id, "referral_id": referral_id, "referral_unix": now_unix, }, ) await session.execute( text( "UPDATE storage_users " "SET user_referrer_id = :referrer_id, " "user_balance = :balance " "WHERE user_id = :user_id" ), { "referrer_id": referrer_id, "balance": round(float(referral["user_balance"]) + bonus, 2), "user_id": referral_id, }, ) # Рефереру — бонус на реферальный баланс. await session.execute( text( "UPDATE storage_users " "SET user_referral_balance = :balance " "WHERE user_id = :user_id" ), { "balance": round( float(referrer["user_referral_balance"]) + bonus, 2 ), "user_id": referrer_id, }, ) await session.execute( text( "INSERT INTO storage_referral_transactions " "(user_id, related_user_id, transaction_type, source_id, amount, transaction_unix) " "VALUES (:user_id, :related_user_id, :transaction_type, :source_id, :amount, :transaction_unix)" ), { "user_id": referral_id, "related_user_id": referrer_id, "transaction_type": "registration_bonus_main", "source_id": source_id, "amount": bonus, "transaction_unix": now_unix, }, ) await session.execute( text( "INSERT INTO storage_referral_transactions " "(user_id, related_user_id, transaction_type, source_id, amount, transaction_unix) " "VALUES (:user_id, :related_user_id, :transaction_type, :source_id, :amount, :transaction_unix)" ), { "user_id": referrer_id, "related_user_id": referral_id, "transaction_type": "registration_bonus_referral", "source_id": source_id, "amount": bonus, "transaction_unix": now_unix, }, ) await session.commit() except Exception: await session.rollback() raise return "ok" async def transfer_to_main_balance(self, user_id: int, amount: float) -> str: try: transfer_amount = round(float(amount), 2) except (TypeError, ValueError): return "INVALID_AMOUNT" if not isfinite(transfer_amount) or transfer_amount <= 0: return "INVALID_AMOUNT" async with session_factory() as session: try: await session.execute(text("BEGIN IMMEDIATE")) user_result = await session.execute( text( "SELECT user_balance, user_referral_balance, user_referral_hold " "FROM storage_users WHERE user_id = :user_id" ), {"user_id": user_id}, ) get_user = user_result.mappings().first() if get_user is None: await session.rollback() return "USER_NOT_FOUND" available_balance = round( float(get_user["user_referral_balance"]) - float(get_user["user_referral_hold"]), 2, ) if transfer_amount > available_balance: await session.rollback() return "INSUFFICIENT_FUNDS" await session.execute( text( "UPDATE storage_users " "SET user_balance = :user_balance, " "user_referral_balance = :referral_balance " "WHERE user_id = :user_id" ), { "user_balance": round( float(get_user["user_balance"]) + transfer_amount, 2, ), "referral_balance": round( float(get_user["user_referral_balance"]) - transfer_amount, 2, ), "user_id": user_id, }, ) await session.execute( text( "INSERT INTO storage_referral_transactions " "(user_id, related_user_id, transaction_type, source_id, amount, transaction_unix) " "VALUES (:user_id, NULL, :transaction_type, :source_id, :amount, :transaction_unix)" ), { "user_id": user_id, "transaction_type": "transfer_to_main", "source_id": f"transfer:{user_id}:{uuid4().hex}", "amount": transfer_amount, "transaction_unix": get_unix(), }, ) await session.commit() except Exception: await session.rollback() raise return "ok" class ReferralTransactionx( BaseRepository[ReferralTransactionModel, ReferralTransaction] ): def __init__(self): super().__init__() self.table_model = ReferralTransactionModel self.entity_model = ReferralTransaction self.storage_name = ReferralTransactionModel.__tablename__ async def get_refill_reward(self, source_id: str) -> Optional[ReferralTransaction]: return await self.get( transaction_type="refill_percent", source_id=source_id, ) async def add( self, user_id: int, transaction_type: str, source_id: str, amount: float, related_user_id: Optional[int] = None, ) -> ReferralTransaction: return await self._insert( user_id=user_id, related_user_id=related_user_id, transaction_type=transaction_type, source_id=source_id, amount=round(amount, 2), transaction_unix=get_unix(), ) class ReferralWithdrawalx(BaseRepository[ReferralWithdrawalModel, ReferralWithdrawal]): def __init__(self): super().__init__() self.table_model = ReferralWithdrawalModel self.entity_model = ReferralWithdrawal self.storage_name = ReferralWithdrawalModel.__tablename__ async def create_pending( self, user_id: int, withdrawal_amount: float, withdrawal_method: str, withdrawal_recipient: str, ) -> Tuple[str, Optional[ReferralWithdrawal]]: try: amount = round(float(withdrawal_amount), 2) except (TypeError, ValueError): return "INVALID_AMOUNT", None recipient = str(withdrawal_recipient).strip() if not isfinite(amount) or amount <= 0: return "INVALID_AMOUNT", None if withdrawal_method not in ("Cryptobot", "Lolzteam"): return "INVALID_METHOD", None if not is_valid_referral_withdrawal_recipient(withdrawal_method, recipient): return "INVALID_RECIPIENT", None withdrawal_id: Optional[int] = None async with session_factory() as session: try: await session.execute(text("BEGIN IMMEDIATE")) user_result = await session.execute( text( "SELECT user_referral_balance, user_referral_hold " "FROM storage_users WHERE user_id = :user_id" ), {"user_id": user_id}, ) get_user = user_result.mappings().first() if get_user is None: await session.rollback() return "USER_NOT_FOUND", None available_balance = round( float(get_user["user_referral_balance"]) - float(get_user["user_referral_hold"]), 2, ) if amount > available_balance: await session.rollback() return "INSUFFICIENT_FUNDS", None now_unix = get_unix() withdrawal_result = await session.execute( text( "INSERT INTO storage_referral_withdrawals " "(user_id, withdrawal_amount, withdrawal_method, withdrawal_recipient, withdrawal_status, withdrawal_unix) " "VALUES (:user_id, :amount, :method, :recipient, 'pending', :unix)" ), { "user_id": user_id, "amount": amount, "method": withdrawal_method, "recipient": recipient, "unix": now_unix, }, ) withdrawal_id = int(withdrawal_result.lastrowid) await session.execute( text( "UPDATE storage_users " "SET user_referral_hold = :hold " "WHERE user_id = :user_id" ), { "hold": round( float(get_user["user_referral_hold"]) + amount, 2, ), "user_id": user_id, }, ) await session.execute( text( "INSERT INTO storage_referral_transactions " "(user_id, related_user_id, transaction_type, source_id, amount, transaction_unix) " "VALUES (:user_id, NULL, :transaction_type, :source_id, :amount, :unix)" ), { "user_id": user_id, "transaction_type": "withdrawal_hold", "source_id": f"withdrawal:{withdrawal_id}", "amount": amount, "unix": now_unix, }, ) await session.commit() except Exception: await session.rollback() raise return "ok", await self.get(increment=withdrawal_id) async def complete( self, withdrawal_id: int, admin_id: int ) -> Tuple[str, Optional[ReferralWithdrawal]]: async with session_factory() as session: try: await session.execute(text("BEGIN IMMEDIATE")) withdrawal_result = await session.execute( text( "SELECT user_id, withdrawal_amount, withdrawal_status " "FROM storage_referral_withdrawals WHERE increment = :withdrawal_id" ), {"withdrawal_id": withdrawal_id}, ) withdrawal = withdrawal_result.mappings().first() if withdrawal is None: await session.rollback() return "NOT_FOUND", None if withdrawal["withdrawal_status"] != "pending": await session.rollback() return "ALREADY_PROCESSED", None user_result = await session.execute( text( "SELECT user_referral_balance, user_referral_hold " "FROM storage_users WHERE user_id = :user_id" ), {"user_id": withdrawal["user_id"]}, ) get_user = user_result.mappings().first() amount = float(withdrawal["withdrawal_amount"]) if ( get_user is None or float(get_user["user_referral_balance"]) < amount or float(get_user["user_referral_hold"]) < amount ): await session.rollback() return "BALANCE_ERROR", None now_unix = get_unix() await session.execute( text( "UPDATE storage_users " "SET user_referral_balance = :balance, user_referral_hold = :hold " "WHERE user_id = :user_id" ), { "balance": round( float(get_user["user_referral_balance"]) - amount, 2, ), "hold": round( float(get_user["user_referral_hold"]) - amount, 2, ), "user_id": withdrawal["user_id"], }, ) await session.execute( text( "UPDATE storage_referral_withdrawals " "SET withdrawal_status = 'completed', processed_unix = :unix, processed_admin_id = :admin_id " "WHERE increment = :withdrawal_id" ), { "unix": now_unix, "admin_id": admin_id, "withdrawal_id": withdrawal_id, }, ) await session.execute( text( "INSERT INTO storage_referral_transactions " "(user_id, related_user_id, transaction_type, source_id, amount, transaction_unix) " "VALUES (:user_id, NULL, :transaction_type, :source_id, :amount, :unix)" ), { "user_id": withdrawal["user_id"], "transaction_type": "withdrawal_completed", "source_id": f"withdrawal:{withdrawal_id}", "amount": amount, "unix": now_unix, }, ) await session.commit() except Exception: await session.rollback() raise return "ok", await self.get(increment=withdrawal_id) async def reject( self, withdrawal_id: int, admin_id: int ) -> Tuple[str, Optional[ReferralWithdrawal]]: async with session_factory() as session: try: await session.execute(text("BEGIN IMMEDIATE")) withdrawal_result = await session.execute( text( "SELECT user_id, withdrawal_amount, withdrawal_status " "FROM storage_referral_withdrawals WHERE increment = :withdrawal_id" ), {"withdrawal_id": withdrawal_id}, ) withdrawal = withdrawal_result.mappings().first() if withdrawal is None: await session.rollback() return "NOT_FOUND", None if withdrawal["withdrawal_status"] != "pending": await session.rollback() return "ALREADY_PROCESSED", None user_result = await session.execute( text( "SELECT user_referral_hold FROM storage_users " "WHERE user_id = :user_id" ), {"user_id": withdrawal["user_id"]}, ) get_user = user_result.mappings().first() amount = float(withdrawal["withdrawal_amount"]) if get_user is None or float(get_user["user_referral_hold"]) < amount: await session.rollback() return "BALANCE_ERROR", None now_unix = get_unix() await session.execute( text( "UPDATE storage_users " "SET user_referral_hold = :hold WHERE user_id = :user_id" ), { "hold": round( float(get_user["user_referral_hold"]) - amount, 2, ), "user_id": withdrawal["user_id"], }, ) await session.execute( text( "UPDATE storage_referral_withdrawals " "SET withdrawal_status = 'rejected', processed_unix = :unix, processed_admin_id = :admin_id " "WHERE increment = :withdrawal_id" ), { "unix": now_unix, "admin_id": admin_id, "withdrawal_id": withdrawal_id, }, ) await session.execute( text( "INSERT INTO storage_referral_transactions " "(user_id, related_user_id, transaction_type, source_id, amount, transaction_unix) " "VALUES (:user_id, NULL, :transaction_type, :source_id, :amount, :unix)" ), { "user_id": withdrawal["user_id"], "transaction_type": "withdrawal_rejected", "source_id": f"withdrawal:{withdrawal_id}", "amount": amount, "unix": now_unix, }, ) await session.commit() except Exception: await session.rollback() raise return "ok", await self.get(increment=withdrawal_id) async def add( self, user_id: int, withdrawal_amount: float, withdrawal_method: str, withdrawal_recipient: str, ) -> ReferralWithdrawal: return await self._insert( user_id=user_id, withdrawal_amount=round(withdrawal_amount, 2), withdrawal_method=withdrawal_method, withdrawal_recipient=withdrawal_recipient, withdrawal_status="pending", withdrawal_unix=get_unix(), ) async def update( self, where: Optional[Union[Dict[str, Any], int]] = None, **kwargs, ) -> int: if isinstance(where, int): where = {"increment": where} return await self._update(where=where, **kwargs)