securitybench 0.1.0__py3-none-any.whl

sb/__init__.py

@@ -0,0 +1,38 @@
+"""Security Bench - Security testing for LLM pipelines.
+
+A comprehensive CLI tool for testing LLM applications against security
+vulnerabilities including prompt injection, jailbreaks, and data leakage.
+"""
+from .config import PipelineConfig, EndpointConfig, InputConfig, OutputConfig, ConfigError
+from .bench import SecurityBench, SecurityTest, ScanResults, TestResult
+from .runner import TestRunner, RunResult
+from .loader import TestLoader, LoaderError
+from .evaluators import PatternEvaluator, EvaluationResult
+from .cli import main
+
+__version__ = "0.1.0"
+
+__all__ = [
+    # Config
+    "PipelineConfig",
+    "EndpointConfig",
+    "InputConfig",
+    "OutputConfig",
+    "ConfigError",
+    # Bench
+    "SecurityBench",
+    "SecurityTest",
+    "ScanResults",
+    "TestResult",
+    # Runner
+    "TestRunner",
+    "RunResult",
+    # Loader
+    "TestLoader",
+    "LoaderError",
+    # Evaluator
+    "PatternEvaluator",
+    "EvaluationResult",
+    # CLI
+    "main",
+]

sb/auditor.py

@@ -0,0 +1,355 @@
+"""Security Auditor for local code, config, and infrastructure checks.
+
+Lynis-inspired security scanner that fetches checks from the API
+and runs them against local files and system.
+"""
+import asyncio
+import fnmatch
+import os
+import re
+import subprocess
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+from typing import Optional
+
+import httpx
+
+from .loader import TestLoader
+
+
+@dataclass
+class SecurityCheck:
+    """A security check definition from the API."""
+    id: str
+    command: str  # "code", "config", or "infra"
+    category: str
+    name: str
+    description: str
+    severity: str  # "critical", "high", "medium", "low"
+    detection_type: str  # "regex" or "command"
+    pattern: str  # Regex pattern or shell command
+    file_patterns: list[str] = field(default_factory=list)
+    cwe: Optional[str] = None
+    owasp_llm: Optional[str] = None
+    weight: int = 5
+    remediation: Optional[str] = None
+
+
+@dataclass
+class Finding:
+    """A security finding from running a check."""
+    check_id: str
+    check_name: str
+    severity: str
+    category: str
+    description: str
+    file_path: Optional[str] = None
+    line_number: Optional[int] = None
+    matched_content: Optional[str] = None
+    remediation: Optional[str] = None
+    cwe: Optional[str] = None
+
+
+@dataclass
+class AuditResults:
+    """Results from running an audit."""
+    scan_path: str
+    started_at: str
+    completed_at: str
+    checks_run: int
+    checks_passed: int
+    checks_failed: int
+    findings: list[Finding] = field(default_factory=list)
+
+    @property
+    def hardening_score(self) -> int:
+        """Calculate hardening score (0-100)."""
+        if self.checks_run == 0:
+            return 100
+        return int((self.checks_passed / self.checks_run) * 100)
+
+    @property
+    def grade(self) -> str:
+        """Letter grade based on hardening score."""
+        score = self.hardening_score
+        if score >= 90:
+            return "A"
+        elif score >= 80:
+            return "B"
+        elif score >= 70:
+            return "C"
+        elif score >= 60:
+            return "D"
+        else:
+            return "F"
+
+    def findings_by_severity(self) -> dict[str, list[Finding]]:
+        """Group findings by severity."""
+        result = {"critical": [], "high": [], "medium": [], "low": []}
+        for f in self.findings:
+            if f.severity in result:
+                result[f.severity].append(f)
+        return result
+
+
+class CheckLoader:
+    """Loads security checks from the API."""
+
+    API_BASE = TestLoader.API_BASE
+
+    async def load_checks(
+        self,
+        command: Optional[str] = None,
+        category: Optional[str] = None,
+        limit: int = 500,
+    ) -> list[SecurityCheck]:
+        """Load checks from the Security Bench API.
+
+        Args:
+            command: Filter by command type (code, config, infra).
+            category: Filter by category.
+            limit: Maximum number of checks to load.
+
+        Returns:
+            List of SecurityCheck objects.
+        """
+        async with httpx.AsyncClient(timeout=30.0) as client:
+            params = {"limit": limit}
+            if command:
+                params["command"] = command
+            if category:
+                params["category"] = category
+
+            response = await client.get(
+                f"{self.API_BASE}/api/checks",
+                params=params,
+            )
+            response.raise_for_status()
+
+            data = response.json()
+            checks = []
+
+            for check_data in data.get("checks", data if isinstance(data, list) else []):
+                check = SecurityCheck(
+                    id=check_data["id"],
+                    command=check_data.get("command", "code"),
+                    category=check_data.get("category", "misc"),
+                    name=check_data.get("name", check_data["id"]),
+                    description=check_data.get("description", ""),
+                    severity=check_data.get("severity", "medium"),
+                    detection_type=check_data.get("detection_type", "regex"),
+                    pattern=check_data.get("pattern", ""),
+                    file_patterns=check_data.get("file_patterns", []),
+                    cwe=check_data.get("cwe"),
+                    owasp_llm=check_data.get("owasp_llm"),
+                    weight=check_data.get("weight", 5),
+                    remediation=check_data.get("remediation"),
+                )
+                checks.append(check)
+
+            return checks
+
+
+class Auditor:
+    """Security auditor that runs checks against local files and system."""
+
+    def __init__(
+        self,
+        scan_path: Path = Path("."),
+        exclude_patterns: Optional[list[str]] = None,
+    ):
+        """Initialize the auditor.
+
+        Args:
+            scan_path: Directory to scan.
+            exclude_patterns: Glob patterns to exclude.
+        """
+        self.scan_path = scan_path.resolve()
+        self.exclude_patterns = exclude_patterns or [
+            "node_modules/*",
+            ".git/*",
+            "__pycache__/*",
+            "*.pyc",
+            ".venv/*",
+            "venv/*",
+            "dist/*",
+            "build/*",
+        ]
+        self.loader = CheckLoader()
+        self._file_cache: dict[str, list[Path]] = {}
+
+    def _should_exclude(self, path: Path) -> bool:
+        """Check if a path should be excluded."""
+        rel_path = str(path.relative_to(self.scan_path))
+        for pattern in self.exclude_patterns:
+            if fnmatch.fnmatch(rel_path, pattern):
+                return True
+        return False
+
+    def _find_files(self, patterns: list[str]) -> list[Path]:
+        """Find files matching glob patterns."""
+        cache_key = tuple(sorted(patterns))
+        if cache_key in self._file_cache:
+            return self._file_cache[cache_key]
+
+        files = set()
+        for pattern in patterns:
+            # Clean the pattern - remove leading wildcards and slashes
+            clean_pattern = pattern.lstrip("*./")
+            if not clean_pattern:
+                continue
+
+            try:
+                for path in self.scan_path.rglob(clean_pattern):
+                    if path.is_file() and not self._should_exclude(path):
+                        files.add(path)
+            except (NotImplementedError, ValueError):
+                # Fall back to fnmatch for complex patterns
+                for path in self.scan_path.rglob("*"):
+                    if path.is_file() and fnmatch.fnmatch(path.name, pattern):
+                        if not self._should_exclude(path):
+                            files.add(path)
+
+        result = list(files)
+        self._file_cache[cache_key] = result
+        return result
+
+    def _run_regex_check(self, check: SecurityCheck) -> list[Finding]:
+        """Run a regex-based check against matching files."""
+        findings = []
+
+        if not check.file_patterns:
+            return findings
+
+        files = self._find_files(check.file_patterns)
+
+        try:
+            pattern = re.compile(check.pattern, re.IGNORECASE | re.MULTILINE)
+        except re.error:
+            return findings
+
+        for file_path in files:
+            try:
+                content = file_path.read_text(errors="ignore")
+                for line_num, line in enumerate(content.splitlines(), 1):
+                    match = pattern.search(line)
+                    if match:
+                        findings.append(Finding(
+                            check_id=check.id,
+                            check_name=check.name,
+                            severity=check.severity,
+                            category=check.category,
+                            description=check.description,
+                            file_path=str(file_path.relative_to(self.scan_path)),
+                            line_number=line_num,
+                            matched_content=line.strip()[:200],
+                            remediation=check.remediation,
+                            cwe=check.cwe,
+                        ))
+            except Exception:
+                continue
+
+        return findings
+
+    def _run_command_check(self, check: SecurityCheck) -> list[Finding]:
+        """Run a command-based check."""
+        findings = []
+
+        try:
+            result = subprocess.run(
+                check.pattern,
+                shell=True,
+                capture_output=True,
+                text=True,
+                timeout=30,
+                cwd=str(self.scan_path),
+            )
+
+            # Command checks typically return non-zero or specific output on failure
+            if result.returncode != 0 or result.stdout.strip():
+                findings.append(Finding(
+                    check_id=check.id,
+                    check_name=check.name,
+                    severity=check.severity,
+                    category=check.category,
+                    description=check.description,
+                    matched_content=result.stdout.strip()[:500] if result.stdout else None,
+                    remediation=check.remediation,
+                    cwe=check.cwe,
+                ))
+        except subprocess.TimeoutExpired:
+            pass
+        except Exception:
+            pass
+
+        return findings
+
+    async def run_check(self, check: SecurityCheck) -> tuple[bool, list[Finding]]:
+        """Run a single security check.
+
+        Args:
+            check: The check to run.
+
+        Returns:
+            Tuple of (passed, findings).
+        """
+        if check.detection_type == "regex":
+            findings = self._run_regex_check(check)
+        elif check.detection_type == "command":
+            findings = self._run_command_check(check)
+        else:
+            findings = []
+
+        return len(findings) == 0, findings
+
+    async def audit(
+        self,
+        command: Optional[str] = None,
+        categories: Optional[list[str]] = None,
+        progress_callback=None,
+    ) -> AuditResults:
+        """Run a full audit.
+
+        Args:
+            command: Filter checks by command type (code, config, infra).
+            categories: Filter checks by category.
+            progress_callback: Called with (current, total, check) for progress updates.
+
+        Returns:
+            AuditResults with all findings.
+        """
+        started_at = datetime.now().isoformat()
+
+        # Load checks
+        checks = await self.loader.load_checks(command=command)
+
+        # Filter by categories if specified
+        if categories:
+            checks = [c for c in checks if c.category in categories]
+
+        all_findings = []
+        passed = 0
+        failed = 0
+
+        for i, check in enumerate(checks):
+            if progress_callback:
+                progress_callback(i + 1, len(checks), check)
+
+            check_passed, findings = await self.run_check(check)
+
+            if check_passed:
+                passed += 1
+            else:
+                failed += 1
+                all_findings.extend(findings)
+
+        return AuditResults(
+            scan_path=str(self.scan_path),
+            started_at=started_at,
+            completed_at=datetime.now().isoformat(),
+            checks_run=len(checks),
+            checks_passed=passed,
+            checks_failed=failed,
+            findings=all_findings,
+        )