attocode-code-intel 0.1.0__py3-none-any.whl

attocode_intel/_internal/integrations/mcp/__init__.py

@@ -0,0 +1 @@
+"""Internal code intelligence support; not a public API."""

attocode_intel/_internal/integrations/mcp/client.py

@@ -0,0 +1,253 @@
+"""MCP (Model Context Protocol) client.
+
+Implements a stdio-based JSON-RPC client for communicating
+with MCP servers that provide additional tools.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+import re
+from dataclasses import dataclass, field
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def _expand_env(env: dict[str, str] | None) -> dict[str, str] | None:
+    """Expand ``${VAR}`` references and merge with the current environment.
+
+    When *env* is provided, we start from ``os.environ`` (so the child
+    process inherits PATH, HOME, etc.) and overlay the user-supplied
+    values after expanding any ``${VAR_NAME}`` patterns against
+    ``os.environ``.
+    """
+    if not env:
+        return None
+    base = os.environ.copy()
+    for key, val in env.items():
+        expanded = re.sub(
+            r"\$\{(?:env:)?(\w+)\}",
+            lambda m: os.environ.get(m.group(1), ""),
+            val,
+        )
+        base[key] = expanded
+    return base
+
+
+@dataclass
+class MCPTool:
+    """A tool provided by an MCP server."""
+
+    name: str
+    description: str
+    input_schema: dict[str, Any] = field(default_factory=dict)
+    server_name: str = ""
+
+
+@dataclass
+class MCPCallResult:
+    """Result of calling an MCP tool."""
+
+    success: bool
+    result: Any = None
+    error: str | None = None
+
+
+class MCPClient:
+    """Stdio-based JSON-RPC MCP client.
+
+    Launches an MCP server process, communicates via stdin/stdout
+    using JSON-RPC 2.0 protocol.
+    """
+
+    def __init__(
+        self,
+        server_command: str,
+        server_args: list[str] | None = None,
+        server_name: str = "",
+        env: dict[str, str] | None = None,
+        cwd: str | None = None,
+    ) -> None:
+        self._command = server_command
+        self._args = server_args or []
+        self._server_name = server_name
+        self._env = env
+        self._cwd = cwd
+        self._process: asyncio.subprocess.Process | None = None
+        self._tools: list[MCPTool] = []
+        self._initialized = False
+        self._request_id = 0
+
+    async def connect(self) -> None:
+        """Start the MCP server process and initialize."""
+        self._process = await asyncio.create_subprocess_exec(
+            self._command,
+            *self._args,
+            stdin=asyncio.subprocess.PIPE,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            env=_expand_env(self._env),
+            cwd=self._cwd,
+        )
+
+        # Send initialize request
+        result = await self._request("initialize", {
+            "protocolVersion": "2024-11-05",
+            "capabilities": {},
+            "clientInfo": {"name": "attocode", "version": "0.1.0"},
+        })
+
+        if result is not None:
+            self._initialized = True
+
+            # Send initialized notification
+            await self._notify("notifications/initialized", {})
+
+            # List tools
+            await self._refresh_tools()
+
+    async def disconnect(self) -> None:
+        """Disconnect and terminate the server process."""
+        if self._process:
+            try:
+                if self._process.stdin:
+                    self._process.stdin.close()
+                self._process.terminate()
+                await asyncio.wait_for(self._process.wait(), timeout=5.0)
+            except (TimeoutError, ProcessLookupError):
+                self._process.kill()
+            finally:
+                self._process = None
+                self._initialized = False
+
+    async def call_tool(self, name: str, arguments: dict[str, Any]) -> MCPCallResult:
+        """Call an MCP tool."""
+        if not self._initialized:
+            return MCPCallResult(success=False, error="Not connected")
+
+        try:
+            result = await self._request("tools/call", {
+                "name": name,
+                "arguments": arguments,
+            })
+            if result is None:
+                return MCPCallResult(success=False, error="No response")
+
+            # MCP tool results have a 'content' array
+            content = result.get("content", [])
+            if content and isinstance(content, list):
+                text_parts = [
+                    c.get("text", "")
+                    for c in content
+                    if c.get("type") == "text"
+                ]
+                return MCPCallResult(success=True, result="\n".join(text_parts))
+
+            return MCPCallResult(success=True, result=str(result))
+        except Exception as e:
+            return MCPCallResult(success=False, error=str(e))
+
+    @property
+    def tools(self) -> list[MCPTool]:
+        """Get available tools."""
+        return list(self._tools)
+
+    @property
+    def is_connected(self) -> bool:
+        return self._initialized
+
+    @property
+    def server_name(self) -> str:
+        return self._server_name
+
+    async def _refresh_tools(self) -> None:
+        """Refresh the tool list from the server."""
+        result = await self._request("tools/list", {})
+        if result and "tools" in result:
+            self._tools = [
+                MCPTool(
+                    name=t["name"],
+                    description=t.get("description", ""),
+                    input_schema=t.get("inputSchema", {}),
+                    server_name=self._server_name,
+                )
+                for t in result["tools"]
+            ]
+
+    async def _request(self, method: str, params: dict[str, Any]) -> Any:
+        """Send a JSON-RPC request and wait for response."""
+        if not self._process or not self._process.stdin or not self._process.stdout:
+            return None
+
+        self._request_id += 1
+        req_id = self._request_id
+        msg = {
+            "jsonrpc": "2.0",
+            "id": req_id,
+            "method": method,
+            "params": params,
+        }
+
+        data = json.dumps(msg) + "\n"
+        self._process.stdin.write(data.encode())
+        await self._process.stdin.drain()
+
+        # Read response
+        try:
+            line = await asyncio.wait_for(
+                self._process.stdout.readline(), timeout=30.0
+            )
+            if not line:
+                logger.warning(
+                    "MCP server %s returned empty response for %s",
+                    self._server_name, method,
+                )
+                return None
+            response = json.loads(line.decode())
+            # Validate JSON-RPC response ID matches request
+            resp_id = response.get("id")
+            if resp_id != req_id:
+                logger.warning(
+                    "MCP server %s response ID mismatch for %s: expected %s, got %s",
+                    self._server_name, method, req_id, resp_id,
+                )
+            if "error" in response:
+                err = response["error"]
+                logger.warning(
+                    "MCP server %s RPC error for %s: code=%s message=%s",
+                    self._server_name, method,
+                    err.get("code", "?"), err.get("message", "unknown"),
+                )
+                return None
+            return response.get("result")
+        except TimeoutError:
+            logger.warning(
+                "MCP server %s timed out on %s (30s)",
+                self._server_name, method,
+            )
+            return None
+        except json.JSONDecodeError as exc:
+            logger.warning(
+                "MCP server %s returned invalid JSON for %s: %s",
+                self._server_name, method, exc,
+            )
+            return None
+
+    async def _notify(self, method: str, params: dict[str, Any]) -> None:
+        """Send a JSON-RPC notification (no response expected)."""
+        if not self._process or not self._process.stdin:
+            return
+
+        msg = {
+            "jsonrpc": "2.0",
+            "method": method,
+            "params": params,
+        }
+
+        data = json.dumps(msg) + "\n"
+        self._process.stdin.write(data.encode())
+        await self._process.stdin.drain()

attocode_intel/_internal/integrations/security/__init__.py

@@ -0,0 +1 @@
+"""Internal code intelligence support; not a public API."""

attocode_intel/_internal/integrations/security/custom_rules.py

@@ -0,0 +1,176 @@
+"""Load user-defined security rules from .attocode/rules/*.yaml files.
+
+Custom rules use the same SecurityPattern dataclass as built-in patterns,
+allowing them to be seamlessly merged into the scanning pipeline.
+
+Rule format (YAML):
+    id: my_rule_name
+    pattern: "regex_pattern_here"
+    message: "What was detected"
+    severity: high           # critical | high | medium | low | info
+    cwe: CWE-79             # optional
+    recommendation: "How to fix"
+    languages:               # optional, empty = all languages
+      - python
+      - javascript
+    scan_comments: false     # optional, default false
+    fix:                     # optional autofix
+      search: "old_code("
+      replace: "new_code("
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+from pathlib import Path
+
+from attocode_intel._internal.integrations.security.patterns import (
+    Category,
+    SecurityPattern,
+    Severity,
+)
+
+logger = logging.getLogger(__name__)
+
+_VALID_SEVERITIES = frozenset(s.value for s in Severity)
+
+
+def load_custom_rules(project_dir: str) -> list[SecurityPattern]:
+    """Load custom security rules from .attocode/rules/*.yaml files.
+
+    Args:
+        project_dir: Project root directory.
+
+    Returns:
+        List of SecurityPattern instances from user-defined rules.
+        Invalid rules are logged and skipped.
+    """
+    rules_dir = Path(project_dir) / ".attocode" / "rules"
+    if not rules_dir.is_dir():
+        return []
+
+    patterns: list[SecurityPattern] = []
+    for yaml_file in sorted(rules_dir.glob("*.yaml")):
+        try:
+            file_patterns = _load_rules_file(yaml_file)
+            patterns.extend(file_patterns)
+        except Exception as exc:
+            logger.warning("Failed to load custom rules from %s: %s", yaml_file, exc)
+
+    if patterns:
+        logger.info("Loaded %d custom security rule(s) from %s", len(patterns), rules_dir)
+
+    return patterns
+
+
+def _load_rules_file(path: Path) -> list[SecurityPattern]:
+    """Parse a single YAML rules file into SecurityPattern instances."""
+    try:
+        import yaml  # type: ignore[import-untyped]
+    except ImportError:
+        logger.debug("PyYAML not installed — skipping custom rules file %s", path)
+        return []
+
+    content = path.read_text(encoding="utf-8")
+    data = yaml.safe_load(content)
+
+    if data is None:
+        return []
+
+    # Support both single-rule and multi-rule files
+    if isinstance(data, dict):
+        rules_list = [data]
+    elif isinstance(data, list):
+        rules_list = data
+    else:
+        logger.warning("Invalid rules file %s: expected dict or list", path)
+        return []
+
+    patterns: list[SecurityPattern] = []
+    for i, rule in enumerate(rules_list):
+        try:
+            pattern = _parse_rule(rule, source=f"{path.name}[{i}]")
+            if pattern is not None:
+                patterns.append(pattern)
+        except Exception as exc:
+            logger.warning("Invalid rule in %s at index %d: %s", path.name, i, exc)
+
+    return patterns
+
+
+def _parse_rule(rule: dict, source: str = "") -> SecurityPattern | None:
+    """Parse a single rule dict into a SecurityPattern.
+
+    Required fields: id, pattern, message, severity, recommendation.
+    Optional fields: cwe, languages, scan_comments, fix.
+    """
+    # Validate required fields
+    missing = [f for f in ("id", "pattern", "message", "severity", "recommendation") if f not in rule]
+    if missing:
+        logger.warning("Rule %s missing required fields: %s", source, ", ".join(missing))
+        return None
+
+    # Validate severity
+    severity_str = str(rule["severity"]).lower()
+    if severity_str not in _VALID_SEVERITIES:
+        logger.warning(
+            "Rule %s has invalid severity '%s' (must be one of: %s)",
+            source, rule["severity"], ", ".join(sorted(_VALID_SEVERITIES)),
+        )
+        return None
+
+    # Compile regex
+    try:
+        compiled = re.compile(rule["pattern"])
+    except re.error as exc:
+        logger.warning("Rule %s has invalid regex '%s': %s", source, rule["pattern"], exc)
+        return None
+
+    return SecurityPattern(
+        name=str(rule["id"]),
+        pattern=compiled,
+        severity=Severity(severity_str),
+        category=Category.ANTI_PATTERN,
+        cwe_id=str(rule.get("cwe", "")),
+        message=str(rule["message"]),
+        recommendation=str(rule["recommendation"]),
+        languages=rule.get("languages", []),
+        scan_comments=bool(rule.get("scan_comments", False)),
+    )
+
+
+def get_autofix_from_rules(
+    project_dir: str,
+) -> dict[str, tuple[str, str]]:
+    """Extract autofix templates from custom rules that define a 'fix' field.
+
+    Returns:
+        Dict mapping rule id -> (search, replace) for rules with fix defined.
+    """
+    rules_dir = Path(project_dir) / ".attocode" / "rules"
+    if not rules_dir.is_dir():
+        return {}
+
+    try:
+        import yaml  # type: ignore[import-untyped]
+    except ImportError:
+        return {}
+
+    fixes: dict[str, tuple[str, str]] = {}
+    for yaml_file in sorted(rules_dir.glob("*.yaml")):
+        try:
+            content = yaml_file.read_text(encoding="utf-8")
+            data = yaml.safe_load(content)
+            if data is None:
+                continue
+            rules_list = [data] if isinstance(data, dict) else data if isinstance(data, list) else []
+            for rule in rules_list:
+                if isinstance(rule, dict) and "fix" in rule and "id" in rule:
+                    fix = rule["fix"]
+                    if isinstance(fix, dict) and "search" in fix and "replace" in fix:
+                        fixes[str(rule["id"])] = (str(fix["search"]), str(fix["replace"]))
+        except Exception:
+            continue
+
+    return fixes