agent-audit 0.1.0__py3-none-any.whl

agent_audit/__init__.py

@@ -0,0 +1,3 @@
+"""Agent Audit - Security scanner for AI agents and MCP configurations."""
+
+__version__ = "0.1.0"

agent_audit/__main__.py

@@ -0,0 +1,13 @@
+"""Entry point for running agent-audit as a module."""
+
+import sys
+
+# Windows event loop policy fix - must be set before any asyncio imports
+if sys.platform == "win32":
+    import asyncio
+    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+
+from agent_audit.cli.main import cli
+
+if __name__ == "__main__":
+    cli()

agent_audit/cli/__init__.py

@@ -0,0 +1 @@
+"""CLI module for agent-audit."""

agent_audit/cli/commands/__init__.py

@@ -0,0 +1 @@
+"""CLI commands for agent-audit."""

agent_audit/cli/commands/init.py

@@ -0,0 +1,44 @@
+"""Init command for creating configuration files."""
+
+import sys
+from pathlib import Path
+
+import click
+from rich.console import Console
+
+from agent_audit.config.ignore import create_default_config
+
+console = Console()
+
+
+@click.command()
+@click.option('--force', '-f', is_flag=True, help='Overwrite existing config')
+def init(force: bool):
+    """
+    Initialize agent-audit configuration.
+
+    Creates a .agent-audit.yaml file in the current directory with
+    default settings and example ignore rules.
+
+    Examples:
+
+        agent-audit init
+
+        agent-audit init --force
+    """
+    config_path = Path('.agent-audit.yaml')
+
+    if config_path.exists() and not force:
+        console.print(f"[yellow]Configuration file already exists: {config_path}[/yellow]")
+        console.print("Use --force to overwrite")
+        sys.exit(1)
+
+    config_content = create_default_config()
+    config_path.write_text(config_content, encoding="utf-8")
+
+    console.print(f"[green]Created configuration file: {config_path}[/green]")
+    console.print()
+    console.print("Edit this file to:")
+    console.print("  - Add allowed hosts for network destinations")
+    console.print("  - Configure ignore rules for false positives")
+    console.print("  - Set scan exclusion patterns")

agent_audit/cli/commands/inspect.py

@@ -0,0 +1,236 @@
+"""Inspect command for probing MCP servers."""
+
+import asyncio
+import sys
+from typing import Optional
+
+import click
+from rich.console import Console
+from rich.table import Table
+from rich.panel import Panel
+
+from agent_audit.scanners.mcp_inspector import MCPInspector, MCPInspectionResult
+from agent_audit.utils.mcp_client import TransportType
+
+console = Console()
+
+
+def render_inspection_result(result: MCPInspectionResult, output_format: str = "terminal"):
+    """Render inspection result to the console."""
+    if output_format == "json":
+        _render_json(result)
+    else:
+        _render_terminal(result)
+
+
+def _render_terminal(result: MCPInspectionResult):
+    """Render result as Rich terminal output."""
+    # Status indicator
+    if result.connected:
+        status = "[green]✓ Connected[/green]"
+        border_style = "blue" if result.risk_score < 5 else "red"
+    else:
+        status = "[red]✗ Failed[/red]"
+        border_style = "red"
+
+    # Header panel
+    header_text = (
+        f"[bold]MCP Server Inspection[/bold]\n"
+        f"Server: {result.server_name}"
+    )
+    if result.server_version:
+        header_text += f" v{result.server_version}"
+    header_text += "\n"
+    header_text += f"Status: {status}  |  Response: {result.response_time_ms:.0f}ms\n"
+    header_text += f"Risk Score: {result.risk_score:.1f}/10"
+
+    console.print(Panel.fit(header_text, border_style=border_style))
+
+    if not result.connected:
+        console.print(f"[red]Error: {result.connection_error}[/red]")
+        return
+
+    # Capabilities
+    if result.capabilities_declared:
+        caps = ", ".join(result.capabilities_declared) or "none"
+        console.print(f"\n[dim]Capabilities:[/dim] {caps}")
+
+    # Tools table
+    console.print(f"\n[bold]Tools ({result.tool_count})[/bold]")
+
+    if result.tools:
+        tool_table = Table(show_header=True, header_style="bold cyan")
+        tool_table.add_column("Tool", style="cyan")
+        tool_table.add_column("Permissions", style="yellow")
+        tool_table.add_column("Risk", justify="center")
+        tool_table.add_column("Validation")
+
+        risk_emoji = {
+            1: "🟢",  # SAFE
+            2: "🟢",  # LOW
+            3: "🟡",  # MEDIUM
+            4: "🟠",  # HIGH
+            5: "🔴",  # CRITICAL
+        }
+
+        for tool in result.tools:
+            perms = ", ".join(p.name for p in tool.permissions) or "none"
+            risk_value = tool.risk_level.value if hasattr(tool.risk_level, 'value') else 1
+            risk = risk_emoji.get(risk_value, "⚪")
+            validation = "✅" if tool.has_input_validation else "❌"
+
+            # Truncate description if needed
+            tool_name = tool.name
+            if len(tool_name) > 30:
+                tool_name = tool_name[:27] + "..."
+
+            tool_table.add_row(tool_name, perms, risk, validation)
+
+        console.print(tool_table)
+    else:
+        console.print("[dim]No tools exposed[/dim]")
+
+    # Resources
+    if result.resources:
+        console.print(f"\n[bold]Resources ({result.resource_count})[/bold]")
+        for res in result.resources[:10]:  # Limit display
+            uri = res.get('uri', 'unknown')
+            console.print(f"  📄 {uri}")
+        if result.resource_count > 10:
+            console.print(f"  [dim]... and {result.resource_count - 10} more[/dim]")
+
+    # Prompts
+    if result.prompts:
+        console.print(f"\n[bold]Prompts ({result.prompt_count})[/bold]")
+        for prompt in result.prompts[:10]:
+            name = prompt.get('name', 'unknown')
+            console.print(f"  💬 {name}")
+        if result.prompt_count > 10:
+            console.print(f"  [dim]... and {result.prompt_count - 10} more[/dim]")
+
+    # Security findings
+    if result.findings:
+        console.print(f"\n[bold red]Security Findings ({len(result.findings)})[/bold red]")
+
+        severity_colors = {
+            "critical": "red",
+            "high": "red",
+            "medium": "yellow",
+            "low": "blue"
+        }
+
+        for finding in result.findings:
+            severity = finding.get("severity", "medium")
+            color = severity_colors.get(severity, "white")
+            desc = finding.get("description", "Unknown issue")
+            tool_name = finding.get("tool", "")
+
+            if tool_name:
+                console.print(f"  [{color}]⚠ {severity.upper()}[/{color}]: {desc} (tool: {tool_name})")
+            else:
+                console.print(f"  [{color}]⚠ {severity.upper()}[/{color}]: {desc}")
+
+
+def _render_json(result: MCPInspectionResult):
+    """Render result as JSON."""
+    import json
+
+    output = {
+        "server_name": result.server_name,
+        "server_version": result.server_version,
+        "transport": result.transport.value,
+        "connected": result.connected,
+        "connection_error": result.connection_error,
+        "response_time_ms": result.response_time_ms,
+        "risk_score": result.risk_score,
+        "capabilities": result.capabilities_declared,
+        "tool_count": result.tool_count,
+        "tools": [t.to_dict() for t in result.tools],
+        "resource_count": result.resource_count,
+        "resources": result.resources,
+        "prompt_count": result.prompt_count,
+        "prompts": result.prompts,
+        "findings": result.findings,
+    }
+
+    console.print_json(json.dumps(output, indent=2))
+
+
+async def run_inspect_async(
+    target: str,
+    transport: Optional[str],
+    timeout: int,
+    output_format: str
+) -> int:
+    """Run the inspection asynchronously."""
+    inspector = MCPInspector(timeout=timeout)
+
+    # Determine transport type
+    transport_type: Optional[TransportType] = None
+    if transport:
+        transport_type = TransportType(transport)
+
+    result = await inspector.inspect(target, transport_type)
+    render_inspection_result(result, output_format)
+
+    # Return exit code based on risk
+    if not result.connected:
+        return 2  # Connection failure
+    elif result.risk_score >= 7.0:
+        return 1  # High risk
+    return 0
+
+
+def run_inspect(
+    target: str,
+    transport: Optional[str],
+    timeout: int,
+    output_format: str
+) -> int:
+    """Run the inspection."""
+    return asyncio.run(run_inspect_async(target, transport, timeout, output_format))
+
+
+@click.command()
+@click.argument('transport_type', type=click.Choice(['stdio', 'sse']), required=True)
+@click.argument('target', nargs=-1, required=True)
+@click.option('--timeout', '-t', default=30, help='Connection timeout in seconds')
+@click.option('--format', '-f', 'output_format',
+              type=click.Choice(['terminal', 'json']), default='terminal',
+              help='Output format')
+def inspect(transport_type: str, target: tuple, timeout: int, output_format: str):
+    """
+    Inspect a running MCP server and analyze its tools.
+
+    TRANSPORT_TYPE is either 'stdio' or 'sse'.
+
+    For stdio, TARGET is the command to run the server:
+
+        agent-audit inspect stdio -- python my_mcp_server.py
+
+    For sse, TARGET is the URL:
+
+        agent-audit inspect sse https://example.com/sse
+
+    The inspector connects to the server, retrieves its tool definitions,
+    and analyzes them for security risks WITHOUT executing any tools.
+    """
+    # Join target parts back together
+    target_str = ' '.join(target)
+
+    # Handle the "--" separator for stdio
+    if target_str.startswith('-- '):
+        target_str = target_str[3:]
+
+    if not target_str:
+        console.print("[red]Error: No target specified[/red]")
+        sys.exit(1)
+
+    exit_code = run_inspect(
+        target=target_str,
+        transport=transport_type,
+        timeout=timeout,
+        output_format=output_format
+    )
+
+    sys.exit(exit_code)

agent_audit/cli/commands/scan.py

@@ -0,0 +1,329 @@
+"""Scan command implementation."""
+
+from pathlib import Path
+from typing import Optional, List
+
+import click
+from rich.console import Console
+
+from agent_audit.models.finding import Finding
+from agent_audit.models.risk import Severity
+from agent_audit.rules.engine import RuleEngine
+
+from agent_audit.scanners.python_scanner import PythonScanner
+from agent_audit.scanners.mcp_config_scanner import MCPConfigScanner
+from agent_audit.scanners.secret_scanner import SecretScanner
+from agent_audit.config.ignore import (
+    IgnoreManager, load_baseline, filter_by_baseline, save_baseline
+)
+from agent_audit.cli.formatters.terminal import format_scan_results
+
+console = Console()
+
+
+def run_scan(
+    path: Path,
+    output_format: str,
+    output_path: Optional[Path],
+    min_severity: str,
+    additional_rules: List[str],
+    fail_on_severity: str,
+    baseline_path: Optional[Path] = None,
+    save_baseline_path: Optional[Path] = None,
+    verbose: bool = False,
+    quiet: bool = False
+) -> int:
+    """
+    Run the security scan.
+
+    Returns exit code: 0 for success, 1 for findings at fail_on level.
+    """
+    # Initialize ignore manager first to get exclude patterns
+    ignore_manager = IgnoreManager()
+    config_loaded = ignore_manager.load(path)
+
+    if verbose and config_loaded:
+        console.print(f"[dim]Loaded config from: {ignore_manager._loaded_from}[/dim]")
+
+    # Get exclude patterns from config
+    exclude_patterns = ignore_manager.get_exclude_patterns()
+
+    # Initialize scanners with exclude patterns
+    python_scanner = PythonScanner(exclude_patterns=exclude_patterns)
+    mcp_scanner = MCPConfigScanner()
+    secret_scanner = SecretScanner(exclude_paths=exclude_patterns)
+
+    # Initialize rule engine
+    rule_engine = RuleEngine()
+
+    # Find rules directory - check multiple possible locations
+    possible_rules_dirs = [
+        # Relative to this file (installed package)
+        Path(__file__).parent.parent.parent.parent.parent.parent.parent / "rules" / "builtin",
+        # Relative to project root (development)
+        Path(__file__).resolve().parent.parent.parent.parent.parent.parent.parent / "rules" / "builtin",
+        # Relative to current working directory
+        Path.cwd() / "rules" / "builtin",
+        # Go up from cwd if in packages/audit
+        Path.cwd().parent.parent / "rules" / "builtin",
+    ]
+
+    for rules_dir in possible_rules_dirs:
+        if rules_dir.exists():
+            rule_engine.add_builtin_rules_dir(rules_dir)
+            break
+
+    rule_engine.load_rules()
+
+    # Collect all findings
+    all_findings: List[Finding] = []
+    scanned_files = 0
+
+    # Run Python scanner
+    if not quiet:
+        console.print("[dim]Scanning Python files...[/dim]")
+
+    python_results = python_scanner.scan(path)
+    for result in python_results:
+        scanned_files += 1
+
+        # Generate findings from dangerous patterns
+        findings = rule_engine.evaluate_dangerous_patterns(
+            result.dangerous_patterns,
+            result.source_file
+        )
+        all_findings.extend(findings)
+
+        # Check for credentials in source
+        try:
+            source = Path(result.source_file).read_text(encoding='utf-8')
+            cred_findings = rule_engine.evaluate_credentials(source, result.source_file)
+            all_findings.extend(cred_findings)
+        except Exception:
+            pass
+
+        # Evaluate tool permissions
+        if result.tools:
+            perm_findings = rule_engine.evaluate_permission_scope(
+                result.tools,
+                result.source_file
+            )
+            all_findings.extend(perm_findings)
+
+    # Run MCP config scanner
+    if not quiet:
+        console.print("[dim]Scanning MCP configurations...[/dim]")
+
+    mcp_results = mcp_scanner.scan(path)
+    for result in mcp_results:
+        scanned_files += 1
+
+        # Convert server configs to dicts for rule engine
+        server_dicts = []
+        for server in result.servers:
+            server_dict = {
+                'name': server.name,
+                'url': server.url,
+                'command': server.command,
+                'args': server.args,
+                'env': server.env,
+                'verified': server.verified,
+                '_line': server._line,
+            }
+            server_dicts.append(server_dict)
+
+        mcp_findings = rule_engine.evaluate_mcp_config(server_dicts, result.source_file)
+        all_findings.extend(mcp_findings)
+
+    # Run secret scanner
+    if not quiet:
+        console.print("[dim]Scanning for secrets...[/dim]")
+
+    secret_results = secret_scanner.scan(path)
+    for result in secret_results:
+        for secret in result.secrets:
+            from agent_audit.models.risk import Location, Category
+            from agent_audit.models.finding import Remediation
+
+            finding = Finding(
+                rule_id="AGENT-004",
+                title="Hardcoded Credentials",
+                description=f"Found {secret.pattern_name}",
+                severity=Severity.CRITICAL if secret.severity == "critical" else
+                         Severity.HIGH if secret.severity == "high" else Severity.MEDIUM,
+                category=Category.CREDENTIAL_EXPOSURE,
+                location=Location(
+                    file_path=result.source_file,
+                    start_line=secret.line_number,
+                    end_line=secret.line_number,
+                    start_column=secret.start_col,
+                    end_column=secret.end_col,
+                    snippet=secret.line_content
+                ),
+                cwe_id="CWE-798",
+                remediation=Remediation(
+                    description="Use environment variables or a secrets manager"
+                )
+            )
+            all_findings.append(finding)
+
+    # Apply ignore rules
+    for finding in all_findings:
+        ignore_manager.apply_to_finding(finding)
+
+    # Filter by baseline if provided
+    if baseline_path and baseline_path.exists():
+        baseline = load_baseline(baseline_path)
+        all_findings = filter_by_baseline(all_findings, baseline)
+        if not quiet:
+            console.print(f"[dim]Filtered by baseline: {baseline_path}[/dim]")
+
+    # Filter by minimum severity
+    severity_order = {
+        'info': 0, 'low': 1, 'medium': 2, 'high': 3, 'critical': 4
+    }
+    min_sev_value = severity_order.get(min_severity.lower(), 0)
+    all_findings = [
+        f for f in all_findings
+        if severity_order.get(f.severity.value, 0) >= min_sev_value
+    ]
+
+    # Save baseline if requested
+    if save_baseline_path:
+        save_baseline(all_findings, save_baseline_path)
+        if not quiet:
+            console.print(f"[dim]Saved baseline to: {save_baseline_path}[/dim]")
+
+    # Output results
+    if output_format == "terminal":
+        format_scan_results(
+            all_findings,
+            str(path),
+            scanned_files,
+            verbose=verbose,
+            quiet=quiet
+        )
+    elif output_format == "json":
+        from agent_audit.cli.formatters.json import format_json
+        json_output = format_json(all_findings, str(path), scanned_files)
+        if output_path:
+            output_path.write_text(json_output, encoding="utf-8")
+        else:
+            console.print(json_output)
+    elif output_format == "sarif":
+        from agent_audit.cli.formatters.sarif import SARIFFormatter
+        formatter = SARIFFormatter()
+        if output_path:
+            formatter.save(all_findings, output_path)
+            if not quiet:
+                console.print(f"[dim]SARIF output saved to: {output_path}[/dim]")
+        else:
+            console.print(formatter.format_to_string(all_findings))
+    elif output_format == "markdown":
+        _output_markdown(all_findings, str(path), output_path)
+
+    # Determine exit code based on fail_on severity
+    fail_sev_value = severity_order.get(fail_on_severity.lower(), 3)  # Default: high
+    actionable_findings = [f for f in all_findings if f.is_actionable()]
+    max_severity = max(
+        (severity_order.get(f.severity.value, 0) for f in actionable_findings),
+        default=0
+    )
+
+    if max_severity >= fail_sev_value:
+        return 1
+    return 0
+
+
+def _output_markdown(findings: List[Finding], scan_path: str, output_path: Optional[Path]):
+    """Output findings as Markdown."""
+    lines = [
+        "# Agent Audit Security Report",
+        "",
+        f"**Scanned:** `{scan_path}`",
+        f"**Findings:** {len(findings)}",
+        "",
+        "## Findings",
+        "",
+    ]
+
+    for finding in findings:
+        sev = finding.severity.value.upper()
+        lines.append(f"### [{sev}] {finding.rule_id}: {finding.title}")
+        lines.append("")
+        lines.append(f"**Location:** `{finding.location.file_path}:{finding.location.start_line}`")
+        lines.append("")
+        if finding.location.snippet:
+            lines.append("```python")
+            lines.append(finding.location.snippet)
+            lines.append("```")
+            lines.append("")
+        lines.append(finding.description)
+        lines.append("")
+        if finding.remediation:
+            lines.append(f"**Fix:** {finding.remediation.description}")
+            lines.append("")
+        lines.append("---")
+        lines.append("")
+
+    md_content = "\n".join(lines)
+
+    if output_path:
+        output_path.write_text(md_content, encoding="utf-8")
+    else:
+        console.print(md_content)
+
+
+@click.command()
+@click.argument('path', type=click.Path(exists=True), default='.')
+@click.option('--format', '-f', 'output_format',
+              type=click.Choice(['terminal', 'json', 'sarif', 'markdown']),
+              default='terminal', help='Output format')
+@click.option('--output', '-o', type=click.Path(), help='Output file path')
+@click.option('--severity', '-s',
+              type=click.Choice(['critical', 'high', 'medium', 'low', 'info']),
+              default='low', help='Minimum severity to report')
+@click.option('--rules', '-r', type=click.Path(exists=True),
+              multiple=True, help='Additional rule files')
+@click.option('--fail-on',
+              type=click.Choice(['critical', 'high', 'medium', 'low']),
+              default='high', help='Exit with error if findings at this level')
+@click.option('--baseline', type=click.Path(),
+              help='Baseline file - only report new findings')
+@click.option('--save-baseline', type=click.Path(),
+              help='Save current findings as baseline')
+@click.pass_context
+def scan(ctx: click.Context, path: str, output_format: str, output: Optional[str],
+         severity: str, rules: tuple, fail_on: str, baseline: Optional[str],
+         save_baseline: Optional[str]):
+    """
+    Scan agent code and configurations for security issues.
+
+    PATH is the directory or file to scan. Defaults to current directory.
+
+    Examples:
+
+        agent-audit scan ./my-agent
+
+        agent-audit scan . --format sarif --output results.sarif
+
+        agent-audit scan . --severity critical --fail-on critical
+
+        agent-audit scan . --baseline baseline.json
+
+        agent-audit scan . --save-baseline baseline.json
+    """
+    exit_code = run_scan(
+        path=Path(path),
+        output_format=output_format,
+        output_path=Path(output) if output else None,
+        min_severity=severity,
+        additional_rules=list(rules),
+        fail_on_severity=fail_on,
+        baseline_path=Path(baseline) if baseline else None,
+        save_baseline_path=Path(save_baseline) if save_baseline else None,
+        verbose=ctx.obj.get('verbose', False),
+        quiet=ctx.obj.get('quiet', False)
+    )
+
+    ctx.exit(exit_code)

agent_audit/cli/formatters/__init__.py

@@ -0,0 +1 @@
+"""Output formatters for agent-audit."""