Spaces:
Running
Running
File size: 4,896 Bytes
157e137 |
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 |
import pymongo
import time
import motor.motor_asyncio
from bson.objectid import ObjectId
from bson.errors import InvalidId
from FileStream.server.exceptions import FIleNotFound
class Database:
def __init__(self, uri, database_name):
self._client = motor.motor_asyncio.AsyncIOMotorClient(uri)
self.db = self._client[database_name]
self.col = self.db.users
self.black = self.db.blacklist
self.file = self.db.file
#---------------------[ NEW USER ]---------------------#
def new_user(self, id):
return dict(
id=id,
join_date=time.time(),
agreed_to_tos=False,
Links=0,
Plan="Free"
)
# ---------------------[ ADD USER ]---------------------#
async def add_user(self, id):
user = self.new_user(id)
await self.col.insert_one(user)
# ---------------------[ GET USER ]---------------------#
async def get_user(self, id):
user = await self.col.find_one({'id': int(id)})
return user
# ---------------------[ CHECK USER ]---------------------#
async def total_users_count(self):
count = await self.col.count_documents({})
return count
async def get_all_users(self):
all_users = self.col.find({})
return all_users
# ---------------------[ REMOVE USER ]---------------------#
async def delete_user(self, user_id):
await self.col.delete_many({'id': int(user_id)})
# ---------------------[ BAN, UNBAN USER ]---------------------#
def black_user(self, id):
return dict(
id=id,
ban_date=time.time()
)
async def ban_user(self, id):
user = self.black_user(id)
await self.black.insert_one(user)
async def unban_user(self, id):
await self.black.delete_one({'id': int(id)})
async def is_user_banned(self, id):
user = await self.black.find_one({'id': int(id)})
return True if user else False
async def total_banned_users_count(self):
count = await self.black.count_documents({})
return count
# ---------------------[ ADD FILE TO DB ]---------------------#
async def add_file(self, file_info):
file_info["time"] = time.time()
fetch_old = await self.get_file_by_fileuniqueid(file_info["user_id"], file_info["file_unique_id"])
if fetch_old:
return fetch_old["_id"]
await self.count_links(file_info["user_id"], "+")
return (await self.file.insert_one(file_info)).inserted_id
# ---------------------[ FIND FILE IN DB ]---------------------#
async def find_files(self, user_id, range):
user_files=self.file.find({"user_id": user_id})
user_files.skip(range[0] - 1)
user_files.limit(range[1] - range[0] + 1)
user_files.sort('_id', pymongo.DESCENDING)
total_files = await self.file.count_documents({"user_id": user_id})
return user_files, total_files
async def get_file(self, _id):
try:
file_info=await self.file.find_one({"_id": ObjectId(_id)})
if not file_info:
raise FIleNotFound
return file_info
except InvalidId:
raise FIleNotFound
async def get_file_by_fileuniqueid(self, id, file_unique_id, many=False):
if many:
return self.file.find({"file_unique_id": file_unique_id})
else:
file_info=await self.file.find_one({"user_id": id, "file_unique_id": file_unique_id})
if file_info:
return file_info
return False
# ---------------------[ TOTAL FILES ]---------------------#
async def total_files(self, id=None):
if id:
return await self.file.count_documents({"user_id": id})
return await self.file.count_documents({})
# ---------------------[ DELETE FILES ]---------------------#
async def delete_one_file(self, _id):
await self.file.delete_one({'_id': ObjectId(_id)})
# ---------------------[ UPDATE FILES ]---------------------#
async def update_file_ids(self, _id, file_ids: dict):
await self.file.update_one({"_id": ObjectId(_id)}, {"$set": {"file_ids": file_ids}})
# ---------------------[ PAID SYS ]---------------------#
async def link_available(self, id):
user = await self.col.find_one({"id": id})
if user.get("Plan") == "Plus":
return "Plus"
elif user.get("Plan") == "Free":
files = await self.file.count_documents({"user_id": id})
if files < 11:
return True
return False
async def count_links(self, id, operation: str):
if operation == "-":
await self.col.update_one({"id": id}, {"$inc": {"Links": -1}})
elif operation == "+":
await self.col.update_one({"id": id}, {"$inc": {"Links": 1}}) |