PR

ai_shell/__init__.py

@@ -1,6 +1,6 @@
 """AI Shell Suite — An intelligent multi-modal command-line assistant."""
 
-__version__ = "0.2.0"
+__version__ = "0.2.3"
 __author__ = "Jon-Arve Constantine Grønsberg-Ovesen"
 __email__ = "jonovesen@gmail.com"
 __description__ = "AI Shell Suite — Intelligent multi-modal command-line assistant with multi-provider LLM support"

ai_shell/plugins/__init__.py

@@ -0,0 +1,32 @@
+"""
+AI Shell Suite – Plugin system.
+
+Importing this package triggers auto-discovery of all tool plugins.
+"""
+
+from .base import (
+    PluginInfo,
+    ToolPlugin,
+    clear_registry,
+    discover_plugins,
+    get_plugin,
+    list_plugin_info,
+    list_plugins,
+    register_plugin,
+    register_plugin_class,
+)
+
+# Trigger discovery of built-in plugins (metasploit, wapiti, nmap, docker, ...)
+discover_plugins()
+
+__all__ = [
+    "ToolPlugin",
+    "PluginInfo",
+    "register_plugin",
+    "register_plugin_class",
+    "get_plugin",
+    "list_plugins",
+    "list_plugin_info",
+    "discover_plugins",
+    "clear_registry",
+]

ai_shell/plugins/ansible.py

@@ -0,0 +1,32 @@
+"""Ansible automation assistant plugin."""
+
+from __future__ import annotations
+
+from .base import ToolPlugin, register_plugin_class
+
+
+ANSIBLE_SYSTEM_PROMPT = (
+    "You are an expert Ansible automation and configuration-management assistant. "
+    "The user has Ansible tools available (`ansible`, `ansible-playbook`, `ansible-galaxy`, "
+    "`ansible-inventory`, `ansible-vault`). Help them write and run playbooks, manage "
+    "inventories, roles, collections, and vault secrets. "
+    "When you provide a command, enclose it in a ```bash ... ``` markdown block. "
+    "Prefer idempotent, check-mode (`--check`) runs before applying changes. "
+    "Warn before any destructive or production-impacting playbook runs. "
+    "Explain modules, handlers, tags, and variable precedence when helpful."
+)
+
+
+@register_plugin_class
+class AnsiblePlugin(ToolPlugin):
+    id = "ansible"
+    name = "Ansible Assistant"
+    description = "AI help for playbooks, roles, inventory and automation"
+    system_prompt = ANSIBLE_SYSTEM_PROMPT
+    start_command = ["bash"]
+    requires_pty = True
+    color_key = "info"
+
+    def check_available(self) -> bool:
+        import shutil
+        return shutil.which("ansible") is not None or shutil.which("ansible-playbook") is not None

ai_shell/plugins/aws.py

@@ -0,0 +1,31 @@
+"""AWS CLI assistant plugin."""
+
+from __future__ import annotations
+
+from .base import ToolPlugin, register_plugin_class
+
+
+AWS_SYSTEM_PROMPT = (
+    "You are an expert AWS cloud administrator. The user has the `aws` CLI "
+    "(and optionally `aws-vault` / SSO). Help them manage EC2, S3, IAM, Lambda, "
+    "EKS, RDS, CloudFormation, VPC, CloudWatch, and other services. "
+    "When you provide a command, enclose it in a ```bash ... ``` markdown block. "
+    "Prefer read-only operations (`describe`, `list`, `get`) first. "
+    "Always warn before delete, terminate, or high-cost operations. "
+    "Remind the user to confirm the correct profile/region and account."
+)
+
+
+@register_plugin_class
+class AwsPlugin(ToolPlugin):
+    id = "aws"
+    name = "AWS Assistant"
+    description = "AI help for AWS CLI, EC2, S3, IAM, EKS and more"
+    system_prompt = AWS_SYSTEM_PROMPT
+    start_command = ["bash"]
+    requires_pty = True
+    color_key = "info"
+
+    def check_available(self) -> bool:
+        import shutil
+        return shutil.which("aws") is not None

ai_shell/plugins/base.py

@@ -0,0 +1,172 @@
+"""
+Plugin system base classes and registry for AI Shell Suite.
+
+Tool plugins provide specialized interactive assistants (Metasploit, Wapiti,
+nmap, docker, kubectl, etc.). New tools can be added by:
+
+1. Subclassing ``ToolPlugin``
+2. Implementing the required attributes/methods
+3. Placing the module under ``ai_shell/plugins/`` (auto-discovered)
+   or registering manually via ``register_plugin()``.
+"""
+
+from __future__ import annotations
+
+import importlib
+import logging
+import pkgutil
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List, Optional, Type
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class PluginInfo:
+    """Lightweight metadata used for menus and discovery."""
+
+    id: str
+    name: str
+    description: str
+    available: bool = True
+    requires_pty: bool = True
+    color_key: str = "info"  # maps to ui.colors or rich style
+
+
+class ToolPlugin(ABC):
+    """
+    Abstract base class for interactive tool plugins.
+
+    Subclasses must define class attributes and may override methods.
+    """
+
+    # --- Required class attributes ---
+    id: str = ""                     # unique short id, e.g. "metasploit"
+    name: str = ""                   # human name shown in menus
+    description: str = ""            # one-line description
+    system_prompt: str = ""          # LLM system prompt for this tool
+    start_command: List[str] = field(default_factory=lambda: ["bash"])  # command to exec in PTY
+    requires_pty: bool = True
+    color_key: str = "info"          # key used by ui for coloring output
+
+    def __init__(self):
+        if not self.id or not self.name:
+            raise ValueError(f"{self.__class__.__name__} must define 'id' and 'name'")
+
+    # --- Optional overrides ---
+
+    def check_available(self) -> bool:
+        """
+        Return True if the underlying tool is installed and usable.
+        Default implementation tries to run the first element of start_command --version / --help.
+        """
+        import shutil
+        import subprocess
+
+        if not self.start_command:
+            return False
+        binary = self.start_command[0]
+        if shutil.which(binary) is None:
+            return False
+        # Light check – many tools support --version
+        try:
+            subprocess.run(
+                [binary, "--version"],
+                capture_output=True,
+                timeout=5,
+                check=False,
+            )
+            return True
+        except Exception:
+            # Still consider it available if the binary exists
+            return True
+
+    def get_info(self) -> PluginInfo:
+        return PluginInfo(
+            id=self.id,
+            name=self.name,
+            description=self.description,
+            available=self.check_available(),
+            requires_pty=self.requires_pty,
+            color_key=self.color_key,
+        )
+
+    def on_start(self) -> None:
+        """Hook called just before the PTY session starts."""
+        pass
+
+    def on_stop(self) -> None:
+        """Hook called after the PTY session ends."""
+        pass
+
+    def preprocess_command(self, command: str) -> str:
+        """Optional: transform an LLM-suggested command before sending to the tool."""
+        return command
+
+
+# ---------------------------------------------------------------------------
+# Registry
+# ---------------------------------------------------------------------------
+
+_REGISTRY: Dict[str, ToolPlugin] = {}
+
+
+def register_plugin(plugin: ToolPlugin) -> None:
+    """Register a plugin instance (or replace an existing one with the same id)."""
+    if not plugin.id:
+        raise ValueError("Plugin must have a non-empty id")
+    _REGISTRY[plugin.id] = plugin
+    logger.debug("Registered plugin: %s (%s)", plugin.id, plugin.name)
+
+
+def register_plugin_class(cls: Type[ToolPlugin]) -> Type[ToolPlugin]:
+    """Decorator to register a plugin class (instantiates it immediately)."""
+    instance = cls()
+    register_plugin(instance)
+    return cls
+
+
+def get_plugin(plugin_id: str) -> Optional[ToolPlugin]:
+    return _REGISTRY.get(plugin_id)
+
+
+def list_plugins(available_only: bool = False) -> List[ToolPlugin]:
+    plugins = list(_REGISTRY.values())
+    if available_only:
+        plugins = [p for p in plugins if p.check_available()]
+    return sorted(plugins, key=lambda p: p.name.lower())
+
+
+def list_plugin_info(available_only: bool = False) -> List[PluginInfo]:
+    return [p.get_info() for p in list_plugins(available_only=available_only)]
+
+
+def discover_plugins() -> None:
+    """
+    Auto-discover and import all modules under ai_shell.plugins
+    (except base and __init__). Modules should register themselves
+    via the @register_plugin_class decorator or by calling register_plugin().
+    """
+    try:
+        import ai_shell.plugins as plugins_pkg
+    except ImportError:
+        logger.warning("Could not import ai_shell.plugins package")
+        return
+
+    package_path = plugins_pkg.__path__
+    prefix = plugins_pkg.__name__ + "."
+
+    for finder, name, ispkg in pkgutil.iter_modules(package_path, prefix):
+        if name.endswith(".base") or name.endswith(".__init__"):
+            continue
+        try:
+            importlib.import_module(name)
+            logger.debug("Loaded plugin module: %s", name)
+        except Exception as e:
+            logger.warning("Failed to load plugin module %s: %s", name, e)
+
+
+def clear_registry() -> None:
+    """Clear all registered plugins (mainly for tests)."""
+    _REGISTRY.clear()

ai_shell/plugins/docker.py

@@ -0,0 +1,30 @@
+"""Docker assistant plugin."""
+
+from __future__ import annotations
+
+from .base import ToolPlugin, register_plugin_class
+
+
+DOCKER_SYSTEM_PROMPT = (
+    "You are an expert Docker and containerization assistant. The user has the `docker` CLI "
+    "available. Help them manage images, containers, networks, volumes, and Compose workflows. "
+    "Provide clear, safe commands. When you suggest a command, enclose it in a ```bash ... ``` "
+    "markdown block. Prefer non-destructive options and always warn before suggesting "
+    "`docker system prune`, forced removals, or operations that delete data. "
+    "Explain concepts briefly when helpful (layers, volumes vs bind mounts, networks, etc.)."
+)
+
+
+@register_plugin_class
+class DockerPlugin(ToolPlugin):
+    id = "docker"
+    name = "Docker Assistant"
+    description = "AI help for Docker images, containers, Compose & troubleshooting"
+    system_prompt = DOCKER_SYSTEM_PROMPT
+    start_command = ["bash"]
+    requires_pty = True
+    color_key = "info"
+
+    def check_available(self) -> bool:
+        import shutil
+        return shutil.which("docker") is not None

ai_shell/plugins/git.py

@@ -0,0 +1,31 @@
+"""Git version control assistant plugin."""
+
+from __future__ import annotations
+
+from .base import ToolPlugin, register_plugin_class
+
+
+GIT_SYSTEM_PROMPT = (
+    "You are an expert Git and version-control assistant. The user has the `git` CLI. "
+    "Help them with everyday workflows (status, add, commit, branch, merge, rebase, "
+    "stash, cherry-pick, bisect) and more advanced topics (reflog, worktrees, submodules, "
+    "hooks, interactive rebase). "
+    "When you provide a command, enclose it in a ```bash ... ``` markdown block. "
+    "Prefer safe, reversible operations. Always warn before force-push, hard reset, "
+    "or history-rewriting commands. Suggest clear, conventional commit messages when asked."
+)
+
+
+@register_plugin_class
+class GitPlugin(ToolPlugin):
+    id = "git"
+    name = "Git Assistant"
+    description = "AI help for Git workflows, branching, and history"
+    system_prompt = GIT_SYSTEM_PROMPT
+    start_command = ["bash"]
+    requires_pty = True
+    color_key = "info"
+
+    def check_available(self) -> bool:
+        import shutil
+        return shutil.which("git") is not None

ai_shell/plugins/helm.py

@@ -0,0 +1,31 @@
+"""Helm Kubernetes package manager assistant plugin."""
+
+from __future__ import annotations
+
+from .base import ToolPlugin, register_plugin_class
+
+
+HELM_SYSTEM_PROMPT = (
+    "You are an expert Helm and Kubernetes packaging assistant. The user has the "
+    "`helm` CLI. Help them search, install, upgrade, rollback, and manage charts "
+    "and releases. Explain values files, chart dependencies, hooks, and "
+    "release history. "
+    "When you provide a command, enclose it in a ```bash ... ``` markdown block. "
+    "Prefer `helm template` / `helm lint` / dry-run installs before applying. "
+    "Warn before uninstall or upgrade that could disrupt production workloads."
+)
+
+
+@register_plugin_class
+class HelmPlugin(ToolPlugin):
+    id = "helm"
+    name = "Helm Assistant"
+    description = "AI help for Helm charts, releases and Kubernetes packaging"
+    system_prompt = HELM_SYSTEM_PROMPT
+    start_command = ["bash"]
+    requires_pty = True
+    color_key = "info"
+
+    def check_available(self) -> bool:
+        import shutil
+        return shutil.which("helm") is not None