FabledSteward/steward/core/plugin_manager.py

# steward/core/plugin_manager.py
from __future__ import annotations
import hashlib
import logging
import os
import shutil
import sys
import zipfile
from pathlib import Path
from typing import TYPE_CHECKING

import yaml
from packaging.version import Version

if TYPE_CHECKING:
    from quart import Quart

logger = logging.getLogger(__name__)

# Track which plugins were successfully loaded so hot-reload knows
# whether to register a blueprint (new plugin) or skip it (already registered).
_LOADED_PLUGINS: set[str] = set()

# Track plugins that failed to load: name → human-readable reason.
_FAILED_PLUGINS: dict[str, str] = {}

# Nav entries contributed by plugins via the optional get_nav() export, so a
# plugin's UI gets a home in the sidebar. Each entry:
#   {"plugin": <name>, "label": str, "href": str, "section": str}
_PLUGIN_NAV: list[dict] = []

# Loaded plugin modules, by name, for the MCP server: a plugin contributes
# agent tools through an optional register_mcp_tools(register) export, which
# steward/mcp/server.py calls when it mounts (and on hot-reload, below).
_PLUGIN_MODULES: dict[str, object] = {}


def get_plugin_failures() -> dict[str, str]:
    """Return a copy of the failed-plugin registry (name → error message)."""
    return dict(_FAILED_PLUGINS)


def get_plugin_nav() -> list[dict]:
    """Plugin-contributed sidebar entries, sorted by section then label."""
    return sorted(_PLUGIN_NAV, key=lambda e: (e.get("section", ""), e.get("label", "")))


def _collect_plugin_nav(name: str, module) -> None:
    """Pull a plugin's optional get_nav() entries into the nav registry.

    Tolerant by design: a missing hook is fine, and a raising or malformed hook
    is logged and skipped — a bad nav contribution must never break loading.
    Idempotent per plugin (drops prior entries first) so hot-reload re-runs cleanly.
    """
    global _PLUGIN_NAV
    _PLUGIN_NAV = [e for e in _PLUGIN_NAV if e.get("plugin") != name]
    if not hasattr(module, "get_nav"):
        return
    try:
        items = module.get_nav() or []
    except Exception:
        logger.exception("Plugin %r: get_nav() raised, skipping its nav entries", name)
        return
    for item in items:
        try:
            label, href = str(item["label"]), str(item["href"])
        except (TypeError, KeyError):
            logger.warning("Plugin %r: malformed nav item %r, skipping", name, item)
            continue
        _PLUGIN_NAV.append({
            "plugin": name,
            "label": label,
            "href": href,
            "section": str(item.get("section", "Infrastructure")),
        })


def get_loaded_plugin_modules() -> dict[str, object]:
    """Loaded plugin modules by name (a copy), for hooks consumed after loading."""
    return dict(_PLUGIN_MODULES)


def _collect_plugin_module(app: "Quart", name: str, module) -> None:
    """Remember a loaded plugin module; hand it to a live MCP server if mounted.

    At boot the MCP server mounts after plugins load and reads them via
    get_loaded_plugin_modules(); a plugin hot-loaded later is registered here.
    """
    _PLUGIN_MODULES[name] = module
    mcp_server = getattr(app, "mcp_server", None)
    if mcp_server is not None:
        mcp_server.register_plugin(name, module)


def resolve_plugin_path(plugin_dirs: list[Path], name: str) -> Path | None:
    """Return the first plugin root that contains `name`, else None.

    Roots are searched in order, so a bundled (first-party) plugin shadows an
    external plugin of the same name. Pure function — no app/IO beyond exists().
    """
    for pd in plugin_dirs:
        candidate = pd / name
        if candidate.exists():
            return candidate
    return None


def _import_plugin(name: str, plugin_path: Path):
    """Load a plugin module by file path, avoiding sys.modules stdlib collisions.

    Using importlib.import_module(name) fails for plugins whose names shadow
    Python stdlib modules (e.g. a plugin named 'json' or 'http' vs the stdlib
    package of the same name).
    This helper loads from the filesystem path directly and registers the module
    under a namespaced key so relative imports within the plugin still work.
    """
    import importlib.util

    module_key = f"_steward_plugin_{name}"
    if module_key in sys.modules:
        return sys.modules[module_key]

    init_file = plugin_path / "__init__.py"
    spec = importlib.util.spec_from_file_location(
        module_key,
        str(init_file),
        submodule_search_locations=[str(plugin_path)],
    )
    if spec is None or spec.loader is None:
        raise ImportError(f"Could not create import spec for plugin {name!r}")

    module = importlib.util.module_from_spec(spec)
    module.__package__ = module_key
    sys.modules[module_key] = module
    try:
        spec.loader.exec_module(module)
    except Exception:
        sys.modules.pop(module_key, None)
        raise
    return module


def load_plugins(app: "Quart") -> None:
    """Import all enabled plugins and register their blueprints and scheduled tasks.

    For each plugin listed as enabled in config.yaml under 'plugins':
    1. Locate the plugin directory under PLUGIN_DIR
    2. Load and validate plugin.yaml (name must match dir, min_app_version check)
    3. Merge plugin.yaml config defaults with user overrides in app.config["PLUGINS"]
    4. Import the plugin Python package
    5. Validate required exports: setup() and get_scheduled_tasks()
    6. Call setup(app)
    7. Register blueprint at /plugins/<name>/ if get_blueprint() exists
    8. Append get_scheduled_tasks() results to app._task_registry
    """
    import steward

    plugin_dirs = [Path(d) for d in app.config["PLUGIN_DIRS"]]
    plugins_cfg: dict = app.config["PLUGINS"]

    # Ensure every plugin root is on sys.path so plugins are importable by name
    for pd in plugin_dirs:
        pd_str = str(pd.resolve())
        if pd_str not in sys.path:
            sys.path.insert(0, pd_str)

    for name, cfg in list(plugins_cfg.items()):
        if not cfg.get("enabled", False):
            continue

        plugin_path = resolve_plugin_path(plugin_dirs, name)
        if plugin_path is None:
            roots = ", ".join(str(d) for d in plugin_dirs)
            _FAILED_PLUGINS[name] = f"Plugin directory not found in: {roots}"
            logger.error("Plugin %r: not found in any plugin root (%s), skipping", name, roots)
            continue

        # Load and validate plugin.yaml
        yaml_path = plugin_path / "plugin.yaml"
        if not yaml_path.exists():
            _FAILED_PLUGINS[name] = "plugin.yaml not found"
            logger.error("Plugin %r: plugin.yaml not found, skipping", name)
            continue

        try:
            with yaml_path.open() as f:
                meta = yaml.safe_load(f) or {}
        except Exception as exc:
            _FAILED_PLUGINS[name] = f"Failed to read plugin.yaml: {exc}"
            logger.exception("Plugin %r: failed to read plugin.yaml, skipping", name)
            continue

        if meta.get("name") != name:
            _FAILED_PLUGINS[name] = (
                f"plugin.yaml name {meta.get('name')!r} does not match directory name"
            )
            logger.error(
                "Plugin %r: plugin.yaml name=%r does not match directory name, skipping",
                name, meta.get("name"),
            )
            continue

        min_ver = meta.get("min_app_version")
        if min_ver:
            try:
                if Version(steward.__version__) < Version(min_ver):
                    _FAILED_PLUGINS[name] = (
                        f"Requires steward>={min_ver}, running {steward.__version__}"
                    )
                    logger.error(
                        "Plugin %r: requires steward>=%s but running %s, skipping",
                        name, min_ver, steward.__version__,
                    )
                    continue
            except Exception as exc:
                _FAILED_PLUGINS[name] = f"Invalid min_app_version {min_ver!r}: {exc}"
                logger.exception("Plugin %r: invalid min_app_version %r, skipping", name, min_ver)
                continue

        # Merge plugin.yaml config defaults with user overrides (non-"enabled" keys).
        # If a stored value's type doesn't match the default (e.g. a list default
        # was corrupted to a string in the DB), fall back to the yaml default.
        defaults = meta.get("config", {})
        user_overrides = {k: v for k, v in cfg.items() if k != "enabled"}
        merged = {"enabled": True, **defaults}
        for k, v in user_overrides.items():
            default_v = defaults.get(k)
            if default_v is not None and type(v) is not type(default_v):
                logger.warning(
                    "Plugin %r: config key %r has wrong type (%s), using yaml default",
                    name, k, type(v).__name__,
                )
            else:
                merged[k] = v
        plugins_cfg[name] = merged

        # Import the plugin package (file-path load avoids stdlib name collisions)
        try:
            module = _import_plugin(name, plugin_path)
        except Exception as exc:
            _FAILED_PLUGINS[name] = f"Import error: {exc}"
            logger.exception("Plugin %r: import failed, skipping", name)
            continue

        # Validate required exports
        if not hasattr(module, "setup"):
            _FAILED_PLUGINS[name] = "Missing required export: setup()"
            logger.error("Plugin %r: missing required export 'setup', skipping", name)
            continue
        if not hasattr(module, "get_scheduled_tasks"):
            _FAILED_PLUGINS[name] = "Missing required export: get_scheduled_tasks()"
            logger.error(
                "Plugin %r: missing required export 'get_scheduled_tasks', skipping", name
            )
            continue

        # Call setup
        try:
            module.setup(app)
        except Exception as exc:
            _FAILED_PLUGINS[name] = f"setup() failed: {exc}"
            logger.exception("Plugin %r: setup() raised, skipping", name)
            continue

        # Register blueprint (optional)
        if hasattr(module, "get_blueprint"):
            try:
                bp = module.get_blueprint()
                app.register_blueprint(bp, url_prefix=f"/plugins/{name}")
                logger.info("Plugin %r: blueprint registered at /plugins/%s/", name, name)
            except Exception as exc:
                _FAILED_PLUGINS[name] = f"Blueprint registration failed: {exc}"
                logger.exception("Plugin %r: get_blueprint() raised", name)
                continue

        # Register scheduled tasks
        try:
            tasks = module.get_scheduled_tasks()
            app._task_registry.extend(tasks)
            logger.info("Plugin %r: registered %d scheduled task(s)", name, len(tasks))
        except Exception as exc:
            _FAILED_PLUGINS[name] = f"get_scheduled_tasks() failed: {exc}"
            logger.exception("Plugin %r: get_scheduled_tasks() raised", name)
            continue

        _FAILED_PLUGINS.pop(name, None)  # clear any stale failure from a previous reload attempt
        _LOADED_PLUGINS.add(name)
        _collect_plugin_nav(name, module)
        _collect_plugin_module(app, name, module)
        logger.info("Plugin %r loaded (v%s)", name, meta.get("version", "?"))


# Bounds on what an install may bring in: the download itself, and what it
# unpacks to (a small zip can expand to fill the disk).
MAX_PLUGIN_DOWNLOAD_BYTES = 50 * 1024 * 1024
MAX_PLUGIN_UNPACKED_BYTES = 200 * 1024 * 1024


class PluginInstallError(Exception):
    """An install refused or failed; the message is shown to the admin."""


def check_plugin_name(name: str) -> None:
    """A plugin is imported by its folder name, so the name must be a plain
    Python identifier — which also keeps it from naming a path."""
    if not (name.isascii() and name.isidentifier()):
        raise PluginInstallError(f"{name!r} is not a valid plugin name")


def verify_checksum(content: bytes, expected_sha256: str) -> None:
    """Refuse a download whose sha256 is missing from the catalog or differs."""
    if not expected_sha256:
        raise PluginInstallError(
            "The catalog lists no checksum for this plugin, so Steward cannot "
            "tell it is the file the catalog meant. Not installed.")
    actual = hashlib.sha256(content).hexdigest()
    if actual != expected_sha256.strip().lower():
        raise PluginInstallError(f"Checksum mismatch: expected {expected_sha256}, got {actual}")


def extract_plugin(content: bytes, name: str, dest: Path) -> None:
    """Unpack a plugin zip into `dest`, which must not exist yet.

    One top-level folder whose name starts with the plugin name is stripped:
        name/                         (clean plugin zip)
        name-v1.0.0/                  (release tag archive)
    otherwise the zip is unpacked as it stands. Every member must land inside
    `dest` — an absolute path or a `..` that climbs out refuses the whole zip —
    and the unpacked total is capped.
    """
    import io
    root = dest.resolve()
    with zipfile.ZipFile(io.BytesIO(content)) as zf:
        infos = zf.infolist()
        if sum(i.file_size for i in infos) > MAX_PLUGIN_UNPACKED_BYTES:
            raise PluginInstallError(
                f"Unpacks to more than {MAX_PLUGIN_UNPACKED_BYTES // (1024 * 1024)} MB")
        top_level = {i.filename.split("/")[0] for i in infos}
        prefix = ""
        if len(top_level) == 1:
            candidate = next(iter(top_level))
            if candidate.lower().startswith(name.lower()) and any(
                    i.filename.startswith(candidate + "/") for i in infos):
                prefix = candidate + "/"
        planned: list[tuple[zipfile.ZipInfo, Path]] = []
        for info in infos:
            if prefix and not info.filename.startswith(prefix):
                continue
            rel = info.filename[len(prefix):]
            if not rel:
                continue
            target = (root / rel).resolve()
            if not target.is_relative_to(root) or target == root:
                raise PluginInstallError(f"Refused: {info.filename!r} points outside the plugin folder")
            planned.append((info, target))
        root.mkdir(parents=True)
        for info, target in planned:
            if info.is_dir():
                target.mkdir(parents=True, exist_ok=True)
                continue
            target.parent.mkdir(parents=True, exist_ok=True)
            with zf.open(info) as src, open(target, "wb") as out:
                shutil.copyfileobj(src, out)


async def _download(url: str) -> bytes:
    import httpx
    chunks: list[bytes] = []
    size = 0
    async with httpx.AsyncClient(timeout=60, follow_redirects=True) as client:
        async with client.stream("GET", url) as resp:
            resp.raise_for_status()
            async for chunk in resp.aiter_bytes():
                size += len(chunk)
                if size > MAX_PLUGIN_DOWNLOAD_BYTES:
                    raise PluginInstallError(
                        f"Download is larger than {MAX_PLUGIN_DOWNLOAD_BYTES // (1024 * 1024)} MB")
                chunks.append(chunk)
    return b"".join(chunks)


async def download_and_install_plugin(
    app: "Quart",
    name: str,
    download_url: str,
    expected_sha256: str,
) -> tuple[bool, str]:
    """Download a plugin zip named by the catalog, verify its checksum, and
    unpack it into the install dir, replacing any earlier copy whole.

    The URL and checksum must come from the catalog entry, never from the
    request. Returns (success, message). Does NOT hot-reload — call
    hot_reload_plugin() after this to activate a new plugin without a restart.
    """
    # Downloads always land in the external (writable, persistent) install dir,
    # never the read-only bundled dir.
    plugin_dir = Path(app.config["PLUGIN_INSTALL_DIR"]).resolve()

    # Ensure the install dir is on sys.path (may not be set yet if no plugins
    # were enabled at startup)
    plugin_dir_str = str(plugin_dir)
    if plugin_dir_str not in sys.path:
        sys.path.insert(0, plugin_dir_str)

    try:
        check_plugin_name(name)
        if not download_url:
            raise PluginInstallError("The catalog lists no download URL for this plugin")
        logger.info("Downloading plugin %r from %s", name, download_url)
        try:
            content = await _download(download_url)
        except PluginInstallError:
            raise
        except Exception as exc:
            raise PluginInstallError(f"Download failed: {exc}") from exc
        verify_checksum(content, expected_sha256)

        # Unpack beside the live copy, then swap: a failed or refused zip
        # leaves the installed plugin untouched, and an update drops files the
        # new version no longer ships.
        plugin_dir.mkdir(parents=True, exist_ok=True)
        dest = plugin_dir / name
        staging = plugin_dir / f".{name}.installing"
        old = plugin_dir / f".{name}.replaced"
        for leftover in (staging, old):
            shutil.rmtree(leftover, ignore_errors=True)
        try:
            extract_plugin(content, name, staging)
        except PluginInstallError:
            raise
        except Exception as exc:
            raise PluginInstallError(f"Extraction failed: {exc}") from exc
        try:
            if dest.exists():
                dest.rename(old)
            staging.rename(dest)
        except OSError as exc:
            if old.exists() and not dest.exists():
                old.rename(dest)  # put the previous version back
            raise PluginInstallError(f"Could not move the plugin into place: {exc}") from exc
        shutil.rmtree(old, ignore_errors=True)
    except PluginInstallError as exc:
        shutil.rmtree(plugin_dir / f".{name}.installing", ignore_errors=True)
        logger.error("Plugin %r: %s", name, exc)
        return False, str(exc)

    # Run migrations for the newly installed plugin.
    # Include all plugin migration dirs so Alembic can resolve any previously-stamped
    # revisions from other plugins already in alembic_version.
    try:
        mdir = plugin_dir / name / "migrations"
        if mdir.exists():
            from steward.core.migration_runner import (
                run_plugin_migrations,
                discover_all_in,
            )
            # Span every plugin root so Alembic can resolve bundled-plugin
            # revisions already stamped in alembic_version.
            all_dirs = discover_all_in([Path(d).resolve() for d in app.config["PLUGIN_DIRS"]])
            run_plugin_migrations(app.config["DATABASE_URL"], all_dirs)
    except Exception:
        logger.exception("Plugin %r: migration failed after install", name)
        return False, "Plugin extracted but migrations failed — check logs"

    logger.info("Plugin %r installed successfully", name)
    return True, "Installed"


def hot_reload_plugin(app: "Quart", name: str) -> tuple[bool, str]:
    """Activate a newly installed plugin without restarting the app.

    This works for plugins that were NOT previously loaded. For plugins that
    are already registered (blueprint already mounted), a full restart is
    required to pick up code changes — use restart_app() in that case.

    Returns (success, message).
    """
    import steward

    if name in _LOADED_PLUGINS:
        return False, "Plugin already loaded — restart required to apply updates"

    plugin_dirs = [Path(d).resolve() for d in app.config["PLUGIN_DIRS"]]
    plugin_path = resolve_plugin_path(plugin_dirs, name)

    if plugin_path is None:
        return False, f"Plugin {name!r} not found in any plugin root"

    yaml_path = plugin_path / "plugin.yaml"
    if not yaml_path.exists():
        return False, "plugin.yaml not found"

    try:
        with yaml_path.open() as f:
            meta = yaml.safe_load(f) or {}
    except Exception as exc:
        return False, f"Could not read plugin.yaml: {exc}"

    if meta.get("name") != name:
        return False, f"plugin.yaml name mismatch: expected {name!r}, got {meta.get('name')!r}"

    min_ver = meta.get("min_app_version")
    if min_ver:
        try:
            if Version(steward.__version__) < Version(min_ver):
                return False, (
                    f"Plugin requires steward>={min_ver} "
                    f"but running {steward.__version__}"
                )
        except Exception:
            return False, f"Invalid min_app_version: {min_ver!r}"

    # Merge config defaults with any stored user config.
    # Discard stored values whose type doesn't match the yaml default (corruption guard).
    defaults = meta.get("config", {})
    stored_cfg = app.config.get("PLUGINS", {}).get(name, {})
    merged = {"enabled": True, **defaults}
    for k, v in stored_cfg.items():
        if k == "enabled":
            continue
        default_v = defaults.get(k)
        if default_v is not None and type(v) is not type(default_v):
            logger.warning("Plugin %r: config key %r has wrong type, using yaml default", name, k)
        else:
            merged[k] = v
    app.config["PLUGINS"][name] = merged

    # Run migrations if the plugin has them (safe to re-run — alembic is idempotent).
    # Pass ALL plugin migration dirs so Alembic can resolve any previously-stamped
    # revisions from other plugins that are already in alembic_version.
    mdir = plugin_path / "migrations"
    if mdir.exists():
        try:
            from steward.core.migration_runner import (
                run_plugin_migrations,
                discover_all_in,
            )
            all_dirs = discover_all_in(plugin_dirs)
            run_plugin_migrations(app.config["DATABASE_URL"], all_dirs)
        except Exception:
            logger.exception("Plugin %r: migration failed during hot-reload", name)
            return False, "Migration failed — check logs"

    # Import (file-path load avoids stdlib name collisions)
    try:
        module = _import_plugin(name, plugin_path)
    except Exception as exc:
        return False, f"Import failed: {exc}"

    if not hasattr(module, "setup") or not hasattr(module, "get_scheduled_tasks"):
        return False, "Plugin missing required exports (setup, get_scheduled_tasks)"

    try:
        module.setup(app)
    except Exception as exc:
        return False, f"setup() raised: {exc}"

    if hasattr(module, "get_blueprint"):
        try:
            bp = module.get_blueprint()
            app.register_blueprint(bp, url_prefix=f"/plugins/{name}")
        except Exception as exc:
            return False, f"Blueprint registration failed: {exc}"

    try:
        tasks = module.get_scheduled_tasks()
        # De-duplicate: skip tasks whose name already exists in the registry
        existing_names = {t.name for t in app._task_registry}
        new_tasks = [t for t in tasks if t.name not in existing_names]
        app._task_registry.extend(new_tasks)
    except Exception as exc:
        return False, f"get_scheduled_tasks() raised: {exc}"

    _LOADED_PLUGINS.add(name)
    _collect_plugin_nav(name, module)
    _collect_plugin_module(app, name, module)
    logger.info("Plugin %r hot-reloaded (v%s)", name, meta.get("version", "?"))
    return True, f"Plugin activated (v{meta.get('version', '?')})"


def restart_app() -> None:
    """Replace the current process with a fresh copy (in-app restart).

    Used when a plugin update requires a full restart to take effect.
    In Docker the container will be relaunched by the restart policy.
    """
    logger.info("Initiating in-app restart via os.execv")
    os.execv(sys.executable, [sys.executable] + sys.argv)