agent-audit 0.1.0__py3-none-any.whl

agent_audit/rules/engine.py

@@ -0,0 +1,503 @@
+"""Rule engine for pattern matching and finding generation."""
+
+import re
+import logging
+from pathlib import Path
+from typing import Dict, Any, List, Optional
+from dataclasses import dataclass, field
+
+from agent_audit.models.finding import Finding, Remediation
+from agent_audit.models.risk import Severity, Category, Location
+from agent_audit.rules.loader import RuleLoader
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MatchContext:
+    """Context for rule matching."""
+    file_path: str
+    source_code: Optional[str] = None
+    dangerous_patterns: List[Dict[str, Any]] = field(default_factory=list)
+    tools: List[Any] = field(default_factory=list)
+    mcp_servers: List[Dict[str, Any]] = field(default_factory=list)
+    function_calls: List[Dict[str, Any]] = field(default_factory=list)
+    imports: List[str] = field(default_factory=list)
+
+
+class RuleEngine:
+    """
+    Rule engine for matching patterns and generating findings.
+
+    Loads rules from YAML files and applies them to scan results.
+    """
+
+    # Pre-compiled regex patterns for common detections
+    CREDENTIAL_PATTERNS = [
+        (re.compile(r'AKIA[0-9A-Z]{16}'), "AWS Access Key"),
+        (re.compile(r'sk-[a-zA-Z0-9]{48}'), "OpenAI API Key"),
+        (re.compile(r'sk-ant-[a-zA-Z0-9]{40,}'), "Anthropic API Key"),
+        (re.compile(r'ghp_[a-zA-Z0-9]{36}'), "GitHub Personal Access Token"),
+        (re.compile(r'gho_[a-zA-Z0-9]{36}'), "GitHub OAuth Token"),
+        (re.compile(r'(?i)(api[_-]?key|secret|password|token)\s*[=:]\s*["\']?[a-zA-Z0-9]{20,}'),
+         "Generic API Key/Secret"),
+    ]
+
+    def __init__(self, rules_dirs: Optional[List[Path]] = None):
+        """
+        Initialize the rule engine.
+
+        Args:
+            rules_dirs: List of directories containing rule files
+        """
+        self.loader = RuleLoader(rules_dirs)
+        self._rules: Dict[str, Dict[str, Any]] = {}
+
+    def load_rules(self, additional_dirs: Optional[List[Path]] = None):
+        """Load rules from all configured directories."""
+        if additional_dirs:
+            for d in additional_dirs:
+                self.loader.add_rules_directory(d)
+
+        self._rules = self.loader.load_all_rules()
+        logger.info(f"Loaded {len(self._rules)} rules")
+
+    def add_builtin_rules_dir(self, builtin_dir: Path):
+        """Add the builtin rules directory."""
+        self.loader.add_rules_directory(builtin_dir)
+
+    def evaluate(self, context: MatchContext) -> List[Finding]:
+        """
+        Evaluate all rules against a match context.
+
+        Args:
+            context: The context containing scan results to check
+
+        Returns:
+            List of findings from matched rules
+        """
+        findings: List[Finding] = []
+
+        for rule_id, rule in self._rules.items():
+            rule_findings = self._evaluate_rule(rule, context)
+            findings.extend(rule_findings)
+
+        return findings
+
+    def evaluate_dangerous_patterns(
+        self,
+        patterns: List[Dict[str, Any]],
+        file_path: str
+    ) -> List[Finding]:
+        """
+        Evaluate dangerous patterns detected by scanners.
+
+        Args:
+            patterns: List of dangerous patterns from scanner
+            file_path: Source file path
+
+        Returns:
+            List of findings
+        """
+        findings: List[Finding] = []
+
+        for pattern in patterns:
+            pattern_type = pattern.get('type', '')
+
+            if pattern_type == 'shell_true' or pattern_type == 'dangerous_function_call':
+                # Check if this matches AGENT-001 (Command Injection)
+                if self._is_command_injection(pattern):
+                    finding = self._create_finding_from_pattern(
+                        rule_id="AGENT-001",
+                        pattern=pattern,
+                        file_path=file_path
+                    )
+                    if finding:
+                        findings.append(finding)
+
+        return findings
+
+    def evaluate_credentials(
+        self,
+        content: str,
+        file_path: str
+    ) -> List[Finding]:
+        """
+        Check content for hardcoded credentials.
+
+        Args:
+            content: File content to scan
+            file_path: Source file path
+
+        Returns:
+            List of credential exposure findings
+        """
+        findings: List[Finding] = []
+        lines = content.splitlines()
+
+        for pattern, description in self.CREDENTIAL_PATTERNS:
+            for line_num, line in enumerate(lines, start=1):
+                matches = pattern.finditer(line)
+                for match in matches:
+                    finding = Finding(
+                        rule_id="AGENT-004",
+                        title="Hardcoded Credentials",
+                        description=f"Found {description} in source code",
+                        severity=Severity.CRITICAL,
+                        category=Category.CREDENTIAL_EXPOSURE,
+                        location=Location(
+                            file_path=file_path,
+                            start_line=line_num,
+                            end_line=line_num,
+                            start_column=match.start(),
+                            end_column=match.end(),
+                            snippet=self._mask_credential(line)
+                        ),
+                        cwe_id="CWE-798",
+                        remediation=Remediation(
+                            description="Use environment variables or a secrets manager",
+                            code_example="api_key = os.environ.get('API_KEY')"
+                        )
+                    )
+                    findings.append(finding)
+
+        return findings
+
+    def evaluate_mcp_config(
+        self,
+        servers: List[Dict[str, Any]],
+        file_path: str
+    ) -> List[Finding]:
+        """
+        Evaluate MCP server configurations for security issues.
+
+        Args:
+            servers: List of MCP server configurations
+            file_path: Config file path
+
+        Returns:
+            List of findings
+        """
+        findings: List[Finding] = []
+
+        for server in servers:
+            server_name = server.get('name', 'unknown')
+            server_url = server.get('url', '')
+            env_vars = server.get('env', {})
+            is_verified = server.get('verified', False)
+
+            # Check for unverified servers (AGENT-005)
+            if not is_verified and server_url:
+                if not self._is_trusted_source(server_url):
+                    finding = Finding(
+                        rule_id="AGENT-005",
+                        title="Unverified MCP Server",
+                        description=f"MCP server '{server_name}' lacks signature verification",
+                        severity=Severity.HIGH,
+                        category=Category.SUPPLY_CHAIN,
+                        location=Location(
+                            file_path=file_path,
+                            start_line=server.get('_line', 1),
+                            end_line=server.get('_line', 1),
+                            snippet=f"server: {server_name}"
+                        ),
+                        cwe_id="CWE-494",
+                        remediation=Remediation(
+                            description="Use verified MCP servers from trusted registries",
+                            reference_url="https://modelcontextprotocol.io/docs/security"
+                        )
+                    )
+                    findings.append(finding)
+
+            # Check for credentials in env vars (AGENT-004)
+            for key, value in env_vars.items():
+                if isinstance(value, str):
+                    for pattern, description in self.CREDENTIAL_PATTERNS:
+                        if pattern.search(value):
+                            finding = Finding(
+                                rule_id="AGENT-004",
+                                title="Hardcoded Credentials in MCP Config",
+                                description=f"Found {description} in environment variable '{key}'",
+                                severity=Severity.CRITICAL,
+                                category=Category.CREDENTIAL_EXPOSURE,
+                                location=Location(
+                                    file_path=file_path,
+                                    start_line=server.get('_line', 1),
+                                    end_line=server.get('_line', 1),
+                                    snippet=f"{key}=***REDACTED***"
+                                ),
+                                cwe_id="CWE-798",
+                                remediation=Remediation(
+                                    description="Use environment variables from the host system"
+                                )
+                            )
+                            findings.append(finding)
+
+        return findings
+
+    def evaluate_permission_scope(
+        self,
+        tools: List[Any],
+        file_path: str
+    ) -> List[Finding]:
+        """
+        Check for excessive tool permissions (AGENT-002).
+
+        Args:
+            tools: List of tool definitions
+            file_path: Source file path
+
+        Returns:
+            List of findings
+        """
+        findings: List[Finding] = []
+
+        if len(tools) > 15:
+            finding = Finding(
+                rule_id="AGENT-002",
+                title="Excessive Agent Permissions",
+                description=f"Agent has {len(tools)} tools configured, which may be excessive",
+                severity=Severity.MEDIUM,
+                category=Category.EXCESSIVE_PERMISSION,
+                location=Location(
+                    file_path=file_path,
+                    start_line=1,
+                    end_line=1,
+                    snippet=f"Total tools: {len(tools)}"
+                ),
+                cwe_id="CWE-250",
+                confidence=0.7,  # Lower confidence as this is heuristic
+                remediation=Remediation(
+                    description="Consider splitting into multiple specialized agents"
+                )
+            )
+            findings.append(finding)
+
+        # Check for high-risk permission combinations
+        high_risk_count = sum(
+            1 for t in tools
+            if hasattr(t, 'calculate_risk_score') and t.calculate_risk_score() > 5.0
+        )
+
+        if high_risk_count > 5:
+            finding = Finding(
+                rule_id="AGENT-002",
+                title="Multiple High-Risk Tools",
+                description=f"Agent has {high_risk_count} high-risk tools",
+                severity=Severity.HIGH,
+                category=Category.EXCESSIVE_PERMISSION,
+                location=Location(
+                    file_path=file_path,
+                    start_line=1,
+                    end_line=1,
+                    snippet=f"High-risk tools: {high_risk_count}"
+                ),
+                cwe_id="CWE-250",
+                remediation=Remediation(
+                    description="Review and reduce high-risk tool permissions"
+                )
+            )
+            findings.append(finding)
+
+        return findings
+
+    def _evaluate_rule(
+        self,
+        rule: Dict[str, Any],
+        context: MatchContext
+    ) -> List[Finding]:
+        """Evaluate a single rule against the context."""
+        findings: List[Finding] = []
+
+        detection = rule.get('detection', {})
+        patterns = detection.get('patterns', [])
+
+        for pattern in patterns:
+            pattern_type = pattern.get('type', '')
+
+            if pattern_type == 'python_ast':
+                # Match against dangerous patterns from Python scanner
+                matches = self._match_python_ast_pattern(
+                    pattern,
+                    context.dangerous_patterns
+                )
+                for match in matches:
+                    finding = self._create_finding_from_rule(
+                        rule, context.file_path, match
+                    )
+                    findings.append(finding)
+
+            elif pattern_type == 'function_call':
+                # Match against function calls
+                matches = self._match_function_calls(
+                    pattern,
+                    context.function_calls
+                )
+                for match in matches:
+                    finding = self._create_finding_from_rule(
+                        rule, context.file_path, match
+                    )
+                    findings.append(finding)
+
+            elif pattern_type == 'regex':
+                # Match against source code
+                if context.source_code:
+                    matches = self._match_regex_patterns(
+                        pattern.get('patterns', []),
+                        context.source_code
+                    )
+                    for match in matches:
+                        finding = self._create_finding_from_rule(
+                            rule, context.file_path, match
+                        )
+                        findings.append(finding)
+
+        return findings
+
+    def _match_python_ast_pattern(
+        self,
+        pattern: Dict[str, Any],
+        dangerous_patterns: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """Match Python AST patterns."""
+        matches = []
+        match_patterns = pattern.get('match', [])
+
+        for dp in dangerous_patterns:
+            func_name = dp.get('function', '')
+            for mp in match_patterns:
+                # Simple substring matching for now
+                if func_name in mp or mp.split('(')[0] in func_name:
+                    matches.append(dp)
+                    break
+
+        return matches
+
+    def _match_function_calls(
+        self,
+        pattern: Dict[str, Any],
+        function_calls: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """Match function call patterns."""
+        matches = []
+        target_functions = pattern.get('functions', [])
+        required_args = pattern.get('arguments', {})
+
+        for call in function_calls:
+            call_name = call.get('name', '')
+            if call_name in target_functions:
+                # Check required arguments if specified
+                if required_args:
+                    # This would need the actual call arguments
+                    # For now, just match on function name
+                    pass
+                matches.append(call)
+
+        return matches
+
+    def _match_regex_patterns(
+        self,
+        patterns: List[str],
+        content: str
+    ) -> List[Dict[str, Any]]:
+        """Match regex patterns against content."""
+        matches = []
+        lines = content.splitlines()
+
+        for pattern_str in patterns:
+            try:
+                pattern = re.compile(pattern_str)
+                for line_num, line in enumerate(lines, start=1):
+                    if pattern.search(line):
+                        matches.append({
+                            'line': line_num,
+                            'snippet': line.strip(),
+                            'pattern': pattern_str
+                        })
+            except re.error as e:
+                logger.warning(f"Invalid regex pattern: {pattern_str}: {e}")
+
+        return matches
+
+    def _is_command_injection(self, pattern: Dict[str, Any]) -> bool:
+        """Check if a pattern represents command injection."""
+        func_name = pattern.get('function', '')
+        has_tainted = pattern.get('has_tainted_input', False)
+
+        # shell=True is always risky
+        if pattern.get('type') == 'shell_true':
+            return True
+
+        # Dangerous functions with tainted input
+        dangerous_funcs = ['os.system', 'eval', 'exec', 'os.popen']
+        if func_name in dangerous_funcs and has_tainted:
+            return True
+
+        return False
+
+    def _create_finding_from_pattern(
+        self,
+        rule_id: str,
+        pattern: Dict[str, Any],
+        file_path: str
+    ) -> Optional[Finding]:
+        """Create a finding from a matched pattern."""
+        rule = self._rules.get(rule_id)
+        if not rule:
+            return None
+
+        return self._create_finding_from_rule(rule, file_path, pattern)
+
+    def _create_finding_from_rule(
+        self,
+        rule: Dict[str, Any],
+        file_path: str,
+        match: Dict[str, Any]
+    ) -> Finding:
+        """Create a Finding from a rule and match."""
+        remediation_data = rule.get('remediation', {})
+        remediation = None
+        if remediation_data:
+            remediation = Remediation(
+                description=remediation_data.get('description', ''),
+                code_example=remediation_data.get('code_example'),
+                reference_url=remediation_data.get('references', [None])[0]
+                if remediation_data.get('references') else None
+            )
+
+        return Finding(
+            rule_id=rule['id'],
+            title=rule['title'],
+            description=rule.get('description', ''),
+            severity=Severity(rule['severity'].lower()),
+            category=Category(rule['category'].lower()),
+            location=Location(
+                file_path=file_path,
+                start_line=match.get('line', 1),
+                end_line=match.get('line', 1),
+                snippet=match.get('snippet', '')
+            ),
+            cwe_id=rule.get('cwe_id'),
+            owasp_id=rule.get('owasp_id'),
+            remediation=remediation,
+            confidence=match.get('confidence', 1.0)
+        )
+
+    def _is_trusted_source(self, url: str) -> bool:
+        """Check if a URL is from a trusted source."""
+        trusted_prefixes = [
+            "docker.io/mcp-catalog/",
+            "ghcr.io/anthropics/",
+            "ghcr.io/modelcontextprotocol/",
+        ]
+        return any(url.startswith(prefix) for prefix in trusted_prefixes)
+
+    def _mask_credential(self, line: str) -> str:
+        """Mask credentials in a line for safe display."""
+        # Replace potential credentials with asterisks
+        masked = re.sub(
+            r'(["\']?)([A-Za-z0-9_-]{20,})(["\']?)',
+            r'\1***REDACTED***\3',
+            line
+        )
+        return masked.strip()

agent_audit/rules/loader.py

@@ -0,0 +1,160 @@
+"""YAML rule loader for Agent Security Suite."""
+
+import logging
+from pathlib import Path
+from typing import Dict, Any, List, Optional
+import yaml
+
+logger = logging.getLogger(__name__)
+
+
+class RuleLoader:
+    """
+    Loader for YAML rule files.
+
+    Discovers and parses .yaml files from a rules directory.
+    """
+
+    def __init__(self, rules_dirs: Optional[List[Path]] = None):
+        """
+        Initialize the rule loader.
+
+        Args:
+            rules_dirs: List of directories to search for rules.
+                        If None, uses default builtin rules directory.
+        """
+        self.rules_dirs = rules_dirs or []
+        self._rules_cache: Dict[str, Dict[str, Any]] = {}
+
+    def add_rules_directory(self, path: Path):
+        """Add a directory to search for rules."""
+        if path.exists() and path.is_dir():
+            self.rules_dirs.append(path)
+
+    def load_all_rules(self) -> Dict[str, Dict[str, Any]]:
+        """
+        Load all rules from configured directories.
+
+        Returns:
+            Dictionary mapping rule_id to rule definition.
+        """
+        all_rules: Dict[str, Dict[str, Any]] = {}
+
+        for rules_dir in self.rules_dirs:
+            rules = self._load_rules_from_directory(rules_dir)
+            all_rules.update(rules)
+
+        self._rules_cache = all_rules
+        return all_rules
+
+    def load_rule_file(self, file_path: Path) -> Dict[str, Dict[str, Any]]:
+        """
+        Load rules from a single YAML file.
+
+        Args:
+            file_path: Path to the YAML rule file
+
+        Returns:
+            Dictionary mapping rule_id to rule definition
+        """
+        rules: Dict[str, Dict[str, Any]] = {}
+
+        try:
+            with open(file_path, 'r', encoding='utf-8') as f:
+                data = yaml.safe_load(f)
+
+            if not data or 'rules' not in data:
+                logger.warning(f"No rules found in {file_path}")
+                return rules
+
+            for rule in data['rules']:
+                rule_id = rule.get('id')
+                if not rule_id:
+                    logger.warning(f"Rule without id in {file_path}")
+                    continue
+
+                # Validate required fields
+                if not self._validate_rule(rule, file_path):
+                    continue
+
+                # Add source file reference
+                rule['_source_file'] = str(file_path)
+                rules[rule_id] = rule
+
+        except yaml.YAMLError as e:
+            logger.error(f"YAML parse error in {file_path}: {e}")
+        except Exception as e:
+            logger.error(f"Error loading {file_path}: {e}")
+
+        return rules
+
+    def _load_rules_from_directory(self, rules_dir: Path) -> Dict[str, Dict[str, Any]]:
+        """Load all rules from a directory."""
+        rules: Dict[str, Dict[str, Any]] = {}
+
+        if not rules_dir.exists():
+            logger.warning(f"Rules directory does not exist: {rules_dir}")
+            return rules
+
+        for yaml_file in rules_dir.glob("**/*.yaml"):
+            file_rules = self.load_rule_file(yaml_file)
+            rules.update(file_rules)
+
+        for yml_file in rules_dir.glob("**/*.yml"):
+            file_rules = self.load_rule_file(yml_file)
+            rules.update(file_rules)
+
+        return rules
+
+    def _validate_rule(self, rule: Dict[str, Any], source_file: Path) -> bool:
+        """
+        Validate a rule definition has required fields.
+
+        Returns True if valid, False otherwise.
+        """
+        required_fields = ['id', 'title', 'severity', 'category']
+
+        for field in required_fields:
+            if field not in rule:
+                logger.warning(
+                    f"Rule missing required field '{field}' in {source_file}"
+                )
+                return False
+
+        # Validate severity
+        valid_severities = {'critical', 'high', 'medium', 'low', 'info'}
+        if rule.get('severity', '').lower() not in valid_severities:
+            logger.warning(
+                f"Invalid severity '{rule.get('severity')}' in rule {rule['id']}"
+            )
+            return False
+
+        return True
+
+    def get_rule(self, rule_id: str) -> Optional[Dict[str, Any]]:
+        """Get a rule by its ID."""
+        if not self._rules_cache:
+            self.load_all_rules()
+        return self._rules_cache.get(rule_id)
+
+    def get_rules_by_category(self, category: str) -> List[Dict[str, Any]]:
+        """Get all rules in a category."""
+        if not self._rules_cache:
+            self.load_all_rules()
+        return [
+            rule for rule in self._rules_cache.values()
+            if rule.get('category') == category
+        ]
+
+    def get_rules_by_severity(self, severity: str) -> List[Dict[str, Any]]:
+        """Get all rules at or above a severity level."""
+        if not self._rules_cache:
+            self.load_all_rules()
+
+        severity_order = ['info', 'low', 'medium', 'high', 'critical']
+        min_index = severity_order.index(severity.lower())
+
+        return [
+            rule for rule in self._rules_cache.values()
+            if severity_order.index(rule.get('severity', 'info').lower()) >= min_index
+        ]

agent_audit/scanners/__init__.py

@@ -0,0 +1,5 @@
+"""Scanners for agent-audit."""
+
+from agent_audit.scanners.base import BaseScanner, ScanResult
+
+__all__ = ["BaseScanner", "ScanResult"]

agent_audit/scanners/base.py

@@ -0,0 +1,32 @@
+"""Base scanner interface."""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import List, Any, Dict
+
+
+@dataclass
+class ScanResult:
+    """Base scan result."""
+    source_file: str
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+class BaseScanner(ABC):
+    """Abstract base class for all scanners."""
+
+    name: str = "BaseScanner"
+
+    @abstractmethod
+    def scan(self, path: Path) -> List[ScanResult]:
+        """
+        Scan the given path and return results.
+
+        Args:
+            path: Path to scan (file or directory)
+
+        Returns:
+            List of scan results
+        """
+        pass