diff --git a/unwind/__main__.py b/unwind/__main__.py index e802831..08da1a0 100644 --- a/unwind/__main__.py +++ b/unwind/__main__.py @@ -1,9 +1,11 @@ import argparse import asyncio import logging +import secrets +from base64 import b64encode from pathlib import Path -from . import config +from . import config, db, models, utils from .db import close_connection_pool, open_connection_pool from .imdb import refresh_user_ratings_from_imdb from .imdb_import import download_datasets, import_from_file @@ -11,6 +13,38 @@ from .imdb_import import download_datasets, import_from_file log = logging.getLogger(__name__) +async def run_add_user(user_id: str, name: str, overwrite_existing: bool): + if not user_id.startswith("ur"): + raise ValueError(f"Invalid IMDb user ID: {user_id!a}") + + await open_connection_pool() + + async with db.new_connection() as conn: + user = await db.get(conn, models.User, imdb_id=user_id) + + if user is not None: + if overwrite_existing: + log.warning("⚠️ Overwriting existing user: %a", user) + else: + log.error("❌ User already exists: %a", user) + return + + secret = secrets.token_bytes() + + user = models.User(name=name, imdb_id=user_id, secret=utils.phc_scrypt(secret)) + async with db.transaction() as conn: + await db.add_or_update_user(conn, user) + + user_data = { + "secret": b64encode(secret), + "user": models.asplain(user), + } + + log.info("✨ User created: %a", user_data) + + await close_connection_pool() + + async def run_load_user_ratings_from_imdb(): await open_connection_pool() @@ -91,6 +125,26 @@ def getargs(): const="load-user-ratings-from-imdb", ) + parser_add_user = commands.add_parser( + "add-user", + help="Add a new user.", + description=""" + Add a new user. + """, + ) + parser_add_user.add_argument( + dest="mode", + action="store_const", + const="add-user", + ) + parser_add_user.add_argument("--name", required=True) + parser_add_user.add_argument("--imdb-id", required=True) + parser_add_user.add_argument( + "--overwrite-existing", + action="store_true", + help="Allow overwriting an existing user. WARNING: This will reset the user's password!", + ) + try: args = parser.parse_args() except TypeError: @@ -115,6 +169,8 @@ def main(): if args.mode == "load-user-ratings-from-imdb": asyncio.run(run_load_user_ratings_from_imdb()) + elif args.mode == "add-user": + asyncio.run(run_add_user(args.imdb_id, args.name, args.overwrite_existing)) elif args.mode == "import-imdb-dataset": asyncio.run(run_import_imdb_dataset(args.basics, args.ratings)) elif args.mode == "download-imdb-dataset":