AstrBot/astrbot/cli/commands/cmd_init.py

import asyncio
import json
import os
import re
from pathlib import Path

import click
from filelock import FileLock, Timeout

from astrbot.cli.utils import check_dashboard
from astrbot.core.config.default import DEFAULT_CONFIG
from astrbot.core.utils.astrbot_path import astrbot_paths

from .cmd_conf import (
    _validate_dashboard_password,
    ensure_config_file,
    prompt_dashboard_password,
    set_dashboard_credentials,
)


async def initialize_astrbot(
    astrbot_root: Path,
    *,
    yes: bool,
    backend_only: bool,
    admin_username: str | None,
    admin_password: str | None,
) -> None:
    """Execute AstrBot initialization logic"""
    dot_astrbot = astrbot_root / ".astrbot"

    if not dot_astrbot.exists():
        if yes or click.confirm(
            f"Install AstrBot to this directory? {astrbot_root}",
            default=True,
            abort=True,
        ):
            dot_astrbot.touch()
            click.echo(f"Created {dot_astrbot}")

    paths = {
        "data": astrbot_root / "data",
        "config": astrbot_root / "data" / "config",
        "plugins": astrbot_root / "data" / "plugins",
        "temp": astrbot_root / "data" / "temp",
        "skills": astrbot_root / "data" / "skills",
    }

    for name, path in paths.items():
        path.mkdir(parents=True, exist_ok=True)
        click.echo(
            f"{'Created' if not path.exists() else f'{name} Directory exists'}: {path}"
        )

    config_path = astrbot_root / "data" / "cmd_config.json"
    if not config_path.exists():
        config_path.write_text(
            json.dumps(DEFAULT_CONFIG, ensure_ascii=False, indent=2),
            encoding="utf-8-sig",
        )
        click.echo(f"Created config file: {config_path}")

    # Generate an .env for this instance from the bundled config.template (if available).
    # The generated file will be written to ASTRBOT_ROOT/.env and will be automatically
    # loaded by `astrbot run` (service-config/.env precedence applies).
    ASTRBOT_ROOT = astrbot_root
    env_file = ASTRBOT_ROOT / ".env"
    if not env_file.exists():
        tmpl_candidates = [
            Path("/opt/astrbot/config.template"),
            # project_root may point to the installed package directory; try it as well
            getattr(astrbot_paths, "project_root", Path.cwd()) / "config.template",
            Path.cwd() / "config.template",
        ]
        tmpl = None
        for t in tmpl_candidates:
            try:
                if t.exists():
                    tmpl = t
                    break
            except Exception:
                continue
        if tmpl is not None:
            try:
                txt = tmpl.read_text(encoding="utf-8")
                # Determine instance name for template replacement (fallback to directory name)
                instance_name = astrbot_root.name or "astrbot"
                # Substitute ${VAR} and ${VAR:-default} for INSTANCE_NAME, PORT, ASTRBOT_ROOT
                txt = re.sub(r"\$\{INSTANCE_NAME(:-[^}]*)?\}", instance_name, txt)
                port_val = (
                    os.environ.get("ASTRBOT_PORT") or os.environ.get("PORT") or "8000"
                )
                txt = re.sub(r"\$\{PORT(:-[^}]*)?\}", str(port_val), txt)
                txt = re.sub(r"\$\{ASTRBOT_ROOT(:-[^}]*)?\}", str(ASTRBOT_ROOT), txt)
                header = (
                    f"# Generated from config.template by astrbot init for instance: {instance_name}\n"
                    "# This file will be auto-loaded by 'astrbot run'\n\n"
                )
                env_file.write_text(header + txt, encoding="utf-8")
                env_file.chmod(0o644)
                click.echo(f"Created environment file from template: {env_file}")
            except Exception as e:
                click.echo(f"Warning: failed to generate .env from template: {e!s}")
        else:
            click.echo("No config.template found; skipping .env generation")

    if admin_password and not admin_username:
        raise click.ClickException(
            "--admin-password requires --admin-username to be provided"
        )

    if admin_username:
        password_hash = (
            _validate_dashboard_password(admin_password)
            if admin_password is not None
            else None
        )
        if password_hash is None:
            if yes or os.environ.get("ASTRBOT_SYSTEMD") == "1":
                raise click.ClickException(
                    "Non-interactive init requires --admin-password when --admin-username is set"
                )
            password_hash = prompt_dashboard_password("Dashboard admin password")

        config = ensure_config_file()
        set_dashboard_credentials(
            config,
            username=admin_username.strip(),
            password_hash=password_hash,
        )
        config_path.write_text(
            json.dumps(config, ensure_ascii=False, indent=2),
            encoding="utf-8-sig",
        )
        click.echo(f"Configured dashboard admin username: {admin_username.strip()}")

    if not backend_only and (
        yes
        or click.confirm(
            "是否需要集成式 WebUI?(个人电脑推荐,服务器不推荐)",
            default=True,
        )
    ):
        # 避免在 systemd 模式下因等待输入而阻塞
        if os.environ.get("ASTRBOT_SYSTEMD") == "1":
            click.echo("Systemd detected: Skipping dashboard check.")
        else:
            await check_dashboard(astrbot_root)
    else:
        click.echo("你可以使用在线面版(需支持配置后端)来控制｡")


@click.command()
@click.option("--yes", "-y", is_flag=True, help="Skip confirmation prompts")
@click.option("--backend-only", "-b", is_flag=True, help="Only initialize the backend")
@click.option("--backup", "-f", help="Initialize from backup file", type=str)
@click.option(
    "-u",
    "--admin-username",
    type=str,
    help="Set dashboard admin username during initialization",
)
@click.option(
    "-p",
    "--admin-password",
    type=str,
    help="Set dashboard admin password during initialization without prompting",
)
@click.option(
    "--root",
    help="ASTRBOT root directory to initialize (overrides ASTRBOT_ROOT env)",
    type=str,
)
def init(
    yes: bool,
    backend_only: bool,
    backup: str | None,
    admin_username: str | None,
    admin_password: str | None,
    root: str | None = None,
) -> None:
    """Initialize AstrBot"""
    click.echo("Initializing AstrBot...")

    if os.environ.get("ASTRBOT_SYSTEMD") == "1":
        yes = True
    from astrbot.core.utils.astrbot_path import astrbot_paths

    astrbot_root = Path(root) if root else astrbot_paths.root
    lock_file = astrbot_root / "astrbot.lock"
    lock = FileLock(lock_file, timeout=5)

    try:
        with lock.acquire():
            asyncio.run(
                initialize_astrbot(
                    astrbot_root,
                    yes=yes,
                    backend_only=backend_only,
                    admin_username=admin_username,
                    admin_password=admin_password,
                )
            )

            if backup:
                from .cmd_bk import import_data_command

                click.echo(f"Restoring from backup: {backup}")
                click.get_current_context().invoke(
                    import_data_command, backup_file=backup, yes=True
                )

            click.echo("Done! You can now run 'astrbot run' to start AstrBot")
    except Timeout:
        raise click.ClickException(
            "Cannot acquire lock file. Please check if another instance is running"
        )

    except Exception as e:
        raise click.ClickException(f"Initialization failed: {e!s}")