123 lines
3.8 KiB
Python
123 lines
3.8 KiB
Python
import argparse
|
|
import getpass
|
|
import sys
|
|
import yaml
|
|
from typing import Any, Dict, Optional
|
|
|
|
from bemani.data import Data, DBCreateException
|
|
|
|
|
|
def create(config: Dict[str, Any]) -> None:
|
|
data = Data(config)
|
|
data.create()
|
|
data.close()
|
|
|
|
|
|
def generate(config: Dict[str, Any], 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: Dict[str, Any]) -> None:
|
|
data = Data(config)
|
|
data.upgrade()
|
|
data.close()
|
|
|
|
|
|
def change_password(config: Dict[str, Any], 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: Dict[str, Any], 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: Dict[str, Any], 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 = yaml.safe_load(open(args.config))
|
|
config['database']['engine'] = Data.create_engine(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()
|