import os import yaml import random import string import asyncio from dotenv import load_dotenv from motor.motor_asyncio import AsyncIOMotorClient try: from . import helpers except ImportError: import helpers load_dotenv() with open(os.path.join(helpers.root, 'api', 'config', 'config.yml'), encoding='utf8') as f: credits_config = yaml.safe_load(f) ## MONGODB Setup class UserManager: """ Manager of all users in the database. """ def __init__(self): self.conn = AsyncIOMotorClient(os.environ['MONGO_URI']) async def _get_collection(self, collection_name: str): return self.conn[os.getenv('MONGO_NAME', 'nova-test')][collection_name] async def get_all_users(self): collection = self.conn[os.getenv('MONGO_NAME', 'nova-test')]['users'] return collection#.find() async def create(self, discord_id: str = '') -> dict: db = await self._get_collection('users') chars = string.ascii_letters + string.digits infix = os.getenv('KEYGEN_INFIX', 'S3LFH0ST') suffix = ''.join(random.choices(chars, k=20)) prefix = ''.join(random.choices(chars, k=20)) new_api_key = f'nv2-{prefix}{infix}{suffix}' existing_user = await self.user_by_discord_id(discord_id) if existing_user: # just change api key await db.update_one({'auth.discord': str(int(discord_id))}, {'$set': {'api_key': new_api_key}}) else: new_user = { 'api_key': new_api_key, 'credits': credits_config['start-credits'], 'role': '', 'level': '', 'status': { 'active': True, 'ban_reason': '', }, 'auth': { 'discord': str(discord_id), 'github': None } } await db.insert_one(new_user) user = await db.find_one({'api_key': new_api_key}) return user async def user_by_id(self, user_id: str): db = await self._get_collection('users') return await db.find_one({'_id': user_id}) async def user_by_discord_id(self, discord_id: str): db = await self._get_collection('users') return await db.find_one({'auth.discord': str(int(discord_id))}) async def user_by_api_key(self, key: str): db = await self._get_collection('users') return await db.find_one({'api_key': key}) async def update_by_id(self, user_id: str, update): db = await self._get_collection('users') return await db.update_one({'_id': user_id}, update) async def update_by_discord_id(self, discord_id: str, update): db = await self._get_collection('users') return await db.update_one({'auth.discord': str(int(discord_id))}, update) async def update_by_filter(self, obj_filter, update): db = await self._get_collection('users') return await db.update_one(obj_filter, update) async def delete(self, user_id: str): db = await self._get_collection('users') await db.delete_one({'_id': user_id}) manager = UserManager() async def demo(): user = await UserManager().create(69420) print(user) if __name__ == '__main__': asyncio.run(demo())