agent-audit 0.1.0__py3-none-any.whl

agent_audit/models/risk.py

@@ -0,0 +1,77 @@
+"""Risk assessment models."""
+
+from enum import Enum
+from dataclasses import dataclass
+from typing import Optional
+
+
+def _normalize_path(path: str) -> str:
+    """Normalize path to use forward slashes for cross-platform consistency."""
+    return path.replace("\\", "/")
+
+
+class Severity(Enum):
+    """Severity levels for findings."""
+    CRITICAL = "critical"
+    HIGH = "high"
+    MEDIUM = "medium"
+    LOW = "low"
+    INFO = "info"
+
+    def __lt__(self, other: "Severity") -> bool:
+        order = [Severity.INFO, Severity.LOW, Severity.MEDIUM, Severity.HIGH, Severity.CRITICAL]
+        return order.index(self) < order.index(other)
+
+    def __le__(self, other: "Severity") -> bool:
+        return self == other or self < other
+
+    def __gt__(self, other: "Severity") -> bool:
+        return not self <= other
+
+    def __ge__(self, other: "Severity") -> bool:
+        return not self < other
+
+
+class Category(Enum):
+    """Categories for security findings."""
+    COMMAND_INJECTION = "command_injection"
+    DATA_EXFILTRATION = "data_exfiltration"
+    PRIVILEGE_ESCALATION = "privilege_escalation"
+    SUPPLY_CHAIN = "supply_chain"
+    CREDENTIAL_EXPOSURE = "credential_exposure"
+    PROMPT_INJECTION = "prompt_injection"
+    EXCESSIVE_PERMISSION = "excessive_permission"
+
+
+@dataclass
+class Location:
+    """Code location for a finding."""
+    file_path: str
+    start_line: int
+    end_line: int
+    start_column: Optional[int] = None
+    end_column: Optional[int] = None
+    snippet: Optional[str] = None
+
+    def __post_init__(self):
+        """Normalize file_path to use forward slashes for cross-platform consistency."""
+        self.file_path = _normalize_path(self.file_path)
+
+
+@dataclass
+class RiskScore:
+    """Risk score calculation result."""
+    score: float  # 0.0 - 10.0
+    factors: dict  # Contributing factors and their weights
+
+    def is_high_risk(self) -> bool:
+        """Check if this represents high risk."""
+        return self.score >= 7.0
+
+    def is_medium_risk(self) -> bool:
+        """Check if this represents medium risk."""
+        return 4.0 <= self.score < 7.0
+
+    def is_low_risk(self) -> bool:
+        """Check if this represents low risk."""
+        return self.score < 4.0

agent_audit/models/tool.py

@@ -0,0 +1,182 @@
+"""Tool definition models for agent analysis."""
+
+from dataclasses import dataclass, field
+from typing import List, Optional, Set
+from enum import Enum, auto
+
+
+class PermissionType(Enum):
+    """Permission types that tools may require."""
+    FILE_READ = auto()
+    FILE_WRITE = auto()
+    FILE_DELETE = auto()
+    SHELL_EXEC = auto()
+    NETWORK_OUTBOUND = auto()
+    NETWORK_INBOUND = auto()
+    DATABASE_READ = auto()
+    DATABASE_WRITE = auto()
+    SECRET_ACCESS = auto()
+    BROWSER_CONTROL = auto()
+    PROCESS_SPAWN = auto()
+
+
+class RiskLevel(Enum):
+    """Risk levels for tools."""
+    SAFE = 1
+    LOW = 2
+    MEDIUM = 3
+    HIGH = 4
+    CRITICAL = 5
+
+
+@dataclass
+class ToolParameter:
+    """Tool parameter definition."""
+    name: str
+    type: str
+    required: bool = False
+    description: Optional[str] = None
+    allows_arbitrary_input: bool = False
+    sanitization_present: bool = False
+
+
+@dataclass
+class ToolDefinition:
+    """
+    Agent tool definition.
+
+    Represents a tool that an agent can call, with its permissions,
+    risk level, and security properties.
+    """
+    name: str
+    description: str
+    source_file: str
+    source_line: int
+
+    permissions: Set[PermissionType] = field(default_factory=set)
+    risk_level: RiskLevel = RiskLevel.LOW
+    parameters: List[ToolParameter] = field(default_factory=list)
+
+    # MCP-specific fields
+    mcp_server: Optional[str] = None
+    mcp_server_verified: bool = False
+
+    # Security properties
+    has_input_validation: bool = False
+    has_output_sanitization: bool = False
+    runs_in_sandbox: bool = False
+    requires_approval: bool = False
+
+    # Capability flags (derived from permissions)
+    can_execute_code: bool = False
+    can_access_filesystem: bool = False
+    can_access_network: bool = False
+    can_access_secrets: bool = False
+
+    def calculate_risk_score(self) -> float:
+        """
+        Calculate risk score (0.0 - 10.0).
+
+        Takes into account:
+        - Permission weights (shell exec is highest risk)
+        - Mitigating factors (input validation, sandbox)
+        - MCP server verification status
+        """
+        score = 0.0
+
+        # Permission weights
+        permission_weights = {
+            PermissionType.SHELL_EXEC: 3.0,
+            PermissionType.SECRET_ACCESS: 2.5,
+            PermissionType.FILE_DELETE: 2.0,
+            PermissionType.DATABASE_WRITE: 2.0,
+            PermissionType.NETWORK_OUTBOUND: 1.5,
+            PermissionType.FILE_WRITE: 1.5,
+            PermissionType.PROCESS_SPAWN: 2.0,
+            PermissionType.BROWSER_CONTROL: 1.5,
+            PermissionType.FILE_READ: 0.5,
+            PermissionType.DATABASE_READ: 0.5,
+            PermissionType.NETWORK_INBOUND: 1.0,
+        }
+
+        for perm in self.permissions:
+            score += permission_weights.get(perm, 0.5)
+
+        # Mitigating factors
+        if self.has_input_validation:
+            score *= 0.7
+        if self.runs_in_sandbox:
+            score *= 0.5
+        if self.requires_approval:
+            score *= 0.6
+
+        # MCP verification penalty
+        if self.mcp_server and not self.mcp_server_verified:
+            score *= 1.3
+
+        return min(10.0, score)
+
+    def infer_risk_level(self) -> RiskLevel:
+        """Infer risk level from calculated score."""
+        score = self.calculate_risk_score()
+        if score >= 8.0:
+            return RiskLevel.CRITICAL
+        elif score >= 6.0:
+            return RiskLevel.HIGH
+        elif score >= 4.0:
+            return RiskLevel.MEDIUM
+        elif score >= 2.0:
+            return RiskLevel.LOW
+        else:
+            return RiskLevel.SAFE
+
+    def update_capability_flags(self):
+        """Update capability flags based on permissions."""
+        self.can_execute_code = PermissionType.SHELL_EXEC in self.permissions
+        self.can_access_filesystem = any(
+            p in self.permissions for p in [
+                PermissionType.FILE_READ,
+                PermissionType.FILE_WRITE,
+                PermissionType.FILE_DELETE
+            ]
+        )
+        self.can_access_network = any(
+            p in self.permissions for p in [
+                PermissionType.NETWORK_OUTBOUND,
+                PermissionType.NETWORK_INBOUND
+            ]
+        )
+        self.can_access_secrets = PermissionType.SECRET_ACCESS in self.permissions
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for JSON serialization."""
+        return {
+            "name": self.name,
+            "description": self.description,
+            "source_file": self.source_file,
+            "source_line": self.source_line,
+            "permissions": [p.name for p in self.permissions],
+            "risk_level": self.risk_level.name,
+            "risk_score": self.calculate_risk_score(),
+            "parameters": [
+                {
+                    "name": p.name,
+                    "type": p.type,
+                    "required": p.required,
+                    "description": p.description,
+                    "allows_arbitrary_input": p.allows_arbitrary_input,
+                    "sanitization_present": p.sanitization_present,
+                }
+                for p in self.parameters
+            ],
+            "mcp_server": self.mcp_server,
+            "mcp_server_verified": self.mcp_server_verified,
+            "has_input_validation": self.has_input_validation,
+            "has_output_sanitization": self.has_output_sanitization,
+            "runs_in_sandbox": self.runs_in_sandbox,
+            "requires_approval": self.requires_approval,
+            "can_execute_code": self.can_execute_code,
+            "can_access_filesystem": self.can_access_filesystem,
+            "can_access_network": self.can_access_network,
+            "can_access_secrets": self.can_access_secrets,
+        }

agent_audit/rules/__init__.py

@@ -0,0 +1,6 @@
+"""Rule engine and loader for Agent Security Suite."""
+
+from agent_audit.rules.engine import RuleEngine
+from agent_audit.rules.loader import RuleLoader
+
+__all__ = ["RuleEngine", "RuleLoader"]