Files

729 lines
28 KiB
Python

# - *- 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)