payment

src/database/model.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 # coding: utf-8
+import logging
 import math
 import os
 from contextlib import contextmanager
@@ -21,7 +22,12 @@ from sqlalchemy.orm import relationship, sessionmaker
 
 from config import FREE_DOWNLOAD
 
-# ytdlbot - model.py
+
+class PaymentStatus:
+    PENDING = "pending"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    REFUNDED = "refunded"
 
 
 Base = declarative_base()
@@ -57,7 +63,16 @@ class Payment(Base):
     id = Column(Integer, primary_key=True, autoincrement=True)
     method = Column(String(50), nullable=False)
     amount = Column(Float, nullable=False)
-    status = Column(Enum("pending", "completed", "failed", "refunded"), nullable=False)
+    status = Column(
+        Enum(
+            PaymentStatus.PENDING,
+            PaymentStatus.COMPLETED,
+            PaymentStatus.FAILED,
+            PaymentStatus.REFUNDED,
+        ),
+        nullable=False,
+    )
+    transaction_id = Column(String(100))
     user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
 
     user = relationship("User", back_populates="payments")
@@ -181,3 +196,26 @@ def reset_free():
         for user in users:
             user.free = FREE_DOWNLOAD
         session.commit()
+
+
+def credit_account(who, total_amount, transaction, method="stripe"):
+    with session_manager() as session:
+        user = session.query(User).filter(User.user_id == who).first()
+        if user:
+            dollar = total_amount / 100
+            price = int(os.getenv("TOKEN_PRICE"))  # per one dollar
+            user.paid += int(dollar * price)
+            logging.info("user %d credited with %d tokens, payment:$%.2f", who, user.paid, dollar)
+            session.add(
+                Payment(
+                    method=method,
+                    amount=total_amount,
+                    status=PaymentStatus.COMPLETED,
+                    transaction_id=transaction,
+                    user_id=user.id,
+                )
+            )
+            session.commit()
+            return user.free, user.paid
+
+        return None, None

src/main.py

@@ -7,8 +7,8 @@
 
 __author__ = "Benny <benny.think@gmail.com>"
 
-import contextlib
 import logging
+import os
 import re
 import threading
 import time
@@ -36,6 +36,7 @@ from config import (
     BotText,
 )
 from database.model import (
+    credit_account,
     get_format_settings,
     get_free_quota,
     get_paid_quota,
@@ -144,6 +145,42 @@ def ping_handler(client: Client, message: types.Message):
     thread.start()
 
 
+@app.on_message(filters.command(["buy"]))
+def buy(client: Client, message: types.Message):
+    chat_id = message.chat.id
+    logging.info("Generating invoice for %s", chat_id)
+    price = os.getenv("TOKEN_PRICE")
+    client.send_invoice(
+        chat_id,
+        f"{price} permanent download",
+        "Please make a payment via Stripe",
+        f"vip-{chat_id}",
+        "USD",
+        [types.LabeledPrice(label="VIP", amount=100)],
+        provider_token=os.getenv("PROVIDER_TOKEN"),
+        protect_content=True,
+        start_parameter="no-forward",
+    )
+
+
+@app.on_pre_checkout_query()
+def pre_checkout(client: Client, query: types.PreCheckoutQuery):
+    client.answer_pre_checkout_query(query.id, ok=True)
+
+
+@app.on_message(filters.successful_payment)
+def successful_payment(client: Client, message: types.Message):
+    who = message.chat.id
+    amount = message.successful_payment.total_amount  # in cents
+    ch = message.successful_payment.provider_payment_charge_id
+    free, paid = credit_account(who, amount, ch)
+    if paid > 0:
+        message.reply_text(f"Payment successful! You now have {free} free and {paid} paid quota.")
+    else:
+        message.reply_text("Something went wrong. Please contact the admin.")
+    message.delete()
+
+
 @app.on_message(filters.command(["stats"]))
 def stats_handler(client: Client, message: types.Message):
     chat_id = message.chat.id

src/test.py

@@ -0,0 +1,14 @@
+#!/usr/bin/env python3
+# coding: utf-8
+
+# ytdlbot - test.py
+
+import yt_dlp
+
+url = "https://www.youtube.com/watch?v=e19kTVgb2c8"
+opts = {
+    "cookiefile": "cookies.txt",
+}
+
+with yt_dlp.YoutubeDL(opts) as ydl:
+    ydl.download([url])