AstrBot/astrbot/cli/commands/cmd_plug.py

import re
import shutil
from pathlib import Path

import click

from astrbot.cli.i18n import t
from astrbot.cli.utils import (
    PluginStatus,
    build_plug_list,
    get_git_repo,
    install_local_plugin,
    manage_plugin,
)


@click.group(name="plugin")
def plug() -> None:
    """Plugin management"""


def display_plugins(plugins, title=None, color=None) -> None:
    if title:
        click.echo(click.style(title, fg=color, bold=True))

    click.echo(
        f"{'Name':<20} {'Version':<10} {'Status':<10} {'Author':<15} {'Description':<30}",
    )
    click.echo("-" * 85)

    for p in plugins:
        desc = p["desc"][:30] + ("..." if len(p["desc"]) > 30 else "")
        click.echo(
            f"{p['name']:<20} {p['version']:<10} {p['status']:<10} "
            f"{p['author']:<15} {desc:<30}",
        )


@plug.command()
@click.argument("name")
def new(name: str) -> None:
    """Create a new plugin"""
    from astrbot.core.utils.astrbot_path import astrbot_paths

    base_path = astrbot_paths.data
    plug_path = base_path / "plugins" / name

    if plug_path.exists():
        raise click.ClickException(t("plugin_already_exists", name=name))

    author = click.prompt("Enter plugin author", type=str)
    desc = click.prompt("Enter plugin description", type=str)
    version = click.prompt("Enter plugin version", type=str)
    if not re.match(r"^\d+\.\d+(\.\d+)?$", version.lower().lstrip("v")):
        raise click.ClickException("Version must be in x.y or x.y.z format")
    repo = click.prompt("Enter plugin repository URL:", type=str)
    if not repo.startswith("http"):
        raise click.ClickException("Repository URL must start with http")

    click.echo("Downloading plugin template...")
    get_git_repo(
        "https://github.com/Soulter/helloworld",
        plug_path,
    )

    click.echo("Rewriting plugin metadata...")
    # Rewrite metadata.yaml
    with open(plug_path / "metadata.yaml", "w", encoding="utf-8") as f:
        f.write(
            f"name: {name}\n"
            f"desc: {desc}\n"
            f"version: {version}\n"
            f"author: {author}\n"
            f"repo: {repo}\n",
        )

    # Rewrite README.md
    with open(plug_path / "README.md", "w", encoding="utf-8") as f:
        f.write(
            f"# {name}\n\n{desc}\n\n# Support\n\n[Documentation](https://docs.astrbot.app)\n",
        )

    # Rewrite main.py
    with open(plug_path / "main.py", encoding="utf-8") as f:
        content = f.read()

    new_content = content.replace(
        '@register("helloworld", "YourName", "一个简单的 Hello World 插件", "1.0.0")',
        f'@register("{name}", "{author}", "{desc}", "{version}")',
    )

    with open(plug_path / "main.py", "w", encoding="utf-8") as f:
        f.write(new_content)

    click.echo(f"Plugin {name} created successfully")


@plug.command()
@click.option("--all", "-a", is_flag=True, help="List uninstalled plugins")
def list(all: bool) -> None:
    """List plugins"""
    from astrbot.core.utils.astrbot_path import astrbot_paths

    base_path = astrbot_paths.data
    plugins = build_plug_list(base_path / "plugins")

    # Unpublished plugins
    not_published_plugins = [
        p for p in plugins if p["status"] == PluginStatus.NOT_PUBLISHED
    ]
    if not_published_plugins:
        display_plugins(not_published_plugins, "Unpublished Plugins", "red")

    # Plugins needing update
    need_update_plugins = [
        p for p in plugins if p["status"] == PluginStatus.NEED_UPDATE
    ]
    if need_update_plugins:
        display_plugins(need_update_plugins, "Plugins Needing Update", "yellow")

    # Installed plugins
    installed_plugins = [p for p in plugins if p["status"] == PluginStatus.INSTALLED]
    if installed_plugins:
        display_plugins(installed_plugins, "Installed Plugins", "green")

    # Uninstalled plugins
    not_installed_plugins = [
        p for p in plugins if p["status"] == PluginStatus.NOT_INSTALLED
    ]
    if not_installed_plugins and all:
        display_plugins(not_installed_plugins, "Uninstalled Plugins", "blue")

    if (
        not any([not_published_plugins, need_update_plugins, installed_plugins])
        and not all
    ):
        click.echo("No plugins installed")


@plug.command()
@click.argument("name", required=False)
@click.option(
    "--editable",
    "-e",
    "local_path",
    type=click.Path(exists=True, file_okay=False, path_type=Path),
    help="Install a plugin from a local directory as a symlink",
)
@click.option("--proxy", help="Proxy server address")
def install(name: str | None, local_path: Path | None, proxy: str | None) -> None:
    """Install a plugin"""
    from astrbot.core.utils.astrbot_path import astrbot_paths

    base_path = astrbot_paths.data
    plug_path = base_path / "plugins"

    if local_path is not None:
        install_local_plugin(local_path, plug_path, editable=True)
        return

    if name is None:
        raise click.ClickException("Missing plugin name or local plugin path")

    local_name_path = Path(name).expanduser()
    if local_name_path.exists() and local_name_path.is_dir():
        install_local_plugin(local_name_path, plug_path, editable=False)
        return

    plugins = build_plug_list(base_path / "plugins")

    plugin = next(
        (
            p
            for p in plugins
            if p["name"] == name and p["status"] == PluginStatus.NOT_INSTALLED
        ),
        None,
    )

    if not plugin:
        raise click.ClickException(t("plugin_not_found_or_installed", name=name))

    manage_plugin(plugin, plug_path, is_update=False, proxy=proxy)


@plug.command()
@click.argument("name")
def remove(name: str) -> None:
    """Uninstall a plugin"""
    from astrbot.core.utils.astrbot_path import astrbot_paths

    base_path = astrbot_paths.data
    plugins = build_plug_list(base_path / "plugins")
    plugin = next((p for p in plugins if p["name"] == name), None)

    if not plugin or not plugin.get("local_path"):
        raise click.ClickException(t("plugin_not_found_or_installed", name=name))

    plugin_path = plugin["local_path"]

    click.confirm(t("plugin_uninstall_confirm", name=name), default=False, abort=True)

    try:
        shutil.rmtree(plugin_path)
        click.echo(t("plugin_uninstall_success", name=name))
    except Exception as e:
        raise click.ClickException(
            t("plugin_uninstall_failed_ex", name=name, error=str(e)),
        ) from e


@plug.command()
@click.argument("name", required=False)
@click.option("--proxy", help="GitHub proxy address")
def update(name: str, proxy: str | None) -> None:
    """Update plugins"""
    from astrbot.core.utils.astrbot_path import astrbot_paths

    base_path = astrbot_paths.data
    plug_path = base_path / "plugins"
    plugins = build_plug_list(base_path / "plugins")

    if name:
        plugin = next(
            (
                p
                for p in plugins
                if p["name"] == name and p["status"] == PluginStatus.NEED_UPDATE
            ),
            None,
        )

        if not plugin:
            raise click.ClickException(
                f"Plugin {name} does not need updating or cannot be updated",
            )

        manage_plugin(plugin, plug_path, is_update=True, proxy=proxy)
    else:
        need_update_plugins = [
            p for p in plugins if p["status"] == PluginStatus.NEED_UPDATE
        ]

        if not need_update_plugins:
            click.echo(t("plugin_no_update_needed"))
            return

        click.echo(t("plugin_found_update", count=str(len(need_update_plugins))))
        for plugin in need_update_plugins:
            plugin_name = plugin["name"]
            click.echo(t("plugin_updating", name=plugin_name))
            manage_plugin(plugin, plug_path, is_update=True, proxy=proxy)


@plug.command()
@click.argument("query")
def search(query: str) -> None:
    """Search for plugins"""
    from astrbot.core.utils.astrbot_path import astrbot_paths

    base_path = astrbot_paths.data
    plugins = build_plug_list(base_path / "plugins")

    matched_plugins = [
        p
        for p in plugins
        if query.lower() in p["name"].lower()
        or query.lower() in p["desc"].lower()
        or query.lower() in p["author"].lower()
    ]

    if not matched_plugins:
        click.echo(t("plugin_search_no_result", query=query))
        return

    display_plugins(matched_plugins, t("plugin_search_results", query=query), "cyan")