bemaniutils/bemani/utils/dbutils.py

import argparse
import getpass
import sys
from typing import Optional

from bemani.data import Config, Data, DBCreateException
from bemani.utils.config import load_config


def create(config: Config) -> None:
    data = Data(config)
    data.create()
    data.close()


def generate(config: Config, message: Optional[str], allow_empty: bool) -> None:
    if message is None:
        raise Exception("Please provide a message!")
    data = Data(config)
    data.generate(message, allow_empty)
    data.close()


def upgrade(config: Config) -> None:
    data = Data(config)
    data.upgrade()
    data.close()


def change_password(config: Config, username: Optional[str]) -> None:
    if username is None:
        raise Exception("Please provide a username!")
    password1 = getpass.getpass("Password: ")
    password2 = getpass.getpass("Re-enter password: ")
    if password1 != password2:
        raise Exception("Passwords don't match!")
    data = Data(config)
    userid = data.local.user.from_username(username)
    if userid is None:
        raise Exception("User not found!")
    data.local.user.update_password(userid, password1)
    print(f"User {username} changed password.")


def add_admin(config: Config, username: Optional[str]) -> None:
    if username is None:
        raise Exception("Please provide a username!")
    data = Data(config)
    userid = data.local.user.from_username(username)
    if userid is None:
        raise Exception("User not found!")
    user = data.local.user.get_user(userid)
    user.admin = True
    data.local.user.put_user(user)
    print(f"User {username} gained admin rights.")


def remove_admin(config: Config, username: Optional[str]) -> None:
    if username is None:
        raise Exception("Please provide a username!")
    data = Data(config)
    userid = data.local.user.from_username(username)
    if userid is None:
        raise Exception("User not found!")
    user = data.local.user.get_user(userid)
    user.admin = False
    data.local.user.put_user(user)
    print(f"User {username} lost admin rights.")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="A utility for working with databases created with this codebase."
    )
    parser.add_argument(
        "operation",
        help="Operation to perform, options include 'create', 'generate', 'upgrade', 'change-password', 'add-admin' and 'remove-admin'.",
        type=str,
    )
    parser.add_argument(
        "-u",
        "--username",
        help="Username of user to add/remove admin rights for.",
        type=str,
    )
    parser.add_argument(
        "-m",
        "--message",
        help="Message to use for auto-generated migration scripts.",
        type=str,
    )
    parser.add_argument(
        "-e",
        "--allow-empty",
        help="Allow empty migration script to be generated. Useful for data-only migrations.",
        action="store_true",
    )
    parser.add_argument(
        "-c",
        "--config",
        help="Core configuration. Defaults to server.yaml",
        type=str,
        default="server.yaml",
    )
    args = parser.parse_args()

    config = Config()
    load_config(args.config, config)

    try:
        if args.operation == "create":
            create(config)
        elif args.operation == "generate":
            generate(config, args.message, args.allow_empty)
        elif args.operation == "upgrade":
            upgrade(config)
        elif args.operation == "add-admin":
            add_admin(config, args.username)
        elif args.operation == "remove-admin":
            remove_admin(config, args.username)
        elif args.operation == "change-password":
            change_password(config, args.username)
        else:
            raise Exception(f"Unknown operation '{args.operation}'")
    except DBCreateException as e:
        print(str(e))
        sys.exit(1)


if __name__ == "__main__":
    main()