zkxw/backend/scripts/reset_weak_password_111111.py

#!/usr/bin/env python3
"""将 sys_user 中密码明文等价于 111111 的账号统一改为 Zkxw@123。

因 bcrypt 每次盐不同，无法用 SQL WHERE password='111111' 命中，必须在应用侧校验。

连接方式（任选其一，优先顺序见下方）：
  1) --database-url / 环境变量 DATABASE_URL
  2) --db-host + --db-user + --db-password + --db-name（可选 --db-port）
  3) --env-file / 环境变量 ENV_FILE（默认 backend/.env）

用法示例（在 backend 目录）:
  # 显式指定库（推荐测试/生产，避免误连本地 .env）
  python scripts/reset_weak_password_111111.py \\
    --database-url 'mysql+pymysql://user:pass@host:3306/dbname' --dry-run

  python scripts/reset_weak_password_111111.py \\
    --db-host 1.2.3.4 --db-user root --db-password 'xxx' --db-name zkxw --dry-run

  # 使用指定环境文件
  python scripts/reset_weak_password_111111.py --env-file ./.env.test --dry-run

  # 正式写入（去掉 --dry-run）
  python scripts/reset_weak_password_111111.py --database-url '...' 

范围：del_flag=0（含停用账号 status!=1）。
"""
from __future__ import annotations

import argparse
import os
import re
import sys
from datetime import datetime
from pathlib import Path
from urllib.parse import quote_plus

WEAK_PASSWORD = "111111"
NEW_PASSWORD = "Zkxw@123"

BACKEND_ROOT = Path(__file__).resolve().parent.parent


def _mask_database_url(url: str) -> str:
    return re.sub(r":([^:@/]+)@", r":***@", url or "")


def _build_url_from_parts(
    *,
    host: str,
    user: str,
    password: str,
    name: str,
    port: int,
) -> str:
    return (
        f"mysql+pymysql://{quote_plus(user)}:{quote_plus(password)}"
        f"@{host}:{port}/{name}"
    )


def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="重置弱密码 111111 → Zkxw@123（支持变量指定目标库）"
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="只列出将修改的账号，不写库",
    )
    parser.add_argument(
        "--env-file",
        default=None,
        help="环境文件路径（等同 ENV_FILE）；未指定 database-url/分项时使用",
    )
    parser.add_argument(
        "--database-url",
        default=None,
        help="完整 SQLAlchemy URL（等同 DATABASE_URL），优先于 env-file",
    )
    parser.add_argument("--db-host", default=None, help="MySQL 主机")
    parser.add_argument("--db-port", type=int, default=3306, help="MySQL 端口，默认 3306")
    parser.add_argument("--db-user", default=None, help="MySQL 用户")
    parser.add_argument("--db-password", default=None, help="MySQL 密码")
    parser.add_argument("--db-name", default=None, help="数据库名")
    return parser.parse_args(argv)


def _resolve_connection(args: argparse.Namespace) -> tuple[str, str]:
    """返回 (来源说明, DATABASE_URL)。在 import app 之前调用并写入 os.environ。"""
    parts = (args.db_host, args.db_user, args.db_password, args.db_name)
    if any(parts) and not all(p is not None and str(p) != "" for p in parts):
        raise SystemExit(
            "使用分项连接时须同时提供: --db-host --db-user --db-password --db-name"
        )

    if all(p is not None and str(p) != "" for p in parts):
        url = _build_url_from_parts(
            host=args.db_host,
            user=args.db_user,
            password=args.db_password,
            name=args.db_name,
            port=args.db_port or 3306,
        )
        os.environ["DATABASE_URL"] = url
        return ("cli --db-*", url)

    if args.database_url:
        url = args.database_url.strip()
        os.environ["DATABASE_URL"] = url
        return ("cli --database-url", url)

    env_url = (os.environ.get("DATABASE_URL") or "").strip()
    if env_url:
        return ("env DATABASE_URL", env_url)

    env_file = args.env_file or os.environ.get("ENV_FILE") or str(BACKEND_ROOT / ".env")
    env_file_path = Path(env_file).expanduser()
    if not env_file_path.is_absolute():
        env_file_path = (Path.cwd() / env_file_path).resolve()
    if not env_file_path.is_file():
        raise SystemExit(
            f"环境文件不存在: {env_file_path}\n"
            "请改用 --database-url 或 --db-host/--db-user/--db-password/--db-name"
        )
    os.environ["ENV_FILE"] = str(env_file_path)
    # DATABASE_URL 由 Settings 从该文件加载；此处仅记录来源
    return (f"env-file {env_file_path}", "")


def main(argv: list[str] | None = None) -> int:
    args = _parse_args(argv)
    source, early_url = _resolve_connection(args)

    sys.path.insert(0, str(BACKEND_ROOT))

    # 必须在设置 ENV_FILE / DATABASE_URL 之后再 import，否则 Settings 已冻结
    from sqlalchemy import create_engine
    from sqlalchemy.orm import sessionmaker

    from app.core.config import settings
    from app.core.security import get_password_hash, verify_password
    from app.models.models import SysUser

    database_url = (os.environ.get("DATABASE_URL") or "").strip() or settings.DATABASE_URL
    if early_url:
        database_url = early_url

    print(f"目标库来源: {source}")
    print(f"DATABASE_URL: {_mask_database_url(database_url)}")

    engine = create_engine(
        database_url,
        pool_pre_ping=True,
        pool_recycle=3600,
        echo=False,
    )
    Session = sessionmaker(autocommit=False, autoflush=False, expire_on_commit=False, bind=engine)
    db = Session()
    try:
        rows = (
            db.query(SysUser)
            .filter(SysUser.del_flag == 0)
            .order_by(SysUser.create_time.desc())
            .all()
        )
        hits = []
        for u in rows:
            pwd = getattr(u, "password", None) or ""
            if verify_password(WEAK_PASSWORD, pwd):
                hits.append(u)

        if not hits:
            print("未发现密码为 111111 的用户（del_flag=0）。")
            return 0

        print(f"命中 {len(hits)} 个弱密码用户：")
        for u in hits:
            print(
                f"  id={u.id} user_no={u.user_no!r} status={u.status!r} "
                f"tenant_no={u.tenant_no!r}"
            )

        if args.dry_run:
            print("dry-run：未写入数据库。去掉 --dry-run 后执行正式更新。")
            return 0

        new_hash = get_password_hash(NEW_PASSWORD)
        now = datetime.now()
        for u in hits:
            u.password = new_hash
            u.update_time = now
        db.commit()
        print(f"已更新 {len(hits)} 个用户密码为强密码（明文 Zkxw@123）。请通知相关用户。")
        return 0
    except Exception as exc:
        db.rollback()
        print(f"执行失败: {exc}", file=sys.stderr)
        return 1
    finally:
        db.close()
        engine.dispose()


if __name__ == "__main__":
    raise SystemExit(main())