document-management/backend/dms/seed_dev.py

"""显式、幂等的本地最小开发数据初始化。

执行：python -m dms.seed_dev
必须通过环境变量提供管理员和普通用户密码；默认不会覆盖已有用户密码。
"""

from __future__ import annotations

import os

from flask import Flask
from sqlalchemy import select

from dms import init_dms
from dms.common.enums import EnabledStatus, RoleCode, SecurityLevel
from dms.extensions import db
from dms.models import Organization, User
from dms.security.passwords import hash_password


PASSWORD_ENVIRONMENTS = {
    "admin": "DMS_SEED_ADMIN_PASSWORD",
    "user": "DMS_SEED_USER_PASSWORD",
}


def _required_passwords() -> dict[str, str]:
    missing = [
        environment
        for environment in PASSWORD_ENVIRONMENTS.values()
        if not os.environ.get(environment)
    ]
    if missing:
        raise RuntimeError("缺少必需的密码环境变量：" + "、".join(missing))
    return {
        username: os.environ[environment]
        for username, environment in PASSWORD_ENVIRONMENTS.items()
    }


def _overwrite_passwords_enabled() -> bool:
    raw_value = os.environ.get("DMS_SEED_OVERWRITE_PASSWORDS", "false").lower()
    if raw_value not in {"true", "false"}:
        raise RuntimeError("DMS_SEED_OVERWRITE_PASSWORDS必须是true或false")
    return raw_value == "true"


def _organization(
    *,
    org_code: str,
    org_name: str,
    org_path: str,
    sort_no: int,
    parent_id: int | None,
) -> tuple[Organization, bool]:
    existing = db.session.scalar(
        select(Organization).where(Organization.org_code == org_code)
    )
    if existing is not None:
        return existing, False
    organization = Organization(
        org_code=org_code,
        org_name=org_name,
        org_path=org_path,
        sort_no=sort_no,
        parent_id=parent_id,
        status=EnabledStatus.ENABLED.value,
    )
    db.session.add(organization)
    db.session.flush()
    return organization, True


def _user(
    *,
    username: str,
    real_name: str,
    password: str,
    role_code: RoleCode,
    security_level: SecurityLevel,
    organization: Organization,
    overwrite_password: bool,
) -> bool:
    existing = db.session.scalar(select(User).where(User.username == username))
    if existing is not None:
        if overwrite_password:
            existing.password_hash = hash_password(password)
            existing.auth_version += 1
            existing.row_version += 1
        return False
    db.session.add(
        User(
            username=username,
            password_hash=hash_password(password),
            real_name=real_name,
            organization_id=organization.id,
            organization_name=organization.org_name,
            role_code=role_code.value,
            security_level=security_level.value,
            status=EnabledStatus.ENABLED.value,
        )
    )
    return True


def seed() -> tuple[int, int]:
    passwords = _required_passwords()
    overwrite_passwords = _overwrite_passwords_enabled()
    created_organizations = 0
    created_users = 0

    root, created = _organization(
        org_code="ORG_ROOT",
        org_name="机关",
        org_path="/机关",
        sort_no=10,
        parent_id=None,
    )
    created_organizations += int(created)
    ops, created = _organization(
        org_code="ORG_OPS",
        org_name="作战部",
        org_path="/机关/作战部",
        sort_no=20,
        parent_id=root.id,
    )
    created_organizations += int(created)
    _, created = _organization(
        org_code="ORG_COMMS",
        org_name="通信部",
        org_path="/机关/通信部",
        sort_no=30,
        parent_id=root.id,
    )
    created_organizations += int(created)

    created_users += int(
        _user(
            username="admin",
            real_name="系统管理员",
            password=passwords["admin"],
            role_code=RoleCode.ADMIN,
            security_level=SecurityLevel.TOP_SECRET,
            organization=root,
            overwrite_password=overwrite_passwords,
        )
    )
    created_users += int(
        _user(
            username="user",
            real_name="普通用户",
            password=passwords["user"],
            role_code=RoleCode.USER,
            security_level=SecurityLevel.SECRET,
            organization=ops,
            overwrite_password=overwrite_passwords,
        )
    )
    db.session.commit()
    return created_organizations, created_users


def main() -> None:
    app = Flask("dms-seed-dev")
    init_dms(app)
    if _overwrite_passwords_enabled():
        print("警告：已显式启用已有开发用户密码覆盖，并将使其旧Token失效。")
    with app.app_context():
        try:
            organization_count, user_count = seed()
        except Exception:
            db.session.rollback()
            raise
    print(f"初始化完成：新增组织{organization_count}个，新增用户{user_count}个。")


if __name__ == "__main__":
    main()