1
0
mirror of synced 2024-11-14 18:07:36 +01:00
bemaniutils/bemani/utils/dbutils.py
Jennifer Taylor 509cb4f0d9 Convert most of the format() string calls to f-strings using libcst.
Exact commands run were:

  python3 -m libcst.tool codemod convert_format_to_fstring.ConvertFormatStringCommand . --no-format
  python3 setup.py build_ext --inplace
2020-01-07 21:29:07 +00:00

123 lines
3.9 KiB
Python

import argparse
import getpass
import sys
import yaml # type: ignore
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)) # type: ignore
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()