mirror of
https://github.com/tgbot-collection/ytdlbot.git
synced 2024-09-20 15:05:56 +08:00
fix first download quota not used issue, add username in db
This commit is contained in:
parent
8d1d8c4276
commit
91547355e9
35
limit.py
35
limit.py
|
@ -12,12 +12,14 @@ import logging
|
|||
import math
|
||||
import os
|
||||
import sqlite3
|
||||
import tempfile
|
||||
import time
|
||||
|
||||
import redis
|
||||
import requests
|
||||
|
||||
QUOTA = 10 * 1024 * 1024 * 1024 # 10G
|
||||
# QUOTA = 10 * 1024 * 1024 * 1024 # 10G
|
||||
QUOTA = 5 * 1024 * 1024 # 10G
|
||||
EX = 24 * 3600
|
||||
MULTIPLY = 5 # VIP1 is 5*10-50G, VIP2 is 100G
|
||||
USD2CNY = 6 # $5 --> ¥30
|
||||
|
@ -43,6 +45,7 @@ class SQLite:
|
|||
user_id integer
|
||||
constraint VIP
|
||||
primary key,
|
||||
username varchar(100),
|
||||
payment_amount integer,
|
||||
payment_id varchar(100),
|
||||
level integer default 1,
|
||||
|
@ -57,6 +60,17 @@ class SQLite:
|
|||
self.con.close()
|
||||
|
||||
|
||||
def get_username(chat_id):
|
||||
from ytdl import create_app
|
||||
with tempfile.NamedTemporaryFile() as tmp:
|
||||
app = create_app(tmp.name, 1)
|
||||
app.start()
|
||||
data = app.get_chat(chat_id).first_name
|
||||
app.stop()
|
||||
|
||||
return data
|
||||
|
||||
|
||||
class VIP(Redis, SQLite):
|
||||
|
||||
def check_vip(self, user_id: "int") -> "tuple":
|
||||
|
@ -65,8 +79,7 @@ class VIP(Redis, SQLite):
|
|||
return data
|
||||
|
||||
def add_vip(self, user_data: "dict") -> ("bool", "str"):
|
||||
user_data["quota"] = QUOTA * user_data["level"] * MULTIPLY
|
||||
sql = "INSERT INTO VIP VALUES (?,?,?,?,?);"
|
||||
sql = "INSERT INTO VIP VALUES (?,?,?,?,?,?);"
|
||||
# first select
|
||||
self.cur.execute("SELECT * FROM VIP WHERE payment_id=?", (user_data["payment_id"],))
|
||||
is_exist = self.cur.fetchone()
|
||||
|
@ -99,7 +112,7 @@ class VIP(Redis, SQLite):
|
|||
if self.r.exists(user_id):
|
||||
self.r.decr(user_id, traffic)
|
||||
else:
|
||||
self.r.set(user_id, user_quota, ex=EX)
|
||||
self.r.set(user_id, user_quota - traffic, ex=EX)
|
||||
|
||||
|
||||
class BuyMeACoffee:
|
||||
|
@ -184,12 +197,18 @@ def verify_payment(user_id, unique) -> "str":
|
|||
f"Talk to @BennyThink if you need any assistant."
|
||||
else:
|
||||
vip = VIP()
|
||||
ud = {"user_id": user_id, "payment_amount": amount, "payment_id": pay_id, "level": level}
|
||||
ud = {
|
||||
"user_id": user_id,
|
||||
"username": get_username(user_id),
|
||||
"payment_amount": amount,
|
||||
"payment_id": pay_id,
|
||||
"level": level,
|
||||
"quota": QUOTA * level * MULTIPLY
|
||||
}
|
||||
|
||||
message = vip.add_vip(ud)
|
||||
return message
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
vip = VIP()
|
||||
ud = {"user_id": 260260121, "username": "bahd", "payment": 5, "payment_id": "dhja767", "level": 1}
|
||||
vip.add_vip(ud)
|
||||
pass
|
||||
|
|
16
ytdl.py
16
ytdl.py
|
@ -23,11 +23,19 @@ from downloader import convert_flac, upload_hook, ytdl_download
|
|||
from limit import verify_payment
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(filename)s [%(levelname)s]: %(message)s')
|
||||
api_id = int(os.getenv("APP_ID", 0))
|
||||
api_hash = os.getenv("APP_HASH")
|
||||
token = os.getenv("TOKEN")
|
||||
|
||||
app = Client("ytdl", api_id, api_hash, bot_token=token, workers=100)
|
||||
|
||||
def create_app(session="ytdl", workers=100):
|
||||
api_id = int(os.getenv("APP_ID", 0))
|
||||
api_hash = os.getenv("APP_HASH")
|
||||
token = os.getenv("TOKEN")
|
||||
|
||||
_app = Client(session, api_id, api_hash,
|
||||
bot_token=token, workers=workers)
|
||||
return _app
|
||||
|
||||
|
||||
app = create_app()
|
||||
bot_text = BotText()
|
||||
|
||||
|
||||
|
|
Loading…
Reference in a new issue