45 lines
1.4 KiB
Python
45 lines
1.4 KiB
Python
# add_user.py
|
|
import asyncio
|
|
import secrets
|
|
from sqlalchemy.future import select
|
|
from sqlalchemy.exc import IntegrityError
|
|
from getpass import getpass
|
|
from data.db import engine, SessionLocal
|
|
from data.models import Base, User, UserRole
|
|
|
|
def create_secret():
|
|
return secrets.token_hex(32)
|
|
|
|
async def create_user():
|
|
|
|
async with SessionLocal() as session:
|
|
username = input("Enter username: ").strip()
|
|
role_input = input("Enter role (admin/user): ").strip().lower()
|
|
print('\n')
|
|
if role_input not in UserRole._value2member_map_:
|
|
print("> Invalid role. Please enter 'admin' or 'user'.")
|
|
return
|
|
|
|
role = UserRole(role_input)
|
|
|
|
# Check if username already exists
|
|
result = await session.execute(select(User).filter_by(username=username))
|
|
existing_user = result.scalars().first()
|
|
if existing_user:
|
|
print(f"> Username '{username}' is already taken.")
|
|
return
|
|
|
|
# Create new user
|
|
api_key = create_secret()
|
|
new_user = User(username=username, role=role, api_key=api_key)
|
|
session.add(new_user)
|
|
await session.commit()
|
|
await session.refresh(new_user)
|
|
|
|
|
|
print(f"> User '{username}' with role '{role.value}' created successfully.")
|
|
print(f"> API Key: {api_key}")
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(create_user())
|