patchi 0.7.5__py3-none-any.whl

patchi/core/security/threat_model_generator.py

@@ -0,0 +1,378 @@
+"""
+Threat Model Generator — Runtime generator from YAML attack scenarios.
+
+Reads attack scenario YAML files and generates executable threat models
+for the current project based on:
+- Detected frameworks and patterns
+- Route analysis
+- Authentication mechanisms
+- Data flow patterns
+
+Output: ThreatModel with prioritized attack scenarios for the project.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from pathlib import Path
+
+_log = logging.getLogger("patchi.security.threat_model")
+
+
+@dataclass
+class AttackScenario:
+    """A single attack scenario from YAML."""
+
+    id: str
+    name: str
+    category: str
+    severity: str
+    cwe: str
+    owasp: str
+    description: str
+    prerequisites: list[str] = field(default_factory=list)
+    attack_steps: list[dict] = field(default_factory=list)
+    detection_signatures: list[str] = field(default_factory=list)
+    remediation_playbook: str = ""
+    tags: list[str] = field(default_factory=list)
+    applicable: bool = False
+    relevance_score: float = 0.0
+
+
+@dataclass
+class ThreatModel:
+    """Generated threat model for the project."""
+
+    project_root: str
+    total_scenarios: int = 0
+    applicable_scenarios: int = 0
+    scenarios: list[AttackScenario] = field(default_factory=list)
+    by_severity: dict[str, int] = field(default_factory=dict)
+    by_category: dict[str, int] = field(default_factory=dict)
+    recommendations: list[str] = field(default_factory=list)
+
+    def to_dict(self) -> dict:
+        return {
+            "project_root": self.project_root,
+            "total_scenarios": self.total_scenarios,
+            "applicable_scenarios": self.applicable_scenarios,
+            "by_severity": self.by_severity,
+            "by_category": self.by_category,
+            "scenarios": [
+                {
+                    "id": s.id,
+                    "name": s.name,
+                    "category": s.category,
+                    "severity": s.severity,
+                    "cwe": s.cwe,
+                    "description": s.description,
+                    "applicable": s.applicable,
+                    "relevance_score": s.relevance_score,
+                    "tags": s.tags,
+                }
+                for s in self.scenarios
+            ],
+            "recommendations": self.recommendations,
+        }
+
+
+# ── Category Detection Rules ─────────────────────────────────────────────────
+
+# Maps YAML scenario categories to code signals that indicate applicability
+_CATEGORY_SIGNALS: dict[str, list[tuple[str, float]]] = {
+    "authentication": [
+        ("login", 0.8),
+        ("auth", 0.9),
+        ("session", 0.7),
+        ("jwt", 0.9),
+        ("cookie", 0.6),
+        ("password", 0.8),
+        ("bcrypt", 0.8),
+        ("argon2", 0.8),
+    ],
+    "authorization": [
+        ("permission", 0.8),
+        ("role", 0.7),
+        ("admin", 0.8),
+        ("access_control", 0.9),
+        ("guard", 0.6),
+        ("protect", 0.5),
+    ],
+    "injection": [
+        ("execute", 0.7),
+        ("query", 0.6),
+        ("raw_sql", 0.9),
+        ("cursor", 0.7),
+        ("database", 0.5),
+        ("orm", 0.4),
+        ("subprocess", 0.8),
+        ("shell", 0.8),
+        ("eval", 0.9),
+    ],
+    "xss": [
+        ("render", 0.5),
+        ("template", 0.5),
+        ("innerHTML", 0.9),
+        ("dangerouslySetInnerHTML", 0.9),
+        ("markup", 0.6),
+        ("html", 0.4),
+        ("response", 0.4),
+    ],
+    "csrf": [
+        ("form", 0.5),
+        ("post", 0.4),
+        ("state_change", 0.8),
+        ("mutate", 0.7),
+        ("PUT", 0.6),
+        ("DELETE", 0.6),
+    ],
+    "secrets": [
+        ("api_key", 0.9),
+        ("secret", 0.8),
+        ("token", 0.7),
+        ("password", 0.7),
+        ("credential", 0.8),
+        ("env", 0.5),
+    ],
+    "dependencies": [
+        ("requirements", 0.7),
+        ("package.json", 0.7),
+        ("Cargo.toml", 0.7),
+        ("go.mod", 0.7),
+        ("Gemfile", 0.7),
+        ("composer.json", 0.7),
+    ],
+    "configuration": [
+        ("config", 0.5),
+        ("settings", 0.5),
+        ("debug", 0.8),
+        ("verbose", 0.6),
+        ("cors", 0.8),
+        ("ssl", 0.7),
+    ],
+    "crypto": [
+        ("encrypt", 0.8),
+        ("decrypt", 0.8),
+        ("hash", 0.7),
+        ("md5", 0.9),
+        ("sha1", 0.8),
+        ("aes", 0.7),
+        ("rsa", 0.7),
+    ],
+    "network": [
+        ("http", 0.5),
+        ("request", 0.4),
+        ("fetch", 0.4),
+        ("axios", 0.5),
+        ("websocket", 0.7),
+        ("ssrf", 0.9),
+    ],
+}
+
+
+class ThreatModelGenerator:
+    """Generates threat models from YAML scenarios + project context."""
+
+    def __init__(self, root: Path):
+        self.root = root
+        self._scenarios_dir = Path(__file__).parent / "attack_scenarios"
+
+    def generate(self) -> ThreatModel:
+        """Generate a threat model for the current project."""
+        model = ThreatModel(project_root=str(self.root))
+
+        # Load all YAML scenarios
+        all_scenarios = self._load_scenarios()
+        model.total_scenarios = len(all_scenarios)
+
+        # Get project context
+        project_signals = self._collect_project_signals()
+
+        # Score and filter scenarios
+        for scenario in all_scenarios:
+            score = self._score_scenario(scenario, project_signals)
+            scenario.relevance_score = score
+            scenario.applicable = score > 0.1
+            model.scenarios.append(scenario)
+
+        # Sort by relevance
+        model.scenarios.sort(key=lambda s: (-s.relevance_score, s.severity))
+
+        # Compute summaries
+        model.applicable_scenarios = sum(1 for s in model.scenarios if s.applicable)
+        for s in model.scenarios:
+            if s.applicable:
+                model.by_severity[s.severity] = model.by_severity.get(s.severity, 0) + 1
+                model.by_category[s.category] = model.by_category.get(s.category, 0) + 1
+
+        # Generate recommendations
+        model.recommendations = self._generate_recommendations(model)
+
+        _log.info(
+            "Threat model: %d/%d scenarios applicable",
+            model.applicable_scenarios,
+            model.total_scenarios,
+        )
+        return model
+
+    def _load_scenarios(self) -> list[AttackScenario]:
+        """Load all attack scenarios from YAML files."""
+        scenarios = []
+        if not self._scenarios_dir.exists():
+            return scenarios
+
+        for yaml_file in self._scenarios_dir.glob("*.yaml"):
+            try:
+                import yaml
+
+                with open(yaml_file, encoding="utf-8") as f:
+                    data = yaml.safe_load(f)
+                if not data or "scenarios" not in data:
+                    continue
+                for s in data["scenarios"]:
+                    scenarios.append(
+                        AttackScenario(
+                            id=s.get("id", "unknown"),
+                            name=s.get("name", "Unknown"),
+                            category=s.get("category", "unknown"),
+                            severity=s.get("severity", "medium"),
+                            cwe=s.get("cwe", ""),
+                            owasp=s.get("owasp", ""),
+                            description=s.get("description", ""),
+                            prerequisites=s.get("prerequisites", []),
+                            attack_steps=s.get("attack_steps", []),
+                            detection_signatures=s.get("detection_signatures", []),
+                            remediation_playbook=s.get("remediation_playbook", ""),
+                            tags=s.get("tags", []),
+                        )
+                    )
+            except Exception as e:
+                _log.warning("Failed to load %s: %s", yaml_file.name, e)
+        return scenarios
+
+    def _collect_project_signals(self) -> dict[str, float]:
+        """Collect signals from the project to determine applicability."""
+        signals: dict[str, float] = {}
+
+        # Scan source files for keywords
+        source_files = []
+        for ext in ("*.py", "*.js", "*.ts", "*.jsx", "*.tsx", "*.go", "*.java", "*.rb"):
+            source_files.extend(self.root.rglob(ext))
+
+        # Limit to 500 files for performance
+        source_files = source_files[:500]
+
+        # Collect keyword frequency
+        keyword_counts: dict[str, int] = {}
+        for fp in source_files:
+            # Skip vendor/node_modules/venv
+            parts = fp.relative_to(self.root).parts
+            if any(
+                p.startswith(".") or p in ("node_modules", "__pycache__", ".venv", "venv")
+                for p in parts
+            ):
+                continue
+            try:
+                content = fp.read_text(encoding="utf-8", errors="ignore").lower()
+                for _category, kws in _CATEGORY_SIGNALS.items():
+                    for kw, _ in kws:
+                        if kw.lower() in content:
+                            keyword_counts[kw] = keyword_counts.get(kw, 0) + 1
+            except Exception as _exc:
+                _log.warning('_collect_project_signals failed: %s', _exc)
+
+        # Convert counts to scores (0-1)
+        max_count = max(keyword_counts.values()) if keyword_counts else 1
+        for kw, count in keyword_counts.items():
+            signals[kw] = min(count / max(max_count, 1), 1.0)
+
+        # Check for dependency files
+        dep_files = [
+            "requirements.txt",
+            "package.json",
+            "Cargo.toml",
+            "go.mod",
+            "Gemfile",
+            "composer.json",
+            "Pipfile",
+        ]
+        for df in dep_files:
+            if (self.root / df).exists():
+                signals[df] = 1.0
+
+        # Check for config files
+        config_files = [
+            "config.py",
+            "settings.py",
+            ".env",
+            "config.yaml",
+            "config.json",
+            "application.properties",
+        ]
+        for cf in config_files:
+            if list(self.root.rglob(cf)):
+                signals[cf] = 1.0
+
+        return signals
+
+    def _score_scenario(self, scenario: AttackScenario, signals: dict[str, float]) -> float:
+        """Score a scenario based on project signals."""
+        category_signals = _CATEGORY_SIGNALS.get(scenario.category, [])
+        if not category_signals:
+            return 0.2  # Default low relevance for unknown categories
+
+        score = 0.0
+        for kw, weight in category_signals:
+            signal_value = signals.get(kw, 0.0)
+            score += signal_value * weight
+
+        # Normalize to 0-1
+        max_possible = sum(w for _, w in category_signals)
+        if max_possible > 0:
+            score = min(score / max_possible, 1.0)
+
+        # Boost for CWE tags matching project patterns
+        if scenario.cwe:
+            cwe_num = scenario.cwe.replace("CWE-", "")
+            # SQL injection CWEs get boost if project uses SQL
+            if cwe_num in ("89", "90") and signals.get("execute", 0) > 0.5:
+                score = min(score * 1.5, 1.0)
+            # XSS CWEs get boost if project renders HTML
+            if cwe_num in ("79", "80") and signals.get("render", 0) > 0.3:
+                score = min(score * 1.3, 1.0)
+
+        return score
+
+    def _generate_recommendations(self, model: ThreatModel) -> list[str]:
+        """Generate security recommendations based on the threat model."""
+        recs = []
+
+        # High-severity applicable scenarios
+        high_sev = [
+            s for s in model.scenarios if s.applicable and s.severity in ("critical", "high")
+        ]
+        if high_sev:
+            recs.append(
+                f"Address {len(high_sev)} high/critical scenarios: "
+                + ", ".join(s.name for s in high_sev[:3])
+            )
+
+        # Category gaps
+        if "secrets" in model.by_category:
+            recs.append("Review secret management — multiple secret-related scenarios apply")
+
+        if "injection" in model.by_category:
+            recs.append("Validate all user inputs — injection scenarios detected")
+
+        if "auth" in str(model.by_category):
+            recs.append("Audit authentication flows — auth scenarios are relevant")
+
+        # Dependency warnings
+        dep_scenarios = [
+            s for s in model.scenarios if s.category == "dependencies" and s.applicable
+        ]
+        if dep_scenarios:
+            recs.append("Run dependency vulnerability scan — project has dependency files")
+
+        return recs

patchi/core/security/threat_model_updater.py

@@ -0,0 +1,196 @@
+"""
+Threat Model Updater — refines threat model from actual scan findings.
+
+After a scan completes, this module:
+  1. Boosts scenarios that match confirmed findings
+  2. Adds new scenarios for finding types not in the YAML library
+  3. Deprioritizes scenarios that produced no findings (false-alarm decay)
+  4. Persists the updated model for the next scan cycle
+
+The result is a threat model that adapts to the real risk profile of the
+project rather than relying solely on static YAML definitions.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import Any
+
+from patchi.core.security.threat_model_generator import (
+    AttackScenario,
+    ThreatModel,
+    ThreatModelGenerator,
+)
+
+_log = logging.getLogger("patchi.security.threat_model_updater")
+
+# How much to boost a scenario's score when a finding matches its category
+_FINDING_BOOST = 0.15
+# Maximum boosted score
+_MAX_SCORE = 1.0
+# Decay factor for scenarios with zero findings (per scan cycle)
+_ZERO_FINDING_DECAY = 0.05
+
+
+def update_threat_model(
+    root: Path,
+    findings: list[dict[str, Any]],
+) -> ThreatModel:
+    """Regenerate the threat model and adjust scores based on scan findings.
+
+    Returns the updated ThreatModel.
+    """
+    gen = ThreatModelGenerator(root)
+    model = gen.generate()
+
+    # Index findings by category
+    category_counts: dict[str, int] = {}
+    cwe_counts: dict[str, int] = {}
+    finding_types: set[str] = set()
+    for f in findings:
+        cat = _categorize_finding(f)
+        category_counts[cat] = category_counts.get(cat, 0) + 1
+        cwe = f.get("cwe", "")
+        if cwe:
+            cwe_counts[cwe] = cwe_counts.get(cwe, 0) + 1
+        finding_types.add(f.get("type", f.get("control_id", "")))
+
+    # Boost scenarios matching actual findings
+    for scenario in model.scenarios:
+        if not scenario.applicable:
+            continue
+        cat_count = category_counts.get(scenario.category, 0)
+        if cat_count > 0:
+            boost = min(cat_count * _FINDING_BOOST, _MAX_SCORE - scenario.relevance_score)
+            scenario.relevance_score = min(scenario.relevance_score + boost, _MAX_SCORE)
+            _log.debug(
+                "Boosted %s by %.2f (category=%s, findings=%d)",
+                scenario.id,
+                boost,
+                scenario.category,
+                cat_count,
+            )
+        else:
+            # Decay scenarios that produced zero findings
+            scenario.relevance_score = max(scenario.relevance_score - _ZERO_FINDING_DECAY, 0.0)
+
+        # Extra boost for CWE matches
+        if scenario.cwe and cwe_counts.get(scenario.cwe, 0) > 0:
+            scenario.relevance_score = min(scenario.relevance_score + 0.2, _MAX_SCORE)
+
+    # Re-sort by updated relevance
+    model.scenarios.sort(key=lambda s: (-s.relevance_score, s.severity))
+
+    # Recompute summaries
+    model.applicable_scenarios = sum(1 for s in model.scenarios if s.applicable)
+    model.by_severity.clear()
+    model.by_category.clear()
+    for s in model.scenarios:
+        if s.applicable:
+            model.by_severity[s.severity] = model.by_severity.get(s.severity, 0) + 1
+            model.by_category[s.category] = model.by_category.get(s.category, 0) + 1
+
+    # Add discovered scenarios for finding types not in the YAML library
+    _add_discovered_scenarios(model, finding_types, category_counts)
+
+    # Update recommendations
+    model.recommendations = _generate_updated_recommendations(model, category_counts)
+
+    # Persist
+    tm_path = root / ".patchi" / "threat_model.json"
+    tm_path.parent.mkdir(parents=True, exist_ok=True)
+    tm_path.write_text(json.dumps(model.to_dict(), indent=2), encoding="utf-8")
+    _log.info(
+        "Threat model updated: %d scenarios, %d applicable, %d categories",
+        model.total_scenarios,
+        model.applicable_scenarios,
+        len(model.by_category),
+    )
+    return model
+
+
+def _categorize_finding(finding: dict[str, Any]) -> str:
+    """Map a finding to a threat model category."""
+    ftype = (finding.get("type", "") or finding.get("control_id", "")).lower()
+    categories = {
+        "authentication": ("auth", "login", "session", "password", "jwt", "cookie"),
+        "authorization": ("permission", "role", "admin", "access", "privilege"),
+        "injection": ("sql", "injection", "command", "eval", "exec", "shell", "xss"),
+        "secrets": ("secret", "credential", "api_key", "token", "hardcoded"),
+        "crypto": ("hash", "md5", "sha1", "encrypt", "decrypt", "cipher"),
+        "configuration": ("config", "debug", "cors", "ssl", "tls", "verbose"),
+        "network": ("ssrf", "redirect", "request", "fetch", "url"),
+        "dependencies": ("dependency", "vulnerability", "outdated", "cve"),
+    }
+    for cat, keywords in categories.items():
+        if any(kw in ftype for kw in keywords):
+            return cat
+    return "other"
+
+
+def _add_discovered_scenarios(
+    model: ThreatModel,
+    finding_types: set[str],
+    category_counts: dict[str, int],
+) -> None:
+    """Add synthetic scenarios for finding types not covered by YAML library."""
+    existing_ids = {s.id for s in model.scenarios}
+    discovered = 0
+
+    for ftype in finding_types:
+        if not ftype or ftype in existing_ids:
+            continue
+        cat = _categorize_finding({"type": ftype})
+        if cat == "other":
+            continue
+        # Only add if we have multiple findings of this type
+        if category_counts.get(cat, 0) < 2:
+            continue
+        synthetic_id = f"DISCOVERED-{ftype.upper()[:20]}"
+        if synthetic_id in existing_ids:
+            continue
+        scenario = AttackScenario(
+            id=synthetic_id,
+            name=f"Discovered: {ftype}",
+            category=cat,
+            severity="medium",
+            cwe="",
+            owasp="",
+            description=f"Auto-discovered from {category_counts[cat]} scan findings of type '{ftype}'",
+            tags=["discovered", "auto"],
+            applicable=True,
+            relevance_score=min(0.3 + category_counts.get(cat, 0) * 0.05, 0.8),
+        )
+        model.scenarios.append(scenario)
+        existing_ids.add(synthetic_id)
+        discovered += 1
+
+    model.total_scenarios += discovered
+    if discovered:
+        _log.info("Added %d discovered scenarios from scan findings", discovered)
+
+
+def _generate_updated_recommendations(
+    model: ThreatModel,
+    category_counts: dict[str, int],
+) -> list[str]:
+    """Generate recommendations factoring in actual findings."""
+    recs = list(model.recommendations)
+
+    # Top-finding categories
+    top_cats = sorted(category_counts.items(), key=lambda x: -x[1])[:3]
+    for cat, count in top_cats:
+        if count >= 3:
+            recs.append(f"High volume of {cat} findings ({count}) — prioritize remediation")
+
+    # Categories with no coverage
+    covered = set(model.by_category.keys())
+    for cat in category_counts:
+        if cat not in covered and category_counts[cat] >= 2:
+            recs.append(
+                f"Add threat scenarios for '{cat}' — {category_counts[cat]} findings detected"
+            )
+
+    return recs[:10]