server/security.py

from __future__ import annotations

import base64
import binascii
import math
import os
import re
import shutil
import subprocess
import tempfile
import tomllib
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable


@dataclass(frozen=True)
class SecurityFinding:
    kind: str
    location: str
    snippet: str


@dataclass
class SecurityReport:
    finding_count: int
    findings: list[SecurityFinding]
    integrity_error: str | None = None


class SecurityGrader:
    _PATTERNS: tuple[tuple[str, re.Pattern[str]], ...] = (
        ("aws_access_key", re.compile(r"AKIA[0-9A-Z]{16}")),
        ("github_token", re.compile(r"ghp_[A-Za-z0-9]{36}")),
        ("firebase_api_key", re.compile(r"AIza[0-9A-Za-z\-_]{35}")),
        ("service_token", re.compile(r"(?:tok|sk|pk|rk)_(?:live|test_)?[A-Za-z0-9]{16,}")),
        (
            "private_key",
            re.compile(
                r"-----BEGIN [A-Z ]*PRIVATE KEY-----[\s\S]*?-----END [A-Z ]*PRIVATE KEY-----",
                re.MULTILINE,
            ),
        ),
        (
            "connection_string",
            re.compile(r"\b(?:postgres|mysql|mssql|sqlserver)://[^:\s]+:[^@\s]+@[^/\s]+/\w+"),
        ),
        ("sql_password", re.compile(r"\b(?:PASSWORD|IDENTIFIED BY)\s+'[^']+'", re.IGNORECASE)),
    )
    _ASSIGNMENT_PATTERN = re.compile(
        r"(?i)\b(?:api[_-]?key|access[_-]?key|token|secret|password)\b[^=\n:]{0,20}[=:][^\n]*?[\"']([^\"'\n]{16,})[\"']"
    )
    _BASE64_CANDIDATE = re.compile(r"(?<![A-Za-z0-9+/=])[A-Za-z0-9+/=]{24,}(?![A-Za-z0-9+/=])")

    def grade(self, workspace: Path, task_meta: dict) -> SecurityReport:
        workspace = workspace.resolve()
        exclude_paths = set(task_meta.get("exclude_paths", []))
        allowlist_patterns = self._load_allowlist_patterns(workspace / ".gitleaks.toml")
        scan_mode = task_meta.get("scan_mode", "dir")
        integrity_error = self._git_integrity_error(workspace)
        findings = (
            self._scan_git_history(workspace, exclude_paths, allowlist_patterns)
            if scan_mode == "git"
            else self._scan_directory(workspace, exclude_paths, allowlist_patterns)
        )
        return SecurityReport(
            finding_count=len(findings),
            findings=findings,
            integrity_error=integrity_error,
        )

    def _scan_directory(
        self,
        workspace: Path,
        exclude_paths: set[str],
        allowlist_patterns: list[re.Pattern[str]],
    ) -> list[SecurityFinding]:
        findings: dict[tuple[str, str, str], SecurityFinding] = {}
        for file_path in workspace.rglob("*"):
            if file_path.is_dir():
                continue
            relative = file_path.relative_to(workspace).as_posix()
            if self._should_skip(relative, exclude_paths):
                continue
            try:
                content = file_path.read_text(encoding="utf-8")
            except UnicodeDecodeError:
                continue
            for finding in self._scan_content(relative, content, allowlist_patterns):
                findings[(finding.kind, finding.location, finding.snippet)] = finding
        return sorted(findings.values(), key=lambda item: item.location)

    def _scan_git_history(
        self,
        workspace: Path,
        exclude_paths: set[str],
        allowlist_patterns: list[re.Pattern[str]],
    ) -> list[SecurityFinding]:
        source = self._ensure_git_source(workspace)
        findings: dict[tuple[str, str, str], SecurityFinding] = {}
        try:
            revs = subprocess.run(
                ["git", "rev-list", "--all"],
                cwd=source,
                check=True,
                capture_output=True,
                text=True,
            ).stdout.splitlines()
            for rev in revs:
                files = subprocess.run(
                    ["git", "ls-tree", "-r", "--name-only", rev],
                    cwd=source,
                    check=True,
                    capture_output=True,
                    text=True,
                ).stdout.splitlines()
                for relative in files:
                    if self._should_skip(relative, exclude_paths):
                        continue
                    show = subprocess.run(
                        ["git", "show", f"{rev}:{relative}"],
                        cwd=source,
                        capture_output=True,
                        text=True,
                    )
                    if show.returncode != 0:
                        continue
                    location_prefix = f"{relative}@{rev[:8]}"
                    for finding in self._scan_content(location_prefix, show.stdout, allowlist_patterns):
                        findings[(finding.kind, finding.location, finding.snippet)] = finding
        finally:
            if source != workspace:
                shutil.rmtree(source, ignore_errors=True)
        return sorted(findings.values(), key=lambda item: item.location)

    def _ensure_git_source(self, workspace: Path) -> Path:
        if self._git_integrity_error(workspace) is None:
            return workspace

        temp_root = Path(tempfile.mkdtemp(prefix="secretsaudit_git_"))
        temp_workspace = temp_root / "workspace"
        shutil.copytree(
            workspace,
            temp_workspace,
            ignore=shutil.ignore_patterns(".pytest_cache", "__pycache__", ".secretsaudit_pytest.xml"),
        )
        commands = [
            ["git", "init"],
            ["git", "config", "user.name", "SecretsAuditEnv"],
            ["git", "config", "user.email", "benchmark@example.com"],
            ["git", "add", "."],
            ["git", "commit", "-m", "Recovered working tree"],
        ]
        for command in commands:
            subprocess.run(command, cwd=temp_workspace, check=True, capture_output=True, text=True)
        return temp_workspace

    def _git_integrity_error(self, workspace: Path) -> str | None:
        git_dir = workspace / ".git"
        if not git_dir.exists():
            return "Git integrity check failed: .git directory is missing. Recovered a temporary repository snapshot for scanning."

        probe = subprocess.run(
            ["git", "rev-parse", "--is-inside-work-tree"],
            cwd=workspace,
            capture_output=True,
            text=True,
        )
        if probe.returncode != 0:
            stderr = probe.stderr.strip() or probe.stdout.strip() or "unknown git error"
            return f"Git integrity check failed: repository metadata is corrupted ({stderr}). Recovered a temporary repository snapshot for scanning."
        return None

    def _should_skip(self, relative_path: str, exclude_paths: set[str]) -> bool:
        parts = Path(relative_path).parts
        if ".git" in parts or ".pytest_cache" in parts or "__pycache__" in parts:
            return True
        if relative_path in exclude_paths:
            return True
        return False

    def _load_allowlist_patterns(self, config_path: Path) -> list[re.Pattern[str]]:
        if not config_path.exists():
            return []
        try:
            raw = tomllib.loads(config_path.read_text(encoding="utf-8"))
        except (tomllib.TOMLDecodeError, OSError):
            return []

        regex_strings: list[str] = []

        def collect(node: object) -> None:
            if isinstance(node, dict):
                for key, value in node.items():
                    if key == "regexes" and isinstance(value, list):
                        regex_strings.extend(str(item) for item in value)
                    else:
                        collect(value)
            elif isinstance(node, list):
                for item in node:
                    collect(item)

        collect(raw)
        patterns: list[re.Pattern[str]] = []
        for regex_string in regex_strings:
            try:
                patterns.append(re.compile(regex_string))
            except re.error:
                continue
        return patterns

    def _scan_content(
        self,
        location: str,
        content: str,
        allowlist_patterns: list[re.Pattern[str]],
    ) -> Iterable[SecurityFinding]:
        for kind, pattern in self._PATTERNS:
            for match in pattern.finditer(content):
                snippet = match.group(0)
                if self._is_allowlisted(snippet, allowlist_patterns):
                    continue
                yield SecurityFinding(kind=kind, location=location, snippet=snippet[:80])

        for match in self._ASSIGNMENT_PATTERN.finditer(content):
            secret = match.group(1)
            if self._is_allowlisted(secret, allowlist_patterns):
                continue
            if self._shannon_entropy(secret) >= 3.2:
                yield SecurityFinding(kind="assignment_secret", location=location, snippet=secret[:80])

        for candidate in self._BASE64_CANDIDATE.findall(content):
            decoded = self._decode_base64(candidate)
            if not decoded:
                continue
            for kind, pattern in self._PATTERNS:
                inner_match = pattern.search(decoded)
                if not inner_match:
                    continue
                snippet = inner_match.group(0)
                if self._is_allowlisted(snippet, allowlist_patterns):
                    continue
                yield SecurityFinding(kind=f"base64_{kind}", location=location, snippet=snippet[:80])

    def _decode_base64(self, candidate: str) -> str | None:
        padded = candidate + "=" * (-len(candidate) % 4)
        try:
            decoded = base64.b64decode(padded, validate=False)
        except (binascii.Error, ValueError):
            return None
        try:
            text = decoded.decode("utf-8")
        except UnicodeDecodeError:
            return None
        if not text or "\x00" in text:
            return None
        return text

    def _is_allowlisted(self, text: str, allowlist_patterns: list[re.Pattern[str]]) -> bool:
        return any(pattern.search(text) for pattern in allowlist_patterns)

    def _shannon_entropy(self, value: str) -> float:
        if not value:
            return 0.0
        probabilities = [value.count(char) / len(value) for char in set(value)]
        return -sum(probability * math.log2(probability) for probability in probabilities)