setup_auth.py

#!/usr/bin/env python3
"""Create or replace the local authentication values in .env safely."""

from __future__ import annotations

import getpass
import os
import re
import secrets
import sys
import tempfile
from pathlib import Path
from typing import Mapping

try:
    import bcrypt
except ImportError:  # pragma: no cover - exercised by the command-line entry point
    bcrypt = None


DEFAULT_USERNAME = "admin"
USERNAME_PATTERN = re.compile(r"[A-Za-z0-9._-]{1,64}")
COMMON_WEAK_PASSWORDS = {
    "123456",
    "password",
    "password123",
    "admin123456",
    "qwerty123456",
}


def validate_username(username: str) -> str:
    """Return a safe username or raise ValueError before writing dotenv syntax."""
    if not USERNAME_PATTERN.fullmatch(username):
        raise ValueError("用户名仅支持 1-64 位字母、数字、点、下划线或连字符")
    return username


def _is_cyclic_sequence(password: str, alphabet: str) -> bool:
    if len(password) < 2 or any(character not in alphabet for character in password):
        return False
    start = alphabet.index(password[0])
    return any(
        all(character == alphabet[(start + direction * index) % len(alphabet)] for index, character in enumerate(password))
        for direction in (1, -1)
    )


def _is_repeated_pattern(password: str) -> bool:
    for pattern_length in range(1, len(password) // 2 + 1):
        if len(password) % pattern_length == 0:
            pattern = password[:pattern_length]
            if pattern * (len(password) // pattern_length) == password:
                return True
    return False


def build_auth_updates(username: str, password: str) -> dict[str, str]:
    """Hash the password and return complete authentication dotenv values."""
    if bcrypt is None:
        raise RuntimeError("缺少 Python bcrypt. 请先安装 media-drop-backend/requirements.txt")
    validate_username(username)
    normalized_password = password.strip().lower()
    if normalized_password in COMMON_WEAK_PASSWORDS:
        raise ValueError("密码过于常见, 请使用其他密码")
    if len(password) < 12:
        raise ValueError("密码至少需要 12 位")
    lower_password = password.lower()
    if _is_cyclic_sequence(password, "0123456789") or _is_cyclic_sequence(
        lower_password, "abcdefghijklmnopqrstuvwxyz"
    ):
        raise ValueError("密码不能是连续序列")
    if _is_repeated_pattern(password):
        raise ValueError("密码不能由重复模式组成")
    password_hash = bcrypt.hashpw(password.encode("utf-8"), bcrypt.gensalt()).decode("ascii")
    return {
        "AUTH_ENABLED": "1",
        "AUTH_USERNAME": username,
        "AUTH_PASSWORD_HASH": password_hash,
        "AUTH_SECRET": secrets.token_hex(32),
        "AUTH_COOKIE_SECURE": "1",
    }


def _assignment_key(line: str) -> str | None:
    stripped = line.strip()
    if not stripped or stripped.startswith("#") or "=" not in line:
        return None
    key, _ = line.split("=", 1)
    return key.strip()


def env_has_auth_config(env_file: Path) -> bool:
    """Whether an existing .env contains any authentication assignment."""
    if not env_file.exists():
        return False
    return any(
        (_assignment_key(line) or "").startswith("AUTH_")
        for line in env_file.read_text(encoding="utf-8").splitlines()
    )


def update_env_file(env_file: Path, template_file: Path, updates: Mapping[str, str]) -> None:
    """Atomically merge auth values into .env while preserving unrelated settings."""
    source_file = env_file if env_file.exists() else template_file
    source = source_file.read_text(encoding="utf-8") if source_file.exists() else ""
    output: list[str] = []
    replaced: set[str] = set()

    for line in source.splitlines():
        key = _assignment_key(line)
        if key in updates:
            if key not in replaced:
                output.append(f"{key}={updates[key]}")
                replaced.add(key)
            continue
        output.append(line)

    output.extend(f"{key}={value}" for key, value in updates.items() if key not in replaced)
    content = "\n".join(output).rstrip("\n") + "\n"

    descriptor, temporary_name = tempfile.mkstemp(prefix=".env.", dir=env_file.parent, text=True)
    try:
        if os.name == "posix":
            os.chmod(temporary_name, 0o600)
        with os.fdopen(descriptor, "w", encoding="utf-8", newline="\n") as temporary_file:
            temporary_file.write(content)
            temporary_file.flush()
            os.fsync(temporary_file.fileno())
        os.replace(temporary_name, env_file)
    finally:
        if os.path.exists(temporary_name):
            os.unlink(temporary_name)


def _confirm_overwrite() -> bool:
    answer = input("检测到现有 AUTH_* 配置, 是否覆盖? [y/N] ").strip().lower()
    return answer in {"y", "yes"}


def main() -> int:
    root = Path(__file__).resolve().parent
    env_file = root / ".env"
    template_file = root / ".env.example"

    if bcrypt is None:
        print("缺少 Python bcrypt. 请先安装 media-drop-backend/requirements.txt", file=sys.stderr)
        return 1
    if env_has_auth_config(env_file) and not _confirm_overwrite():
        print("未修改 .env")
        return 0

    entered_username = input(f"登录用户名 [{DEFAULT_USERNAME}]: ")
    username = entered_username or DEFAULT_USERNAME
    try:
        validate_username(username)
    except ValueError as error:
        print(error, file=sys.stderr)
        return 1

    password = getpass.getpass("登录密码 (至少 12 位): ")
    password_confirmation = getpass.getpass("再次输入密码: ")
    if password != password_confirmation:
        print("两次输入的密码不一致", file=sys.stderr)
        return 1

    try:
        updates = build_auth_updates(username, password)
    except (RuntimeError, ValueError) as error:
        print(error, file=sys.stderr)
        return 1
    finally:
        password = ""
        password_confirmation = ""

    update_env_file(env_file, template_file, updates)
    if os.name == "posix":
        print("登录配置已写入 .env (权限已尽力设为 600).")
    else:
        print("登录配置已写入 .env. Windows 不会自动修改 ACL, 请自行限制文件访问权限.")
    print("请继续填写 DB_* 等部署配置, 再启动服务.")
    return 0


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