File size: 6,158 Bytes
35e7795 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 | #!/usr/bin/env python
"""创建超级用户脚本"""
import argparse
import getpass
import sys
from pathlib import Path
# 添加项目根目录到路径（脚本在 scripts 目录下，需要指向上一级目录）
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))
# 在添加路径后再导入项目模块
from qa_annotate.database.base import SessionLocal, init_db # noqa: E402
from qa_annotate.database.crud import UserCRUD # noqa: E402
from qa_annotate.schema.user import UserCreate, UserUpdate # noqa: E402
from qa_annotate.utils.password import hash_password # noqa: E402
def create_superuser(
username=None, password=None, full_name=None, update_existing=False
):
"""创建超级用户
Args:
username: 用户名（如果为None则从命令行输入）
password: 密码（如果为None则从命令行输入）
full_name: 全名（可选）
update_existing: 如果用户已存在，是否将其提升为超级用户
Returns:
bool: 是否成功
"""
# 初始化数据库（确保表存在）
init_db()
# 获取数据库会话
db = SessionLocal()
try:
# 获取用户输入
print("=" * 50)
print("创建超级用户")
print("=" * 50)
# 获取用户名
if username is None:
username = input("请输入用户名: ").strip()
else:
print(f"用户名: {username}")
if not username:
print("错误: 用户名不能为空")
return False
# 检查用户名是否已存在
existing_user = UserCRUD.get_by_username(db, username=username)
if existing_user:
if update_existing:
# 将现有用户提升为超级用户
print(f"用户 '{username}' 已存在，正在将其提升为超级用户...")
user_update = UserUpdate(is_superuser=True, is_active=True)
updated_user = UserCRUD.update(
db=db, user_id=existing_user.id, user_update=user_update
)
print("\n" + "=" * 50)
print("用户已成功提升为超级用户！")
print("=" * 50)
print(f"用户ID: {updated_user.id}")
print(f"用户名: {updated_user.username}")
print(f"全名: {updated_user.full_name or '(未设置)'}")
print(f"是否激活: {updated_user.is_active}")
print(f"是否超级用户: {updated_user.is_superuser}")
print("=" * 50)
return True
else:
print(f"错误: 用户名 '{username}' 已存在")
print("提示: 使用 --update 参数可以将现有用户提升为超级用户")
return False
# 获取密码
is_interactive = password is None
if password is None:
password = getpass.getpass("请输入密码: ")
else:
print("密码: ***")
if not password:
print("错误: 密码不能为空")
return False
if len(password) < 6:
print("错误: 密码长度至少为6位")
return False
# 确认密码（仅在交互模式下）
if is_interactive:
password_confirm = getpass.getpass("请再次输入密码: ")
if password != password_confirm:
print("错误: 两次输入的密码不一致")
return False
# 获取全名（可选）
if full_name is None:
full_name = input("请输入全名（可选，直接回车跳过）: ").strip()
if not full_name:
full_name = None
# 对密码进行SHA-256哈希
password_hash = hash_password(password)
# 创建超级用户
user_create = UserCreate(
username=username,
password=password_hash, # 存储哈希值
full_name=full_name,
is_active=True,
is_superuser=True,
)
user = UserCRUD.create(db=db, user=user_create)
print("\n" + "=" * 50)
print("超级用户创建成功！")
print("=" * 50)
print(f"用户ID: {user.id}")
print(f"用户名: {user.username}")
print(f"全名: {user.full_name or '(未设置)'}")
print(f"是否激活: {user.is_active}")
print(f"是否超级用户: {user.is_superuser}")
print("=" * 50)
return True
except KeyboardInterrupt:
print("\n\n操作已取消")
return False
except Exception as e:
print(f"\n错误: 创建超级用户失败 - {str(e)}")
import traceback
traceback.print_exc()
return False
finally:
db.close()
def main():
"""主函数"""
parser = argparse.ArgumentParser(
description="创建或更新超级用户",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
示例:
# 交互式创建超级用户（从项目根目录运行）
python scripts/create_superuser.py
# 使用命令行参数创建
python scripts/create_superuser.py --username admin --password admin123
# 将现有用户提升为超级用户
python scripts/create_superuser.py --username existing_user --update
""",
)
parser.add_argument("--username", "-u", type=str, help="用户名")
parser.add_argument(
"--password",
"-p",
type=str,
help="密码（不推荐在命令行中使用，建议留空以交互式输入）",
)
parser.add_argument("--full-name", "-n", type=str, dest="full_name", help="全名")
parser.add_argument(
"--update", action="store_true", help="如果用户已存在，将其提升为超级用户"
)
args = parser.parse_args()
success = create_superuser(
username=args.username,
password=args.password,
full_name=args.full_name,
update_existing=args.update,
)
sys.exit(0 if success else 1)
if __name__ == "__main__":
main()
|