mirror of
https://github.com/NovaOSS/nova-api.git
synced 2024-11-25 20:43:56 +01:00
116 lines
2.7 KiB
Python
116 lines
2.7 KiB
Python
"""User system."""
|
|
|
|
import os
|
|
import uuid
|
|
import time
|
|
import string
|
|
import random
|
|
import asyncio
|
|
|
|
from helpers import databases
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
async def prepare() -> None:
|
|
"""Creates the database tables"""
|
|
|
|
users_db = await databases.connect('users')
|
|
await users_db.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
id TEXT PRIMARY KEY,
|
|
api_key TEXT,
|
|
active BOOLEAN,
|
|
created_at INTEGER,
|
|
last_used INTEGER DEFAULT 0,
|
|
uses_count INTEGER DEFAULT 0,
|
|
tokens_generated INTEGER DEFAULT 0,
|
|
discord_id INTEGER DEFAULT 0,
|
|
credit INTEGER DEFAULT 0,
|
|
tags TEXT DEFAULT ''
|
|
)
|
|
"""
|
|
)
|
|
await users_db.commit()
|
|
|
|
async def add_user(
|
|
discord_id: int=0,
|
|
tags: list=None,
|
|
) -> dict:
|
|
"""Adds a new key to the database"""
|
|
|
|
chars = string.ascii_letters + string.digits
|
|
|
|
infix = os.getenv('KEYGEN_INFIX')
|
|
suffix = ''.join(random.choices(chars, k=20))
|
|
prefix = ''.join(random.choices(chars, k=20))
|
|
|
|
key = f'nv-{prefix}{infix}{suffix}'
|
|
|
|
tags = tags or []
|
|
new_user = {
|
|
'id': str(uuid.uuid4()),
|
|
'api_key': key,
|
|
'created_at': int(time.time()),
|
|
'last_used': 0,
|
|
'uses_count': 0,
|
|
'tokens_generated': 0,
|
|
'active': True,
|
|
'discord_id': discord_id,
|
|
'credit': 0,
|
|
'tags': '/'.join(tags)
|
|
}
|
|
|
|
await databases.insert_dict(new_user, 'users', 'users')
|
|
return new_user
|
|
|
|
async def get_user(
|
|
by_id: str='',
|
|
by_discord_id: int=0,
|
|
):
|
|
users_db = await databases.connect('users')
|
|
|
|
async with users_db.execute(
|
|
'SELECT * FROM users WHERE id = :id OR discord_id = :discord_id',
|
|
{'id': by_id, 'discord_id': by_discord_id}
|
|
) as cursor:
|
|
async for row in cursor:
|
|
result = await databases.row_to_dict(row, cursor)
|
|
return result
|
|
|
|
return None
|
|
|
|
async def get_all_users():
|
|
users_db = await databases.connect('users')
|
|
results = []
|
|
|
|
async with users_db.execute(
|
|
'SELECT * FROM users'
|
|
) as cursor:
|
|
async for row in cursor:
|
|
result = await databases.row_to_dict(row, cursor)
|
|
results.append(result)
|
|
|
|
return results
|
|
|
|
async def demo():
|
|
await prepare()
|
|
|
|
users = await get_all_users()
|
|
print(users)
|
|
|
|
example_id = 133769420
|
|
user = await add_user(discord_id=example_id)
|
|
print(user)
|
|
|
|
del user
|
|
print('Fetching user...')
|
|
|
|
user = await get_user(by_discord_id=example_id)
|
|
print(user['api_key'])
|
|
|
|
if __name__ == '__main__':
|
|
asyncio.run(demo())
|
|
os.system(f'pkill -f {os.path.basename(__file__)}')
|