Maas-Model-platform/backend/scripts/reset_user_password.py

"""
普通用户密码重置脚本

用法:
    python scripts/reset_user_password.py <email_or_username> <new_password>

示例:
    python scripts/reset_user_password.py user@example.com newpass123
    python scripts/reset_user_password.py myusername newpass123
"""
import os
import sys
from pathlib import Path
from dotenv import load_dotenv
import psycopg2
from passlib.context import CryptContext

# 添加项目根目录到路径
sys.path.insert(0, str(Path(__file__).parent.parent))

# 加载环境变量
env_path = Path(__file__).parent.parent / '.env'
load_dotenv(dotenv_path=env_path)

# 普通用户：直接 bcrypt，不做 sha256
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")


def get_db_connection():
    try:
        conn = psycopg2.connect(
            host=os.getenv("DB_HOST", "localhost"),
            port=int(os.getenv("DB_PORT", "5432")),
            database=os.getenv("DB_NAME", "aigcspace"),
            user=os.getenv("DB_USER", "postgres"),
            password=os.getenv("DB_PASSWORD", "")
        )
        return conn
    except Exception as e:
        print(f"❌ 数据库连接失败: {e}")
        sys.exit(1)


def reset_password(identifier: str, new_password: str):
    conn = get_db_connection()
    cur = conn.cursor()

    # 先查用户（支持用户名、手机号或邮箱）
    cur.execute(
        "SELECT id, username, email, phone FROM aigcspace.users WHERE username = %s OR phone = %s OR email = %s",
        (identifier, identifier, identifier)
    )
    row = cur.fetchone()

    if not row:
        print(f"❌ 未找到用户: {identifier}")
        cur.close()
        conn.close()
        sys.exit(1)

    user_id, username, email, phone = row
    print(f"找到用户: id={user_id}, username={username}, email={email}, phone={phone}")

    # 生成新密码哈希，确保密码是字符串类型
    new_hash = pwd_context.hash(str(new_password))

    cur.execute(
        "UPDATE aigcspace.users SET password_hash = %s WHERE id = %s",
        (new_hash, user_id)
    )
    conn.commit()

    print(f"✓ 密码已更新")
    cur.close()
    conn.close()


if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("用法: python scripts/reset_user_password.py <email_or_username> <new_password>")
        sys.exit(1)

    identifier = sys.argv[1]
    new_password = sys.argv[2]
    reset_password(identifier, new_password)