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}})