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()