Spaces:
Paused
Paused
File size: 3,980 Bytes
e99a4e1 4369d38 e99a4e1 3104a9c 4369d38 3104a9c 2552c4a a9cba1d 3104a9c a9cba1d 3104a9c e99a4e1 3104a9c e99a4e1 3104a9c e99a4e1 3104a9c e99a4e1 9d3be26 3104a9c e99a4e1 9d3be26 3104a9c e99a4e1 3104a9c e99a4e1 3104a9c e99a4e1 3104a9c 1d407e5 9d3be26 3104a9c e99a4e1 9d3be26 3104a9c e99a4e1 3104a9c e99a4e1 3104a9c 12f6d26 3104a9c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | #!/usr/bin/env python3
# coding: utf-8
import logging
import math
import os
from contextlib import contextmanager
from typing import Literal
from pymongo import MongoClient
from datetime import datetime
from config import ENABLE_VIP, FREE_DOWNLOAD, MONGO_URI
class PaymentStatus:
PENDING = "pending"
COMPLETED = "completed"
FAILED = "failed"
REFUNDED = "refunded"
def get_database():
client = MongoClient(MONGO_URI)
try:
# Verify connection
client.admin.command('ping')
except Exception as e:
logging.error("Failed to connect to database: %s", e)
raise
return client[os.getenv("DB_NAME", "ytdl_bot")]
db = get_database()
def get_quality_settings(tgid) -> Literal["high", "medium", "low", "audio", "custom"]:
user = db.users.find_one({"user_id": tgid})
if user and "settings" in user:
return user["settings"].get("quality", "high")
return "high"
def get_format_settings(tgid) -> Literal["video", "audio", "document"]:
user = db.users.find_one({"user_id": tgid})
if user and "settings" in user:
return user["settings"].get("format", "video")
return "video"
def set_user_settings(tgid: int, key: str, value: str):
db.users.update_one(
{"user_id": tgid},
{"$set": {f"settings.{key}": value}},
upsert=True
)
def get_free_quota(uid: int):
if not ENABLE_VIP:
return math.inf
user = db.users.find_one({"user_id": uid})
if user:
return user.get("free", FREE_DOWNLOAD)
return FREE_DOWNLOAD
def get_paid_quota(uid: int):
if ENABLE_VIP:
user = db.users.find_one({"user_id": uid})
if user:
return user.get("paid", 0)
return 0
return math.inf
def reset_free_quota(uid: int):
db.users.update_one(
{"user_id": uid},
{"$set": {"free": 5}}
)
def add_paid_quota(uid: int, amount: int):
db.users.update_one(
{"user_id": uid},
{"$inc": {"paid": amount}}
)
def check_quota(uid: int):
if not ENABLE_VIP:
return
user = db.users.find_one({"user_id": uid})
if user and (user.get("free", 0) + user.get("paid", 0)) <= 0:
raise Exception("Quota exhausted. Please /buy or wait until free quota is reset")
def use_quota(uid: int):
if not ENABLE_VIP:
return
user = db.users.find_one({"user_id": uid})
if user:
if user.get("free", 0) > 0:
db.users.update_one(
{"user_id": uid},
{"$inc": {"free": -1}}
)
elif user.get("paid", 0) > 0:
db.users.update_one(
{"user_id": uid},
{"$inc": {"paid": -1}}
)
else:
raise Exception("Quota exhausted. Please /buy or wait until free quota is reset")
def init_user(uid: int):
db.users.update_one(
{"user_id": uid},
{"$setOnInsert": {
"user_id": uid,
"free": FREE_DOWNLOAD,
"paid": 0,
"created_at": datetime.utcnow()
}},
upsert=True
)
def reset_free():
db.users.update_many(
{},
{"$set": {"free": FREE_DOWNLOAD}}
)
def credit_account(who, total_amount: int, quota: int, transaction, method="stripe"):
user = db.users.find_one_and_update(
{"user_id": who},
{"$inc": {"paid": quota}},
return_document=True
)
if user:
dollar = total_amount / 100
logging.info("user %d credited with %d tokens, payment:$%.2f", who, user["paid"], dollar)
payment = {
"method": method,
"amount": total_amount,
"status": PaymentStatus.COMPLETED,
"transaction_id": transaction,
"user_id": who,
"created_at": datetime.utcnow()
}
db.payments.insert_one(payment)
return user.get("free", 0), user.get("paid", 0)
return None, None |