agent-audit 0.1.0__py3-none-any.whl

agent_audit/utils/mcp_client.py

@@ -0,0 +1,343 @@
+"""MCP protocol client with STDIO and SSE transports."""
+
+import asyncio
+import json
+import logging
+import sys
+from typing import Dict, Any, Optional
+from abc import ABC, abstractmethod
+from enum import Enum
+
+logger = logging.getLogger(__name__)
+
+# Windows subprocess creation flags
+_IS_WINDOWS = sys.platform == "win32"
+_CREATE_NO_WINDOW = 0x08000000 if _IS_WINDOWS else 0
+
+
+class TransportType(Enum):
+    """MCP transport types."""
+    STDIO = "stdio"
+    SSE = "sse"
+    STREAMABLE_HTTP = "streamable_http"
+
+
+class BaseMCPTransport(ABC):
+    """Base class for MCP transport implementations."""
+
+    @abstractmethod
+    async def connect(self):
+        """Establish connection to MCP server."""
+        pass
+
+    @abstractmethod
+    async def send(self, method: str, params: Dict[str, Any]) -> Dict[str, Any]:
+        """Send a JSON-RPC request and wait for response."""
+        pass
+
+    @abstractmethod
+    async def notify(self, method: str, params: Dict[str, Any]):
+        """Send a JSON-RPC notification (no response expected)."""
+        pass
+
+    @abstractmethod
+    async def close(self):
+        """Close the connection."""
+        pass
+
+
+class StdioTransport(BaseMCPTransport):
+    """
+    STDIO transport for local MCP servers.
+
+    Communicates with MCP server via stdin/stdout using JSON-RPC.
+    """
+
+    def __init__(self, command: str, args: Optional[list] = None, env: Optional[dict] = None):
+        """
+        Initialize STDIO transport.
+
+        Args:
+            command: Command to execute (e.g., "python", "node")
+            args: Command arguments (e.g., ["server.py"])
+            env: Environment variables
+        """
+        self.command = command
+        self.args = args or []
+        self.env = env
+        self.process: Optional[asyncio.subprocess.Process] = None
+        self._request_id = 0
+        self._pending_requests: Dict[int, asyncio.Future] = {}
+        self._reader_task: Optional[asyncio.Task] = None
+
+    async def connect(self):
+        """Start the MCP server process."""
+        import os
+
+        # Prepare environment
+        process_env = os.environ.copy()
+        if self.env:
+            process_env.update(self.env)
+
+        # Build kwargs for subprocess creation
+        kwargs: Dict[str, Any] = {
+            "stdin": asyncio.subprocess.PIPE,
+            "stdout": asyncio.subprocess.PIPE,
+            "stderr": asyncio.subprocess.PIPE,
+            "env": process_env,
+        }
+
+        # On Windows, prevent console window from appearing
+        if _IS_WINDOWS:
+            kwargs["creationflags"] = _CREATE_NO_WINDOW
+
+        self.process = await asyncio.create_subprocess_exec(
+            self.command, *self.args,
+            **kwargs
+        )
+
+        # Start background reader
+        self._reader_task = asyncio.create_task(self._read_responses())
+
+    async def _read_responses(self):
+        """Background task to read responses from server."""
+        try:
+            while self.process and self.process.stdout:
+                line = await self.process.stdout.readline()
+                if not line:
+                    break
+
+                try:
+                    response = json.loads(line.decode())
+                    request_id = response.get('id')
+
+                    if request_id is not None and request_id in self._pending_requests:
+                        future = self._pending_requests.pop(request_id)
+                        if not future.done():
+                            future.set_result(response)
+
+                except json.JSONDecodeError:
+                    logger.warning(f"Invalid JSON from server: {line}")
+
+        except asyncio.CancelledError:
+            pass
+        except Exception as e:
+            logger.error(f"Error reading from server: {e}")
+
+    async def send(self, method: str, params: Dict[str, Any]) -> Dict[str, Any]:
+        """Send a JSON-RPC request and wait for response."""
+        if not self.process or not self.process.stdin:
+            raise RuntimeError("Not connected")
+
+        self._request_id += 1
+        request_id = self._request_id
+
+        request = {
+            "jsonrpc": "2.0",
+            "id": request_id,
+            "method": method,
+            "params": params
+        }
+
+        # Create future for response
+        future: asyncio.Future = asyncio.get_event_loop().create_future()
+        self._pending_requests[request_id] = future
+
+        # Send request
+        request_bytes = json.dumps(request).encode() + b"\n"
+        self.process.stdin.write(request_bytes)
+        await self.process.stdin.drain()
+
+        # Wait for response with timeout
+        try:
+            response = await asyncio.wait_for(future, timeout=30)
+        except asyncio.TimeoutError:
+            self._pending_requests.pop(request_id, None)
+            raise TimeoutError(f"Request {method} timed out")
+
+        if "error" in response:
+            error = response["error"]
+            raise RuntimeError(f"MCP Error: {error.get('message', error)}")
+
+        return response.get("result", {})
+
+    async def notify(self, method: str, params: Dict[str, Any]):
+        """Send a JSON-RPC notification."""
+        if not self.process or not self.process.stdin:
+            raise RuntimeError("Not connected")
+
+        notification = {
+            "jsonrpc": "2.0",
+            "method": method,
+            "params": params
+        }
+
+        notification_bytes = json.dumps(notification).encode() + b"\n"
+        self.process.stdin.write(notification_bytes)
+        await self.process.stdin.drain()
+
+    async def close(self):
+        """Terminate the server process."""
+        if self._reader_task:
+            self._reader_task.cancel()
+            try:
+                await self._reader_task
+            except asyncio.CancelledError:
+                pass
+
+        if self.process:
+            self.process.terminate()
+            try:
+                await asyncio.wait_for(self.process.wait(), timeout=5)
+            except asyncio.TimeoutError:
+                self.process.kill()
+                await self.process.wait()
+
+
+class SSETransport(BaseMCPTransport):
+    """
+    SSE (Server-Sent Events) transport for remote MCP servers.
+
+    Uses HTTP POST for requests and SSE for responses.
+    """
+
+    def __init__(self, url: str):
+        """
+        Initialize SSE transport.
+
+        Args:
+            url: SSE endpoint URL
+        """
+        self.url = url
+        self.session = None
+        self._request_id = 0
+        self._endpoint: Optional[str] = None
+
+    async def connect(self):
+        """Connect to the SSE endpoint and get the messages URL."""
+        import aiohttp
+
+        self.session = aiohttp.ClientSession()
+
+        # Connect to SSE endpoint to get messages URL
+        try:
+            async with self.session.get(self.url) as response:
+                if response.status != 200:
+                    raise RuntimeError(f"SSE connection failed: {response.status}")
+
+                # Read SSE events to find endpoint
+                async for line in response.content:
+                    decoded = line.decode().strip()
+
+                    if decoded.startswith("event: endpoint"):
+                        next_line = await response.content.readline()
+                        data = next_line.decode().strip()
+                        if data.startswith("data: "):
+                            self._endpoint = data[6:]
+                            break
+
+        except Exception as e:
+            await self.close()
+            raise RuntimeError(f"Failed to connect to SSE: {e}")
+
+        if not self._endpoint:
+            raise RuntimeError("Failed to get messages endpoint from SSE")
+
+    async def send(self, method: str, params: Dict[str, Any]) -> Dict[str, Any]:
+        """Send a JSON-RPC request via HTTP POST."""
+        if not self.session or not self._endpoint:
+            raise RuntimeError("Not connected")
+
+        self._request_id += 1
+        request = {
+            "jsonrpc": "2.0",
+            "id": self._request_id,
+            "method": method,
+            "params": params
+        }
+
+        async with self.session.post(self._endpoint, json=request) as response:
+            if response.status != 200:
+                raise RuntimeError(f"Request failed: {response.status}")
+
+            result = await response.json()
+
+        if "error" in result:
+            error = result["error"]
+            raise RuntimeError(f"MCP Error: {error.get('message', error)}")
+
+        return result.get("result", {})
+
+    async def notify(self, method: str, params: Dict[str, Any]):
+        """Send a JSON-RPC notification via HTTP POST."""
+        if not self.session or not self._endpoint:
+            raise RuntimeError("Not connected")
+
+        notification = {
+            "jsonrpc": "2.0",
+            "method": method,
+            "params": params
+        }
+
+        async with self.session.post(self._endpoint, json=notification):
+            pass  # Notifications don't expect a response
+
+    async def close(self):
+        """Close the HTTP session."""
+        if self.session:
+            await self.session.close()
+            self.session = None
+
+
+async def create_client(target: str, transport_type: TransportType) -> BaseMCPTransport:
+    """
+    Factory function to create an appropriate MCP transport.
+
+    Args:
+        target: Target specification
+            - For STDIO: "python server.py" or command with args
+            - For SSE: "https://example.com/sse"
+        transport_type: Type of transport to use
+
+    Returns:
+        Connected MCP transport instance
+    """
+    if transport_type == TransportType.STDIO:
+        # Parse command and arguments
+        parts = target.split()
+        if not parts:
+            raise ValueError("Empty command for STDIO transport")
+
+        command = parts[0]
+        args = parts[1:]
+
+        transport = StdioTransport(command, args)
+        await transport.connect()
+        return transport
+
+    elif transport_type == TransportType.SSE:
+        transport = SSETransport(target)
+        await transport.connect()
+        return transport
+
+    else:
+        raise ValueError(f"Unsupported transport type: {transport_type}")
+
+
+def infer_transport_type(target: str) -> TransportType:
+    """
+    Infer transport type from target string.
+
+    Args:
+        target: Target specification
+
+    Returns:
+        Inferred TransportType
+    """
+    if target.startswith(("http://", "https://")):
+        return TransportType.SSE
+    elif target.startswith("stdio"):
+        return TransportType.STDIO
+    else:
+        # Assume it's a command for STDIO
+        return TransportType.STDIO

agent_audit/version.py

@@ -0,0 +1,3 @@
+"""Version information for agent-audit."""
+
+__version__ = "0.1.0"

agent_audit-0.1.0.dist-info/METADATA

@@ -0,0 +1,219 @@
+Metadata-Version: 2.1
+Name: agent-audit
+Version: 0.1.0
+Summary: Security scanner for AI agents and MCP configurations - Based on OWASP Agentic Top 10
+Home-page: https://github.com/HeadyZhang/agent-audit
+License: MIT
+Keywords: ai,agent,security,mcp,audit,owasp,vulnerability,scanner
+Author: Agent Security Team
+Author-email: security@example.com
+Requires-Python: >=3.9,<4.0
+Classifier: Development Status :: 4 - Beta
+Classifier: Environment :: Console
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Information Technology
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Security
+Classifier: Topic :: Software Development :: Quality Assurance
+Classifier: Topic :: Software Development :: Testing
+Requires-Dist: aiofiles (>=23.0,<24.0)
+Requires-Dist: aiohttp (>=3.9,<4.0)
+Requires-Dist: click (>=8.1.0,<9.0.0)
+Requires-Dist: pydantic (>=2.0,<3.0)
+Requires-Dist: pyyaml (>=6.0,<7.0)
+Requires-Dist: rich (>=13.0.0,<14.0.0)
+Project-URL: Bug Tracker, https://github.com/HeadyZhang/agent-audit/issues
+Project-URL: Changelog, https://github.com/HeadyZhang/agent-audit/releases
+Project-URL: Documentation, https://github.com/HeadyZhang/agent-audit#readme
+Project-URL: Repository, https://github.com/HeadyZhang/agent-audit
+Description-Content-Type: text/markdown
+
+# Agent Audit
+
+[![PyPI version](https://badge.fury.io/py/agent-audit.svg)](https://badge.fury.io/py/agent-audit)
+[![Python](https://img.shields.io/pypi/pyversions/agent-audit.svg)](https://pypi.org/project/agent-audit/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![CI](https://github.com/HeadyZhang/agent-audit/actions/workflows/ci.yml/badge.svg)](https://github.com/HeadyZhang/agent-audit/actions/workflows/ci.yml)
+[![codecov](https://codecov.io/gh/HeadyZhang/agent-audit/branch/master/graph/badge.svg)](https://codecov.io/gh/HeadyZhang/agent-audit)
+
+> 🛡️ Security scanner for AI agents and MCP configurations. Detects vulnerabilities based on the **OWASP Agentic Top 10**.
+
+<p align="center">
+  <img src="docs/demo.gif" alt="Agent Audit Demo" width="800">
+</p>
+
+## ✨ Features
+
+- **🔍 Python AST Scanning** - Detects dangerous patterns like `shell=True`, `eval()`, and tainted input flows
+- **⚙️ MCP Configuration Scanning** - Validates MCP server configurations for security issues
+- **🔐 Secret Detection** - Finds hardcoded credentials (AWS keys, API tokens, private keys)
+- **🌐 Runtime MCP Inspection** - Probes MCP servers without executing tools ("Agent Nmap")
+- **📊 Multiple Output Formats** - Terminal, JSON, SARIF (for GitHub Code Scanning), Markdown
+
+## 🚀 Quick Start
+
+### Installation
+
+```bash
+pip install agent-audit
+```
+
+### Basic Usage
+
+```bash
+# Scan current directory
+agent-audit scan .
+
+# Scan with JSON output
+agent-audit scan ./my-agent --format json
+
+# Scan with SARIF output for GitHub Code Scanning
+agent-audit scan . --format sarif --output results.sarif
+
+# Fail CI on critical findings only
+agent-audit scan . --fail-on critical
+
+# Inspect an MCP server at runtime
+agent-audit inspect stdio -- npx -y @modelcontextprotocol/server-filesystem /tmp
+```
+
+## 🔗 GitHub Action
+
+Add Agent Audit to your CI/CD pipeline with just a few lines:
+
+```yaml
+name: Security Scan
+on: [push, pull_request]
+
+jobs:
+  agent-audit:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      
+      - name: Run Agent Audit
+        uses: HeadyZhang/agent-audit@v1
+        with:
+          path: '.'
+          fail-on: 'high'
+          upload-sarif: 'true'
+```
+
+### Action Inputs
+
+| Input | Description | Default |
+|-------|-------------|---------|
+| `path` | Path to scan | `.` |
+| `format` | Output format: `terminal`, `json`, `sarif`, `markdown` | `sarif` |
+| `severity` | Minimum severity to report: `info`, `low`, `medium`, `high`, `critical` | `low` |
+| `fail-on` | Exit with error if findings at this severity | `high` |
+| `baseline` | Path to baseline file for incremental scanning | - |
+| `upload-sarif` | Upload SARIF to GitHub Security tab | `true` |
+
+## 🎯 Detected Issues
+
+| Rule ID | Title | Severity |
+|---------|-------|----------|
+| AGENT-001 | Command Injection via Unsanitized Input | 🔴 Critical |
+| AGENT-002 | Excessive Agent Permissions | 🟡 Medium |
+| AGENT-003 | Potential Data Exfiltration Chain | 🟠 High |
+| AGENT-004 | Hardcoded Credentials | 🔴 Critical |
+| AGENT-005 | Unverified MCP Server | 🟠 High |
+
+## ⚙️ Configuration
+
+Create a `.agent-audit.yaml` file to customize scanning:
+
+```yaml
+# Allowed network hosts (reduces AGENT-003 confidence)
+allowed_hosts:
+  - "*.internal.company.com"
+  - "api.openai.com"
+
+# Ignore rules
+ignore:
+  - rule_id: AGENT-003
+    paths:
+      - "auth/**"
+    reason: "Auth module legitimately communicates externally"
+
+# Scan settings
+scan:
+  exclude:
+    - "tests/**"
+    - "venv/**"
+  min_severity: low
+  fail_on: high
+```
+
+## 📈 Baseline Scanning
+
+Track new findings incrementally:
+
+```bash
+# Save current findings as baseline
+agent-audit scan . --save-baseline baseline.json
+
+# Only report new findings
+agent-audit scan . --baseline baseline.json
+```
+
+## 📖 CLI Reference
+
+```
+Usage: agent-audit [OPTIONS] COMMAND [ARGS]...
+
+Commands:
+  scan     Scan agent code and configurations
+  inspect  Inspect an MCP server at runtime
+  init     Initialize configuration file
+
+Options:
+  --version   Show version
+  -v          Enable verbose output
+  -q          Only show errors
+  --help      Show this message
+```
+
+## 🛠️ Development
+
+See [CONTRIBUTING.md](CONTRIBUTING.md) for development setup and guidelines.
+
+```bash
+# Clone the repository
+git clone https://github.com/HeadyZhang/agent-audit
+cd agent-security-suite
+
+# Install dependencies
+cd packages/core && poetry install
+cd ../audit && poetry install
+
+# Run tests
+poetry run pytest tests/ -v
+
+# Run the scanner
+poetry run agent-audit scan .
+```
+
+## 📄 License
+
+MIT License - see [LICENSE](LICENSE) for details.
+
+## 🙏 Acknowledgments
+
+- Based on the [OWASP Agentic Security Top 10](https://owasp.org/www-project-agentic-security/)
+- Inspired by the need for better AI agent security tooling
+
+---
+
+<p align="center">
+  Made with ❤️ for the AI agent security community
+</p>
+

agent_audit-0.1.0.dist-info/RECORD

@@ -0,0 +1,37 @@
+agent_audit/__init__.py,sha256=D8k1t9MieJMyq0in1RsJAOgZfGW1F5r-djHTspnXEG0,98
+agent_audit/__main__.py,sha256=juDPvuN3eWRpjRm94fxhPUyM8VBMmg8xOS59iEYdAMM,340
+agent_audit/cli/__init__.py,sha256=rsgbDj0L-wbwEfkY4bWRbr2QPUlDrtMp3z9k7qEtXAo,34
+agent_audit/cli/commands/__init__.py,sha256=YKhyx4t_78Ti9_utLsgSo_iFNqXBCxCPyD0MZtGZwRo,36
+agent_audit/cli/commands/init.py,sha256=agP2lvZoNwcVSrNJGatAb7boQZL5ma3Xcq5K3y0Scsw,1274
+agent_audit/cli/commands/inspect.py,sha256=HFxPrLCGRP1Chtccip_9fr4l49Mvx3rpIbaOHPrwlC8,7697
+agent_audit/cli/commands/scan.py,sha256=5JvH8HsGtDnNsJRwWkR5A1IH10zFg-tuGfbzHAeAd8w,11642
+agent_audit/cli/formatters/__init__.py,sha256=BspmN4xcVYvmSY5fClccAIqwMUqQP7MO9auVwJJbWIM,41
+agent_audit/cli/formatters/json.py,sha256=Wy-cKPmH62oXeHyM4rjQySgB7X4GwJX_7XEcWHK1Pcg,3997
+agent_audit/cli/formatters/sarif.py,sha256=5j4dfuEBIy5HkusWLT0Ha_vW6sJsxcNcTfm4jq-CPH8,5361
+agent_audit/cli/formatters/terminal.py,sha256=naNy-jAqWlrlO_SsL9iSZPsN-4jY1EsHcepVFIE2Fpg,6917
+agent_audit/cli/main.py,sha256=JHKm0VXsNsSWwNLFOurwaUR54PuKVLod1Mat6ypqwss,971
+agent_audit/config/__init__.py,sha256=BHW5R9IU5aMumvd5JLDaeuUu3ZH_kU7qODTKs3tjdak,48
+agent_audit/config/ignore.py,sha256=6tLxp89zR5FOst1Ci6GvVGZhghNm_Q6UyXkFYmAZu1Q,15019
+agent_audit/core_utils/__init__.py,sha256=XtUKNre_-5Ryf9u3brruaWu8kmhWklWYn1brRQg0vJg,49
+agent_audit/models/__init__.py,sha256=0KBDUfUlXOfQPNkx8VDL89DxCjjbQgw2vO-v5xTf2Xs,475
+agent_audit/models/finding.py,sha256=2M8H4gNlZGUMEUTj4svvNGJ3kGu-OagOQS49YH0K5ug,5512
+agent_audit/models/risk.py,sha256=mDpvi5WpvcppTktPgTMw2mSgnfgusn7fU8HXIjMBvN8,2216
+agent_audit/models/tool.py,sha256=hVQW4Q9ZvNDEDNtFMUFvvxwcXTiMw0ZFLb_IQLNSihE,5790
+agent_audit/rules/__init__.py,sha256=L-sssDeN5jVP8_mi6zF4JMGxq4jtvJOQZAt3jKnDnhs,192
+agent_audit/rules/engine.py,sha256=9S6TMfW-0v8XpzGW4vVscTJUo8OUf2K5p7V5qUyoyPE,18179
+agent_audit/rules/loader.py,sha256=Fdwl_0L89s97gP-6TD27S6V-jD_A5TkQMsmr3xtC-is,5157
+agent_audit/scanners/__init__.py,sha256=My758N4VyF-g2ARYPxWLIlyvZBAMm7KFQx3w6FDrZR4,136
+agent_audit/scanners/base.py,sha256=pdSu0lxS5gJtvI7o7c3MFcytzOXezytID0KDUlCyxks,687
+agent_audit/scanners/config_scanner.py,sha256=hkpq-5-Qbm7gCTjdUrcXrQGklK7D89gTmpPaMcMTl-Y,13136
+agent_audit/scanners/mcp_config_scanner.py,sha256=LAtW9F7HFgDIb1W9U4dSFz27YOB3ls_ifF8_XG-lDbU,10666
+agent_audit/scanners/mcp_inspector.py,sha256=8tfeLm_9t5c4KPkaEoRSQwB-94HN3DXbm_86Mbc88Js,15316
+agent_audit/scanners/python_scanner.py,sha256=BO6sfp1lq6WARACqQzueDFbvnSCBHSP0i5SnW7tIhso,20691
+agent_audit/scanners/secret_scanner.py,sha256=moNaesAqc3PFfWEvfu8V8OmO0bhZsGPj8dz_Ha_Cqvc,18686
+agent_audit/utils/__init__.py,sha256=3leLd2NVV013ymtXxJjUg-THvZwUmfYiCBTwJaSsDqg,402
+agent_audit/utils/compat.py,sha256=-TC2CCq1gBislYbVxdqKq7XE7c325GoPbbDF9BovOA8,3040
+agent_audit/utils/mcp_client.py,sha256=TjW-aUcwnOgnSuAvk352HxiHO0sfd8Vb3Y2ySqA3-9M,10631
+agent_audit/version.py,sha256=YeuRVUF33dlS0mzQu2Huk9qG-_by3kSZEAwoBB2VjCA,66
+agent_audit-0.1.0.dist-info/METADATA,sha256=XjaOj5QxB2MeA-zv1DYSagyLliwmsI0dDHAcUM0mkrM,6588
+agent_audit-0.1.0.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+agent_audit-0.1.0.dist-info/entry_points.txt,sha256=Fc0jPCn3KKoZttTJATiXY2QljohsIdTv-J29TRNU_Gc,56
+agent_audit-0.1.0.dist-info/RECORD,,

agent_audit-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 1.9.1
+Root-Is-Purelib: true
+Tag: py3-none-any

agent_audit-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,3 @@
+[console_scripts]
+agent-audit=agent_audit.cli.main:cli
+