social-post-writer-bot/payments.py

"""Реальная оплата Telegram Stars (XTR) за подписку.

Цикл: sendInvoice -> pre_checkout_query -> answerPreCheckoutQuery ->
successful_payment -> activate_subscription.
`telegram_payment_charge_id` каждого платежа пишется в таблицу payments,
поэтому один платёж не может активировать подписку дважды, а по нему потом
можно сделать возврат через refundStarPayment.
"""
from __future__ import annotations

import asyncio
import datetime
import logging
import os
from contextlib import closing

from aiogram import F, Router
from aiogram.types import LabeledPrice, Message, PreCheckoutQuery, SuccessfulPayment

from user_store import (
    SUBSCRIPTION_PLANS,
    activate_subscription_days,
    activate_subscription,
    connect,
    ensure_subscription,
    subscription_row,
    subscription_status_text,
)

logger = logging.getLogger(__name__)
router = Router(name="payments")

STARS_CURRENCY = "XTR"
INVOICE_PREFIX = "sub:"


def plan_for_payload(payload: str) -> tuple[int, int] | None:
    if payload.startswith("sub30:"):
        try:
            timestamp = payload.split(":", 1)[1]
            datetime.datetime.strptime(timestamp, "%Y%m%d%H%M%S")
        except (ValueError, IndexError):
            return None
        return (30, SUBSCRIPTION_PLANS[30])
    try:
        prefix, days_part, timestamp = payload.split(":", 2)
        if prefix != "sub" or not timestamp.isdigit():
            return None
        days = int(days_part)
    except (ValueError, IndexError):
        return None
    price = SUBSCRIPTION_PLANS.get(days)
    if price is None or (days == 0 and price != 1):
        return None
    return days, price


def _db_path(db_path: str | None = None) -> str | None:
    return db_path if db_path is not None else os.getenv("BOT_DB_PATH")


def init_payments_table(db_path: str | None = None) -> None:
    with closing(connect(db_path)) as connection, connection:
        connection.execute(
            """
            CREATE TABLE IF NOT EXISTS payments (
                telegram_payment_charge_id TEXT PRIMARY KEY,
                telegram_user_id INTEGER NOT NULL,
                invoice_payload TEXT NOT NULL,
                currency TEXT NOT NULL,
                total_amount INTEGER NOT NULL,
                created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
            )
            """
        )
        connection.execute(
            "CREATE INDEX IF NOT EXISTS idx_payments_user ON payments(telegram_user_id)"
        )


def payment_recorded(charge_id: str, db_path: str | None = None) -> bool:
    with closing(connect(db_path)) as connection:
        row = connection.execute(
            "SELECT 1 FROM payments WHERE telegram_payment_charge_id = ?", (charge_id,)
        ).fetchone()
    return row is not None


def save_payment(
    charge_id: str,
    user_id: int,
    invoice_payload: str,
    currency: str,
    total_amount: int,
    db_path: str | None = None,
) -> None:
    with closing(connect(db_path)) as connection, connection:
        connection.execute(
            """
            INSERT OR IGNORE INTO payments
                (telegram_payment_charge_id, telegram_user_id, invoice_payload, currency, total_amount)
            VALUES (?, ?, ?, ?, ?)
            """,
            (charge_id, user_id, invoice_payload, currency, total_amount),
        )


async def send_subscription_invoice(
    bot, chat_id: int, thread_id: int | None = None, days: int = 30
) -> None:
    """Отправляет выбранный инвойс на разовую подписку в Stars."""
    price = SUBSCRIPTION_PLANS[days]
    payload = f"{INVOICE_PREFIX}{days}:{datetime.datetime.now(datetime.timezone.utc):%Y%m%d%H%M%S}"
    kwargs: dict[str, object] = {
        "chat_id": chat_id,
        "title": f"Подписка на {days} дней",
        "description": "Доступ к боту по генерации текста постов",
        "payload": payload,
        "currency": STARS_CURRENCY,
        "prices": [LabeledPrice(label=f"{days} дней", amount=price)],
        "provider_token": "",
    }
    if thread_id is not None:
        kwargs["message_thread_id"] = thread_id
    await bot.send_invoice(**kwargs)


@router.pre_checkout_query()
async def handle_pre_checkout(query: PreCheckoutQuery) -> None:
    """Подтверждаем только свой инвойс, в правильной валюте и по правильной цене."""
    plan = plan_for_payload(query.invoice_payload)
    if plan is None or query.currency != STARS_CURRENCY or query.total_amount != plan[1]:
        await query.answer(ok=False, error_message="Некорректная сумма или валюта платежа.")
        return
    await query.answer(ok=True)


@router.message(F.successful_payment)
async def handle_successful_payment(message: Message) -> None:
    payment = message.successful_payment
    if payment is None or message.from_user is None:
        return

    charge_id = payment.telegram_payment_charge_id
    plan = plan_for_payload(payment.invoice_payload)
    if plan is None or payment.currency != STARS_CURRENCY or payment.total_amount != plan[1]:
        await _reject(message, charge_id, "неверная сумма или валюта")
        return
    days, _price = plan

    db_path = _db_path()
    if await asyncio.to_thread(payment_recorded, charge_id, db_path):
        await message.answer("Этот платёж уже учтён — подписка не продлена повторно.")
        return

    if days == 0:
        days = 1

    await asyncio.to_thread(ensure_subscription, message.from_user.id)
    await asyncio.to_thread(activate_subscription_days, message.from_user.id, days)
    sub = await asyncio.to_thread(subscription_row, message.from_user.id, db_path)
    await asyncio.to_thread(
        save_payment,
        charge_id,
        message.from_user.id,
        payment.invoice_payload,
        payment.currency,
        payment.total_amount,
        db_path,
    )
    await message.answer(subscription_status_text(sub))


async def _reject(message: Message, charge_id: str, reason: str) -> None:
    logger.error(
        "Rejected payment charge=%s user=%s reason=%s",
        charge_id,
        message.from_user.id if message.from_user else None,
        reason,
    )
    await message.answer(
        "Платёж получен, но активировать подписку не получилось. Напиши, пожалуйста, поддержке."
    )