Init admin user

Artrajz · Oct 31, 2023 · 5632361 · 5632361
1 parent a403bf2
commit 5632361
Show file tree

Hide file tree

Showing 4 changed files with 97 additions and 46 deletions.
diff --git a/app.py b/app.py
@@ -6,14 +6,13 @@
 from flask_wtf import CSRFProtect
 
 from tts_app import frontend, voice_api, auth, admin
-from tts_app.auth.models import users
-from utils.config_manager import global_config as config
+from utils.config_manager import global_config
 
 app = Flask(__name__, template_folder=os.path.join(os.path.dirname(__file__), 'tts_app', 'templates'),
             static_folder=os.path.join(os.path.dirname(__file__), 'tts_app', 'static'))
 
 app.config.from_pyfile("config.py")
-app.config.update(config)
+app.config.update(global_config)
 
 login_manager = LoginManager()
 login_manager.init_app(app)
@@ -24,6 +23,7 @@
 
 @login_manager.user_loader
 def load_user(user_id):
+    users = app.config["users"]["admin"]
     for user in users.values():
         if user.get_id() == user_id:
             return user

diff --git a/tts_app/auth/models.py b/tts_app/auth/models.py
@@ -16,4 +16,4 @@ def is_anonymous(self):
     def get_id(self):
         return str(self.id)
 
-users = {'admin': User(1, 'admin', 'password')}
+# users = {'group': {'username': User(1, 'username', 'password')}}
diff --git a/tts_app/auth/views.py b/tts_app/auth/views.py
@@ -1,8 +1,7 @@
-from flask import render_template, redirect, url_for, flash, Blueprint
+from flask import render_template, redirect, url_for, flash, Blueprint, current_app
 from flask_login import login_user, logout_user, login_required
 
 from tts_app.auth.forms import LoginForm
-from tts_app.auth.models import users
 
 auth = Blueprint('auth', __name__)
 
@@ -11,6 +10,7 @@
 def login():
     form = LoginForm()
     if form.validate_on_submit():
+        users = current_app.config["users"]["admin"]
         user = users.get(form.username.data)
         if user and user.password == form.password.data:
             login_user(user)
@@ -19,11 +19,10 @@ def login():
         flash('Wrong username or password.')
     return render_template('login.html', form=form)
 
+
 @auth.route('/logout')
 @login_required
 def logout():
     logout_user()
     flash('You have been logged out.')
     return redirect(url_for('auth.login'))
-
-
diff --git a/utils/config_manager.py b/utils/config_manager.py
@@ -2,25 +2,29 @@
 import secrets
 import shutil
 import logging
+import string
 
 import torch
 import yaml
+from flask import current_app
 
 import config as default_config
+from tts_app.auth.models import User
 from utils.data_utils import check_is_none
 
 YAML_CONFIG_FILE = os.path.join(default_config.ABS_PATH, 'config.yml')
 
 logging.getLogger().setLevel(logging.DEBUG)
 
+
 class Config(dict):
     def __init__(self, *args, **kwargs):
         super(Config, self).__init__(*args, **kwargs)
 
     def __getattr__(self, key):
         if key in self:
             return self[key]
-        raise AttributeError(f"'Config' object has no attribute '{key}'")
+        # raise AttributeError(f"'Config' object has no attribute '{key}'")
 
     def __setattr__(self, key, value):
         self[key] = value
@@ -29,7 +33,7 @@ def __setattr__(self, key, value):
 global_config = Config()
 
 
-# 修改表示和构造器的定义
+# torch.device
 def represent_torch_device(dumper, device_obj):
     return dumper.represent_scalar('!torch.device', str(device_obj))
 
@@ -43,11 +47,27 @@ def construct_torch_device(loader, node):
 yaml.add_constructor('!torch.device', construct_torch_device, Loader=yaml.SafeLoader)
 
 
+def represent_user(dumper, user_obj):
+    return dumper.represent_mapping('!User', {
+        'id': user_obj.id,
+        'username': user_obj.username,
+        'password': user_obj.password
+    })
+
+
+# User
+def construct_user(loader, node):
+    user_data = loader.construct_mapping(node, deep=True)
+    return User(user_data['id'], user_data['username'], user_data['password'])
+
+
+yaml.add_representer(User, represent_user, Dumper=yaml.SafeDumper)
+yaml.add_constructor('!User', construct_user, Loader=yaml.SafeLoader)
+
+
 def load_yaml_config(filename):
     with open(filename, 'r') as f:
         yaml_config = yaml.safe_load(f)
-    global global_config
-    global_config.update(yaml_config)
     logging.info(f"Loading yaml from {YAML_CONFIG_FILE}")
     return Config(yaml_config)
 
@@ -58,10 +78,9 @@ def save_yaml_config(filename, data):
         dict_data = dict(data)
         with open(temp_filename, 'w') as f:
             yaml.safe_dump(dict_data, f, default_style="'")
-        global global_config
-        global_config.update(dict_data)
         shutil.move(temp_filename, filename)
         logging.info(f"Saving yaml to {YAML_CONFIG_FILE}")
+        current_app.config.update(data)
     except Exception as e:
         logging.error(f"Error while saving yaml: {e}")
         if os.path.exists(temp_filename):
@@ -72,35 +91,68 @@ def generate_secret_key(length=32):
     return secrets.token_hex(length)
 
 
-model_path = ["MODEL_LIST", "HUBERT_SOFT_MODEL", "DIMENSIONAL_EMOTION_NPY", "DIMENSIONAL_EMOTION_MODEL"]
-default_parameter = ["ID", "FORMAT", "LANG", "LENGTH", "NOISE", "NOISEW", "MAX", "SDP_RATIO"]
-
-if os.path.exists(YAML_CONFIG_FILE):
-    global_config.update(load_yaml_config(YAML_CONFIG_FILE))
-else:
-    global_config.setdefault("model_path", {})
-    global_config.setdefault("default_parameter", {})
-
-    for key, value in vars(default_config).items():
-        if key.islower():
-            continue
-        if key in model_path:
-            global_config["model_path"][key.lower()] = value
-        elif key in default_parameter:
-            global_config["default_parameter"][key.lower()] = value
-        else:
-            global_config[key] = value
-    logging.info("config.yml not found. Generating a new config.yml based on config.py.")
-    save_yaml_config(YAML_CONFIG_FILE, global_config)
-
-if check_is_none(global_config.SECRET_KEY):
-    secret_key = generate_secret_key()
-    global_config["SECRET_KEY"] = secret_key
-    logging.info(f"SECRET_KEY is not found or is None. Generating a new SECRET_KEY:{secret_key}")
-    save_yaml_config(YAML_CONFIG_FILE, global_config)
-
-if check_is_none(global_config.API_KEY):
-    secret_key = generate_secret_key()
-    global_config["API_KEY"] = secret_key
-    logging.info(f"API_KEY is not found or is None. Generating a new API_KEY:{secret_key}")
-    save_yaml_config(YAML_CONFIG_FILE, global_config)
+def generate_random_username(length=8):
+    characters = string.ascii_letters + string.digits
+    username = ''.join(secrets.choice(characters) for _ in range(length))
+    return username
+
+
+def generate_random_password(length=16):
+    characters = string.ascii_letters + string.digits
+    password = ''.join(secrets.choice(characters) for _ in range(length))
+    return password
+
+
+def init_config():
+    global global_config
+    model_path = ["MODEL_LIST", "HUBERT_SOFT_MODEL", "DIMENSIONAL_EMOTION_NPY", "DIMENSIONAL_EMOTION_MODEL"]
+    default_parameter = ["ID", "FORMAT", "LANG", "LENGTH", "NOISE", "NOISEW", "MAX", "SDP_RATIO"]
+
+    if os.path.exists(YAML_CONFIG_FILE):
+        global_config.update(load_yaml_config(YAML_CONFIG_FILE))
+    else:
+        global_config.setdefault("model_path", {})
+        global_config.setdefault("default_parameter", {})
+
+        for key, value in vars(default_config).items():
+            if key.islower():
+                continue
+            if key in model_path:
+                global_config["model_path"][key.lower()] = value
+            elif key in default_parameter:
+                global_config["default_parameter"][key.lower()] = value
+            else:
+                global_config[key] = value
+        logging.info("config.yml not found. Generating a new config.yml based on config.py.")
+        save_yaml_config(YAML_CONFIG_FILE, global_config)
+
+    if check_is_none(global_config.SECRET_KEY):
+        secret_key = generate_secret_key()
+        global_config["SECRET_KEY"] = secret_key
+        logging.info(f"SECRET_KEY is not found or is None. Generating a new SECRET_KEY:{secret_key}")
+        save_yaml_config(YAML_CONFIG_FILE, global_config)
+
+    if check_is_none(global_config.API_KEY):
+        secret_key = generate_secret_key()
+        global_config["API_KEY"] = secret_key
+        logging.info(f"API_KEY is not found or is None. Generating a new API_KEY:{secret_key}")
+        save_yaml_config(YAML_CONFIG_FILE, global_config)
+
+    if getattr(global_config, "users") is None:
+        random_username = generate_random_username()
+        random_password = generate_random_password()
+        logging.info(
+            f"New admin user created:\n"
+            f"{'-' * 40}\n"
+            f"| Username: {random_username:<26} |\n"
+            f"| Password: {random_password:<26} |\n"
+            f"{'-' * 40}\n"
+            f"Please do not share this information.")
+        global_config["users"] = {}
+        global_config["users"]["admin"] = {f"{random_username}": User(1, random_username, random_password)}
+        save_yaml_config(YAML_CONFIG_FILE, global_config)
+
+    return global_config
+
+
+init_config()