258 lines
11 KiB
Python
258 lines
11 KiB
Python
from decimal import Decimal
|
||
import json
|
||
from instance.model import User, Subscription, Transaction
|
||
from app.services.billing_service import BillingAdapter
|
||
from app.services.marzban import MarzbanService, MarzbanUser
|
||
from .postgres_rep import PostgresRepository
|
||
from instance.model import Transaction,TransactionType, Plan
|
||
from dateutil.relativedelta import relativedelta
|
||
from datetime import datetime, timezone
|
||
import random
|
||
import string
|
||
from typing import Optional
|
||
import logging
|
||
from uuid import UUID
|
||
|
||
|
||
class DatabaseManager:
|
||
def __init__(self, session_generator,marzban_username,marzban_password,marzban_url,billing_base_url):
|
||
"""
|
||
Инициализация с асинхронным генератором сессий (например, get_postgres_session).
|
||
"""
|
||
self.logger = logging.getLogger(__name__)
|
||
self.postgres_repo = PostgresRepository(session_generator, self.logger)
|
||
self.marzban_service = MarzbanService(marzban_url,marzban_username,marzban_password)
|
||
self.billing_adapter = BillingAdapter(billing_base_url)
|
||
|
||
async def create_user(self, telegram_id: int, invented_by: Optional[int]= None):
|
||
"""
|
||
Создаёт пользователя.
|
||
"""
|
||
try:
|
||
username = self.generate_string(6)
|
||
return await self.postgres_repo.create_user(telegram_id, username, invented_by)
|
||
except Exception as e:
|
||
self.logger.error(f"Ошибка при создании пользователя:{e}")
|
||
|
||
async def get_user_by_telegram_id(self, telegram_id: int):
|
||
"""
|
||
Возвращает пользователя по Telegram ID.
|
||
"""
|
||
return await self.postgres_repo.get_user_by_telegram_id(telegram_id)
|
||
|
||
async def add_transaction(self, telegram_id: int, amount: float):
|
||
"""
|
||
Добавляет транзакцию.
|
||
"""
|
||
tran = Transaction(
|
||
user_id=telegram_id,
|
||
amount=Decimal(amount),
|
||
type=TransactionType.DEPOSIT
|
||
)
|
||
return await self.postgres_repo.add_record(tran)
|
||
async def add_referal(self,referrer_id: int, new_user_telegram_id: int):
|
||
"""
|
||
Добавление рефералу пользователей
|
||
"""
|
||
return await self.postgres_repo.add_referral(referrer_id,new_user_telegram_id)
|
||
async def get_transaction(self, telegram_id: int, limit: int = 10):
|
||
"""
|
||
Возвращает транзакции.
|
||
"""
|
||
return await self.postgres_repo.get_last_transactions(telegram_id, limit)
|
||
|
||
async def get_referrals_count(self,telegram_id: int) -> int:
|
||
"""
|
||
Docstring for get_referrals_count
|
||
|
||
:param self: Description
|
||
:param telegram_id: Description
|
||
:type telegram_id: int
|
||
:return: Description
|
||
:rtype: int
|
||
"""
|
||
return await self.postgres_repo.get_referrals_count(telegram_id)
|
||
|
||
# async def update_balance(self, telegram_id: int, amount: float):
|
||
# """
|
||
# Обновляет баланс пользователя и добавляет транзакцию.
|
||
# """
|
||
# self.logger.info(f"Попытка обновления баланса: telegram_id={telegram_id}, amount={amount}")
|
||
# user = await self.get_user_by_telegram_id(telegram_id)
|
||
# if not user:
|
||
# self.logger.warning(f"Пользователь с Telegram ID {telegram_id} не найден.")
|
||
# return "ERROR"
|
||
|
||
# updated = await self.postgres_repo.update_balance(user, amount)
|
||
# if not updated:
|
||
# self.logger.error(f"Не удалось обновить баланс пользователя {telegram_id}")
|
||
# return "ERROR"
|
||
|
||
# self.logger.info(f"Баланс пользователя {telegram_id} обновлен на {amount}, добавление транзакции")
|
||
# await self.add_transaction(user.telegram_id, amount)
|
||
# return "OK"
|
||
|
||
|
||
async def get_active_subscription(self, telegram_id: int):
|
||
"""
|
||
Проверяет наличие активной подписки.
|
||
"""
|
||
try:
|
||
|
||
return await self.postgres_repo.get_active_subscription(telegram_id)
|
||
except Exception as e:
|
||
self.logger.error(f"Неожиданная ошибка в get_active_subscription: {str(e)}")
|
||
return "ERROR"
|
||
|
||
async def get_plan_by_id(self, plan_id):
|
||
"""
|
||
Ищет по названию плана.
|
||
"""
|
||
try:
|
||
return await self.postgres_repo.get_plan_by_id(plan_id)
|
||
except Exception as e:
|
||
self.logger.error(f"Неожиданная ошибка в get_plan_by_name: {str(e)}")
|
||
return None
|
||
|
||
async def get_last_subscriptions(self, telegram_id: int, limit: int = 1):
|
||
"""
|
||
Возвращает список последних подписок.
|
||
"""
|
||
return await self.postgres_repo.get_last_subscription_by_user_id(telegram_id)
|
||
|
||
async def buy_sub(self, telegram_id: int, plan_name: str):
|
||
"""
|
||
Покупка подписки: сначала создаем подписку, потом списываем деньги
|
||
"""
|
||
try:
|
||
self.logger.info(f"Покупка подписки: user={telegram_id}, plan={plan_name}")
|
||
|
||
# 1. Проверка активной подписки
|
||
if await self.get_active_subscription(telegram_id):
|
||
return "ACTIVE_SUBSCRIPTION_EXISTS"
|
||
|
||
# 2. Получаем план
|
||
plan = await self.postgres_repo.get_subscription_plan(plan_name)
|
||
if not plan:
|
||
return "TARIFF_NOT_FOUND"
|
||
|
||
# 3. Проверяем пользователя
|
||
user = await self.get_user_by_telegram_id(telegram_id)
|
||
if not user:
|
||
return "USER_NOT_FOUND"
|
||
|
||
# 4. Проверяем баланс (только для информации)
|
||
balance_result = await self.billing_adapter.get_balance(telegram_id)
|
||
if balance_result["status"] == "error":
|
||
return "BILLING_SERVICE_ERROR"
|
||
|
||
if balance_result["balance"] < plan.price:
|
||
return "INSUFFICIENT_FUNDS"
|
||
|
||
# 5. СОЗДАЕМ ПОДПИСКУ (самое важное - сначала!)
|
||
new_subscription = await self._create_subscription_and_add_client(user, plan)
|
||
if not new_subscription:
|
||
return "SUBSCRIPTION_CREATION_FAILED"
|
||
|
||
# 6. ТОЛЬКО ПОСЛЕ УСПЕШНОГО СОЗДАНИЯ ПОДПИСКИ - списываем деньги
|
||
withdraw_result = await self.billing_adapter.withdraw_funds(
|
||
telegram_id,
|
||
float(plan.price),
|
||
f"Оплата подписки {plan_name}"
|
||
)
|
||
|
||
if withdraw_result["status"] == "error":
|
||
await self.postgres_repo.delete_subscription(new_subscription.id)
|
||
self.logger.error(f"Payment failed but subscription created: {new_subscription.id}")
|
||
return "PAYMENT_FAILED_AFTER_SUBSCRIPTION"
|
||
|
||
# 7. ВСЕ УСПЕШНО
|
||
self.logger.info(f"Подписка успешно создана и оплачена: {new_subscription.id}")
|
||
return {"status": "OK", "subscription_id": str(new_subscription.id)}
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"Ошибка в buy_sub: {str(e)}")
|
||
return "ERROR"
|
||
|
||
async def _create_subscription_and_add_client(self, user: User, plan: Plan):
|
||
"""Создаёт подписку и добавляет клиента на сервер."""
|
||
try:
|
||
self.logger.info(f"Создание подписки для user_id={user.telegram_id}, plan={plan.name}")
|
||
|
||
# Проверяем типы объектов
|
||
self.logger.info(f"Тип user: {type(user)}, тип plan: {type(plan)}")
|
||
|
||
expiry_date = datetime.utcnow() + relativedelta(days=plan.duration_days)
|
||
|
||
new_subscription = Subscription(
|
||
user_id=user.telegram_id,
|
||
vpn_server_id="BASE SERVER NEED TO UPDATE",
|
||
plan_id=plan.id,
|
||
end_date=expiry_date,
|
||
start_date=datetime.utcnow()
|
||
)
|
||
|
||
self.logger.info(f"Создан объект подписки: {new_subscription}")
|
||
|
||
response = await self.marzban_service.create_user(user, new_subscription)
|
||
self.logger.info(f"Ответ от Marzban: {response}")
|
||
|
||
if response == "USER_ALREADY_EXISTS":
|
||
response = await self.marzban_service.get_user_status(user)
|
||
result = await self.marzban_service.update_user(user, new_subscription)
|
||
|
||
# if not isinstance(response,MarzbanUser) or not isinstance(result,MarzbanUser):
|
||
# self.logger.error(f"Ошибка при добавлении клиента: {response}, {result}")
|
||
# return None
|
||
|
||
await self.postgres_repo.add_record(new_subscription)
|
||
self.logger.info(f"Подписка сохранена в БД с ID: {new_subscription.id}")
|
||
return new_subscription
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"Неожиданная ошибка в _create_subscription_and_add_client: {str(e)}")
|
||
import traceback
|
||
self.logger.error(f"Трассировка: {traceback.format_exc()}")
|
||
return None
|
||
|
||
async def generate_uri(self, telegram_id: int):
|
||
"""
|
||
Генерация URI для пользователя.
|
||
|
||
:param telegram_id: Telegram ID пользователя.
|
||
:return: Строка URI или None в случае ошибки.
|
||
"""
|
||
try:
|
||
user = await self.get_user_by_telegram_id(telegram_id)
|
||
if user == False or user == None:
|
||
self.logger.error(f"Ошибка при получении клиента: user = {user}")
|
||
return "ERROR"
|
||
|
||
result = await self.marzban_service.get_config_links(user)
|
||
if result == None:
|
||
self.logger.error(f"Ошибка при получении ссылки клиента: result = {user}")
|
||
return "ERROR"
|
||
|
||
self.logger.info(f"Итог generate_uri: result = {result}")
|
||
|
||
return result
|
||
except Exception as e:
|
||
self.logger.error(f"Неожиданная ошибка в generate_uri: {str(e)}")
|
||
return "ERROR"
|
||
|
||
|
||
@staticmethod
|
||
def generate_string(length):
|
||
"""
|
||
Генерирует случайную строку заданной длины.
|
||
"""
|
||
return ''.join(random.choices(string.ascii_lowercase + string.digits, k=length))
|
||
|
||
@staticmethod
|
||
def _is_subscription_expired(expire_timestamp: int) -> bool:
|
||
"""Проверяет, истекла ли подписка"""
|
||
|
||
current_time = datetime.now(timezone.utc)
|
||
expire_time = datetime.fromtimestamp(expire_timestamp, tz=timezone.utc)
|
||
|
||
return expire_time < current_time |