shellmate-cli 3.0.0__py3-none-any.whl

shellmate/agents/base.py

@@ -0,0 +1,248 @@
+import asyncio
+import logging
+import uuid
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional
+from enum import Enum
+from abc import ABC, abstractmethod
+
+logger = logging.getLogger("shellmate.agents")
+
+
+class AgentStatus(Enum):
+    STARTING = "starting"
+    READY = "ready"
+    ERROR = "error"
+    STOPPED = "stopped"
+
+
+@dataclass
+class AgentRequest:
+    request_id: str
+    command: str
+    data: Dict[str, Any] = field(default_factory=dict)
+    timeout: float = 5.0
+
+
+@dataclass
+class AgentResponse:
+    request_id: str
+    agent_name: str
+    status: str
+    data: Dict[str, Any] = field(default_factory=dict)
+    error: Optional[str] = None
+    latency_ms: float = 0.0
+
+
+class Agent(ABC):
+    def __init__(self, name: str):
+        self.name = name
+        self._status = AgentStatus.STARTING
+        self._request_queue: asyncio.Queue = asyncio.Queue()
+        self._response_futures: Dict[str, asyncio.Future] = {}
+        self._running = False
+        self._worker_task: Optional[asyncio.Task] = None
+
+    @property
+    def status(self) -> AgentStatus:
+        return self._status
+
+    @abstractmethod
+    async def initialize(self) -> None:
+        pass
+
+    @abstractmethod
+    async def process(self, request: AgentRequest) -> AgentResponse:
+        pass
+
+    @abstractmethod
+    async def shutdown(self) -> None:
+        pass
+
+    async def start(self) -> None:
+        if self._running:
+            return
+        self._running = True
+        await self.initialize()
+        self._status = AgentStatus.READY
+        self._worker_task = asyncio.create_task(self._worker())
+        logger.info(f"Agent {self.name} started")
+
+    async def stop(self) -> None:
+        if not self._running:
+            return
+        self._running = False
+        if self._worker_task:
+            self._worker_task.cancel()
+            try:
+                await self._worker_task
+            except asyncio.CancelledError:
+                pass
+        await self.shutdown()
+        self._status = AgentStatus.STOPPED
+        logger.info(f"Agent {self.name} stopped")
+
+    async def _worker(self) -> None:
+        while self._running:
+            try:
+                request = await asyncio.wait_for(
+                    self._request_queue.get(), timeout=1.0
+                )
+                await self._handle_request(request)
+            except asyncio.TimeoutError:
+                continue
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"Agent {self.name} worker error: {e}", exc_info=True)
+
+    async def _handle_request(self, request: AgentRequest) -> None:
+        response_future = self._response_futures.get(request.request_id)
+        if not response_future:
+            logger.warning(f"Agent {self.name}: No response future for request {request.request_id}")
+            return
+
+        start_time = asyncio.get_event_loop().time()
+        try:
+            response = await asyncio.wait_for(
+                self.process(request), timeout=request.timeout
+            )
+        except asyncio.TimeoutError:
+            response = AgentResponse(
+                request_id=request.request_id,
+                agent_name=self.name,
+                status="timeout",
+                error=f"Agent {self.name} timed out after {request.timeout}s",
+                latency_ms=(asyncio.get_event_loop().time() - start_time) * 1000,
+            )
+        except Exception as e:
+            logger.error(f"Agent {self.name} process error: {e}", exc_info=True)
+            response = AgentResponse(
+                request_id=request.request_id,
+                agent_name=self.name,
+                status="error",
+                error=str(e),
+                latency_ms=(asyncio.get_event_loop().time() - start_time) * 1000,
+            )
+        else:
+            response.latency_ms = (asyncio.get_event_loop().time() - start_time) * 1000
+
+        if not response_future.done():
+            response_future.set_result(response)
+
+    async def send_request(self, command: str, data: Dict[str, Any] = None, timeout: float = 5.0) -> AgentResponse:
+        if not self._running or self._status != AgentStatus.READY:
+            return AgentResponse(
+                request_id=str(uuid.uuid4()),
+                agent_name=self.name,
+                status="error",
+                error=f"Agent {self.name} not ready",
+            )
+
+        request_id = str(uuid.uuid4())
+        request = AgentRequest(
+            request_id=request_id,
+            command=command,
+            data=data or {},
+            timeout=timeout,
+        )
+
+        response_future = asyncio.get_running_loop().create_future()
+        self._response_futures[request_id] = response_future
+
+        try:
+            await self._request_queue.put(request)
+            response = await asyncio.wait_for(response_future, timeout=timeout + 1.0)
+            return response
+        except asyncio.TimeoutError:
+            if not response_future.done():
+                response_future.cancel()
+            return AgentResponse(
+                request_id=request_id,
+                agent_name=self.name,
+                status="timeout",
+                error=f"Request timed out after {timeout}s",
+            )
+        finally:
+            self._response_futures.pop(request_id, None)
+
+
+class AgentManager:
+    def __init__(self):
+        self._agents: Dict[str, Agent] = {}
+        self._running = False
+
+    def register(self, agent: Agent) -> None:
+        self._agents[agent.name] = agent
+
+    def get(self, name: str) -> Optional[Agent]:
+        return self._agents.get(name)
+
+    async def start_all(self) -> None:
+        if self._running:
+            return
+        self._running = True
+        for agent in self._agents.values():
+            await agent.start()
+        logger.info("All agents started")
+
+    async def stop_all(self) -> None:
+        if not self._running:
+            return
+        for agent in self._agents.values():
+            await agent.stop()
+        self._running = False
+        logger.info("All agents stopped")
+
+    async def send_request(self, agent_name: str, command: str, data: Dict[str, Any] = None, timeout: float = 5.0) -> AgentResponse:
+        agent = self._agents.get(agent_name)
+        if not agent:
+            return AgentResponse(
+                request_id=str(uuid.uuid4()),
+                agent_name=agent_name,
+                status="error",
+                error=f"Agent {agent_name} not found",
+            )
+        return await agent.send_request(command, data, timeout)
+
+    async def send_concurrent(self, requests: Dict[str, tuple], timeout: float = 5.0) -> Dict[str, AgentResponse]:
+        tasks = {}
+        for agent_name, (command, data) in requests.items():
+            agent = self._agents.get(agent_name)
+            if agent:
+                tasks[agent_name] = asyncio.create_task(
+                    agent.send_request(command, data, timeout)
+                )
+
+        results = {}
+        for agent_name, task in tasks.items():
+            try:
+                results[agent_name] = await asyncio.wait_for(task, timeout=timeout + 1.0)
+            except asyncio.TimeoutError:
+                results[agent_name] = AgentResponse(
+                    request_id=str(uuid.uuid4()),
+                    agent_name=agent_name,
+                    status="timeout",
+                    error=f"Request timed out after {timeout}s",
+                )
+            except Exception as e:
+                results[agent_name] = AgentResponse(
+                    request_id=str(uuid.uuid4()),
+                    agent_name=agent_name,
+                    status="error",
+                    error=str(e),
+                )
+        return results
+
+    def get_status(self) -> Dict[str, str]:
+        return {name: agent.status.value for name, agent in self._agents.items()}
+
+
+_agent_manager: Optional[AgentManager] = None
+
+
+def get_agent_manager() -> AgentManager:
+    global _agent_manager
+    if _agent_manager is None:
+        _agent_manager = AgentManager()
+    return _agent_manager

shellmate/agents/chain_patterns.yaml

@@ -0,0 +1,98 @@
+patterns:
+  # Terraform workflows
+  - name: terraform_workflow
+    sequence:
+      - pattern: "terraform plan"
+        required: true
+      - pattern: "terraform apply"
+        expected: true
+    risk_if_broken: HIGH
+    description: "Terraform apply should follow plan"
+  
+  # Git workflows
+  - name: git_force_push_workflow
+    sequence:
+      - pattern: "git (pull|fetch)"
+        required: true
+      - pattern: "git push.*--force"
+        expected: true
+    risk_if_broken: CRITICAL
+    description: "Force push should follow pull/fetch"
+  
+  - name: git_commit_workflow
+    sequence:
+      - pattern: "git add"
+        required: true
+      - pattern: "git commit"
+        expected: true
+      - pattern: "git push"
+        expected: true
+    risk_if_broken: MEDIUM
+    description: "Standard git workflow: add → commit → push"
+  
+  # Kubernetes workflows
+  - name: kubernetes_apply_workflow
+    sequence:
+      - pattern: "kubectl (create|apply) configmap"
+        required: false
+      - pattern: "kubectl apply.*deployment"
+        expected: true
+    risk_if_broken: HIGH
+    description: "Apply deployment after configmap"
+  
+  # Docker workflows
+  - name: docker_build_push_workflow
+    sequence:
+      - pattern: "docker build"
+        required: true
+      - pattern: "docker tag"
+        required: true
+      - pattern: "docker push"
+        expected: true
+    risk_if_broken: MEDIUM
+    description: "Docker image should be built and tagged before push"
+  
+  - name: docker_compose_workflow
+    sequence:
+      - pattern: "docker-compose (build|pull)"
+        required: true
+      - pattern: "docker-compose up"
+        expected: true
+    risk_if_broken: LOW
+    description: "Docker compose up should follow build/pull"
+  
+  # Database workflows
+  - name: database_backup_workflow
+    sequence:
+      - pattern: "(pg_dump|mysqldump|mongodump)"
+        required: true
+      - pattern: "DROP DATABASE"
+        expected: true
+    risk_if_broken: CRITICAL
+    description: "Database should be backed up before dropping"
+  
+  # Dangerous escalations
+  - name: dangerous_escalation_rm
+    sequence:
+      - pattern: "cd (/|/home|/var|/etc|/usr)"
+        required: true
+      - pattern: "rm.*-rf"
+        dangerous: true
+    risk_if_detected: CRITICAL
+    description: "rm -rf in sensitive directory"
+  
+  - name: dangerous_escalation_chmod
+    sequence:
+      - pattern: "chmod (777|666)"
+        dangerous: true
+    risk_if_detected: CRITICAL
+    description: "Overly permissive chmod"
+  
+  - name: dangerous_sudo_escalation
+    sequence:
+      - pattern: "sudo su"
+        required: true
+      - pattern: "(rm|del|DROP|truncate)"
+        dangerous: true
+    risk_if_detected: CRITICAL
+    description: "Destructive command after sudo su"

shellmate/agents/command_chain.py

@@ -0,0 +1,322 @@
+"""
+CommandChainAgent: Tracks command history and detects workflow patterns.
+
+Maintains a persistent JSONL history of commands and analyzes patterns
+to detect expected workflows, broken workflows, and dangerous escalations.
+"""
+
+import asyncio
+import json
+import re
+import time
+from datetime import datetime, timedelta
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import yaml
+
+from shellmate.agents.base import Agent, AgentRequest, AgentResponse
+from shellmate.config.defaults import get_default_config
+from shellmate.daemon.session import get_command_history_file
+
+
+class CommandChainAgent(Agent):
+    """Tracks command history and detects workflow patterns."""
+
+    def __init__(self):
+        super().__init__(name="command_chain")
+        self.history_file = get_command_history_file()
+        self.patterns = []
+        self.patterns_file = Path(__file__).parent / "chain_patterns.yaml"
+
+    async def initialize(self) -> None:
+        """Initialize the agent and load pattern definitions."""
+        self._load_patterns()
+
+    async def shutdown(self) -> None:
+        """Shutdown the agent."""
+        pass
+
+    async def process(self, request: AgentRequest) -> AgentResponse:
+        """Process agent requests by dispatching to appropriate methods."""
+        start_time = time.time()
+        
+        try:
+            if request.command == "append_command":
+                data = request.data
+                await self.append_command(
+                    data.get("command", ""),
+                    data.get("risk", "MEDIUM"),
+                    data.get("exit_code")
+                )
+                latency = (time.time() - start_time) * 1000
+                return AgentResponse(
+                    request.request_id,
+                    self.name,
+                    status="ok",
+                    data={"success": True},
+                    latency_ms=latency,
+                )
+            elif request.command == "get_chain_patterns":
+                data = request.data
+                recent_commands = await self.get_recent_commands(
+                    count=data.get("count", 10),
+                    minutes=data.get("minutes", 30)
+                )
+                current_command = data.get("command", "")
+                patterns = await self.detect_patterns(recent_commands, current_command)
+                latency = (time.time() - start_time) * 1000
+                return AgentResponse(
+                    request.request_id,
+                    self.name,
+                    status="ok",
+                    data={
+                        "command_chain": recent_commands,
+                        "chain_patterns": patterns,
+                    },
+                    latency_ms=latency,
+                )
+            else:
+                return AgentResponse(
+                    request.request_id,
+                    self.name,
+                    status="error",
+                    error=f"Unknown command: {request.command}",
+                )
+        except Exception as e:
+            latency = (time.time() - start_time) * 1000
+            return AgentResponse(
+                request.request_id,
+                self.name,
+                status="error",
+                error=str(e),
+                latency_ms=latency,
+            )
+
+    def _load_patterns(self) -> None:
+        """Load workflow patterns from YAML file."""
+        if not self.patterns_file.exists():
+            self.patterns = []
+            return
+
+        try:
+            with open(self.patterns_file, "r", encoding="utf-8") as f:
+                data = yaml.safe_load(f)
+                self.patterns = data.get("patterns", [])
+        except Exception:
+            self.patterns = []
+
+    def _redact_secrets(self, command: str) -> str:
+        """Redact secrets from command before storage."""
+        config = get_default_config()
+        secret_patterns = (
+            config.get("v3", {})
+            .get("command_history", {})
+            .get("secret_patterns", [])
+        )
+
+        redacted = command
+        for pattern in secret_patterns:
+            try:
+                redacted = re.sub(pattern, r"\1***", redacted)
+            except Exception:
+                pass
+
+        return redacted
+
+    async def append_command(
+        self, command: str, risk: str, exit_code: Optional[int] = None
+    ) -> None:
+        """
+        Append command to history file (JSONL format).
+
+        Args:
+            command: The command that was executed
+            risk: Risk level (SAFE, LOW, MEDIUM, HIGH, CRITICAL)
+            exit_code: Command exit code (optional)
+        """
+        entry = {
+            "timestamp": datetime.utcnow().isoformat() + "Z",
+            "command": self._redact_secrets(command),
+            "risk": risk,
+            "exit_code": exit_code,
+            "directory": str(Path.cwd()),
+        }
+
+        # Append to JSONL file
+        try:
+            with open(self.history_file, "a", encoding="utf-8") as f:
+                f.write(json.dumps(entry) + "\n")
+        except OSError:
+            # Fail silently if we can't write to history
+            pass
+
+    async def get_recent_commands(
+        self, count: int = 10, minutes: Optional[int] = None
+    ) -> List[Dict[str, Any]]:
+        """
+        Query recent commands from history.
+
+        Args:
+            count: Maximum number of commands to return
+            minutes: Optional time window in minutes
+
+        Returns:
+            List of command entries
+        """
+        if not self.history_file.exists():
+            return []
+
+        commands = []
+        cutoff = None
+        if minutes:
+            cutoff = datetime.utcnow() - timedelta(minutes=minutes)
+
+        try:
+            with open(self.history_file, "r", encoding="utf-8") as f:
+                for line in f:
+                    try:
+                        entry = json.loads(line.strip())
+                        timestamp = datetime.fromisoformat(
+                            entry["timestamp"].replace("Z", "")
+                        )
+
+                        if cutoff and timestamp < cutoff:
+                            continue
+
+                        commands.append(entry)
+                    except (json.JSONDecodeError, KeyError, ValueError):
+                        continue
+        except OSError:
+            return []
+
+        # Return last N commands
+        return commands[-count:]
+
+    async def detect_patterns(
+        self, recent_commands: List[Dict[str, Any]], current_command: str
+    ) -> Dict[str, Any]:
+        """
+        Detect workflow patterns in command chain.
+
+        Args:
+            recent_commands: List of recent command entries
+            current_command: The command being analyzed
+
+        Returns:
+            Dictionary with detected, broken, and dangerous_escalations lists
+        """
+        detected_patterns = {
+            "detected": [],
+            "broken": [],
+            "dangerous_escalations": [],
+        }
+
+        for pattern_def in self.patterns:
+            match_result = self._match_pattern(
+                pattern_def, recent_commands, current_command
+            )
+
+            if match_result["type"] == "expected_workflow":
+                detected_patterns["detected"].append(
+                    {
+                        "name": pattern_def["name"],
+                        "description": pattern_def["description"],
+                        "risk_adjustment": -1,  # Lower risk
+                    }
+                )
+            elif match_result["type"] == "broken_workflow":
+                detected_patterns["broken"].append(
+                    {
+                        "name": pattern_def["name"],
+                        "description": pattern_def["description"],
+                        "risk_adjustment": +1,  # Higher risk
+                        "missing_step": match_result["missing_step"],
+                    }
+                )
+            elif match_result["type"] == "dangerous_escalation":
+                detected_patterns["dangerous_escalations"].append(
+                    {
+                        "name": pattern_def["name"],
+                        "description": pattern_def["description"],
+                        "risk_adjustment": +2,  # Much higher risk
+                        "previous_command": match_result.get("previous_command"),
+                    }
+                )
+
+        return detected_patterns
+
+    def _match_pattern(
+        self,
+        pattern_def: Dict[str, Any],
+        history: List[Dict[str, Any]],
+        current: str,
+    ) -> Dict[str, str]:
+        """
+        Match a single pattern definition against history.
+
+        Args:
+            pattern_def: Pattern definition from YAML
+            history: List of recent commands
+            current: Current command being analyzed
+
+        Returns:
+            Dictionary with match type and additional info
+        """
+        sequence = pattern_def.get("sequence", [])
+        if not sequence:
+            return {"type": "no_match"}
+
+        matched_steps = []
+
+        # Check if current command matches final step
+        final_step = sequence[-1]
+        try:
+            if not re.search(final_step["pattern"], current):
+                return {"type": "no_match"}
+        except re.error:
+            return {"type": "no_match"}
+
+        # Check if required prior steps are in history
+        for step in sequence[:-1]:
+            if not step.get("required", False):
+                continue
+
+            found = False
+            for entry in reversed(history):
+                try:
+                    if re.search(step["pattern"], entry["command"]):
+                        matched_steps.append(entry)
+                        found = True
+                        break
+                except re.error:
+                    pass
+
+            if not found:
+                return {
+                    "type": "broken_workflow",
+                    "missing_step": step["pattern"],
+                }
+
+        # Check for dangerous escalations
+        if pattern_def.get("risk_if_detected") == "CRITICAL":
+            return {
+                "type": "dangerous_escalation",
+                "previous_command": (
+                    matched_steps[0]["command"] if matched_steps else None
+                ),
+            }
+
+        return {"type": "expected_workflow"}
+
+
+# Singleton instance
+_command_chain_agent: Optional[CommandChainAgent] = None
+
+
+def get_command_chain_agent() -> CommandChainAgent:
+    """Get the singleton CommandChainAgent instance."""
+    global _command_chain_agent
+    if _command_chain_agent is None:
+        _command_chain_agent = CommandChainAgent()
+    return _command_chain_agent