FabledSteward/steward/core/plugin_manager.py

# steward/core/plugin_manager.py
from __future__ import annotations
import hashlib
import importlib
import logging
import os
import sys
import tempfile
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] = {}


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


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. the 'http' plugin vs stdlib's 'http' package).
    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)
        logger.info("Plugin %r loaded (v%s)", name, meta.get("version", "?"))


async def download_and_install_plugin(
    app: "Quart",
    name: str,
    download_url: str,
    expected_sha256: str,
) -> tuple[bool, str]:
    """Download a plugin zip, verify its checksum, and extract it to PLUGIN_DIR.

    Returns (success, message). Does NOT hot-reload — call hot_reload_plugin()
    after this to activate the plugin without a restart (new plugins only).
    """
    import httpx

    # 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:
        logger.info("Downloading plugin %r from %s", name, download_url)
        async with httpx.AsyncClient(timeout=60, follow_redirects=True) as client:
            resp = await client.get(download_url)
            resp.raise_for_status()
        content = resp.content
    except Exception as exc:
        msg = f"Download failed: {exc}"
        logger.error("Plugin %r: %s", name, msg)
        return False, msg

    # Verify checksum if provided
    if expected_sha256:
        actual = hashlib.sha256(content).hexdigest()
        if actual != expected_sha256.lower():
            msg = f"Checksum mismatch: expected {expected_sha256}, got {actual}"
            logger.error("Plugin %r: %s", name, msg)
            return False, msg

    # Extract the zip
    try:
        with tempfile.TemporaryDirectory() as tmpdir:
            zip_path = Path(tmpdir) / f"{name}.zip"
            zip_path.write_bytes(content)
            with zipfile.ZipFile(zip_path) as zf:
                members = zf.namelist()
                # Collect unique top-level names (dirs and loose files)
                top_level = {m.split("/")[0] for m in members}

                dest = plugin_dir / name
                dest.mkdir(parents=True, exist_ok=True)

                # Detect a single top-level directory to strip. Handles:
                #   name/               (clean plugin zip)
                #   name-v1.0.0/        (GitHub release tag archive)
                #   steward-plugins-main/name/  (whole-repo archive)
                # Strategy: if there is exactly one top-level item and it is a
                # directory whose name starts with the plugin name (case-insensitive),
                # strip it. Otherwise extract flat into dest.
                prefix: str | None = None
                if len(top_level) == 1:
                    candidate = list(top_level)[0]
                    if candidate.lower().startswith(name.lower()):
                        prefix = candidate + "/"

                for member in members:
                    if prefix:
                        if not member.startswith(prefix):
                            continue
                        rel = member[len(prefix):]
                    else:
                        rel = member
                    if not rel or rel.endswith("/"):
                        if rel:
                            (dest / rel.rstrip("/")).mkdir(parents=True, exist_ok=True)
                        continue
                    target = dest / rel
                    target.parent.mkdir(parents=True, exist_ok=True)
                    target.write_bytes(zf.read(member))
    except Exception as exc:
        msg = f"Extraction failed: {exc}"
        logger.error("Plugin %r: %s", name, msg)
        return False, msg

    # 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)
    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)