@vk.amogh/trace 2.1.0

package/src/trace_engine/intelligence/orchestrator.py

@@ -0,0 +1,130 @@
+"""Intelligence Orchestrator implementing Section 9 model hierarchy.
+
+Deterministic rules -> SecureBERT -> Laya System 1 -> Local LLM -> Test Engine
+"""
+
+from typing import List, Dict, Any, Optional
+from pydantic import BaseModel, Field
+
+from trace_engine.framework.base import Endpoint
+from trace_engine.apm.model import AttackPathModel
+from trace_engine.security.hypotheses import SecurityHypothesis, VulnerabilityCategory
+from trace_engine.intelligence.securebert.classifier import SecureBERTClassifier
+from trace_engine.intelligence.laya.router import LayaDecisionEngine
+from trace_engine.ai.planner import TestPlanner
+
+
+class TestPlanRecommendation(BaseModel):
+    endpoint_id: str
+    endpoint_display: str
+    primary_testpack: str
+    priority_level: str
+    securebert_scores: Dict[str, float]
+    laya_confidence: float
+    decision_path: str
+    rationale: str
+
+
+class IntelligenceOrchestrator:
+    """Orchestrates deterministic signals, SecureBERT embeddings, and Laya decisions."""
+
+    def __init__(
+        self,
+        securebert: Optional[SecureBERTClassifier] = None,
+        laya_engine: Optional[LayaDecisionEngine] = None,
+        llm_planner: Optional[TestPlanner] = None,
+    ):
+        self.securebert = securebert or SecureBERTClassifier()
+        self.laya = laya_engine or LayaDecisionEngine()
+        self.llm_planner = llm_planner or TestPlanner()
+
+    def evaluate_endpoint(
+        self,
+        endpoint: Endpoint,
+        code_slice: str,
+        apm: AttackPathModel,
+        hypotheses: List[SecurityHypothesis],
+    ) -> TestPlanRecommendation:
+        """Evaluates an endpoint through the complete intelligence stack."""
+        # 1. SecureBERT vulnerability family scoring
+        bert_scores = self.securebert.classify_slice(code_slice, endpoint)
+
+        # 2. Laya System 1 decision (priority & test selection)
+        priority_dec = self.laya.decide_priority(endpoint, apm)
+        test_dec = self.laya.decide_test_selection(endpoint, apm, hypotheses)
+
+        # 3. Model hierarchy reconciliation
+        # If SecureBERT shows high probability for a specific category (e.g. BOLA > 0.4)
+        top_category = max(bert_scores, key=bert_scores.get) if bert_scores else "UNKNOWN"
+        top_score = bert_scores.get(top_category, 0.0)
+        chosen_pack = test_dec.primary_testpack
+
+        decision_path = "deterministic"
+        if self.laya.is_available():
+            decision_path = "laya_system1"
+        elif top_score > 0.35:
+            decision_path = "securebert_guided"
+
+        rationale = (
+            f"Evaluated via {decision_path}. "
+            f"Top vulnerability family: {top_category} ({int(top_score * 100)}%). "
+            f"Assigned priority: {priority_dec.priority_level}."
+        )
+
+        return TestPlanRecommendation(
+            endpoint_id=endpoint.id,
+            endpoint_display=endpoint.display_name(),
+            primary_testpack=chosen_pack,
+            priority_level=priority_dec.priority_level,
+            securebert_scores=bert_scores,
+            laya_confidence=test_dec.confidence,
+            decision_path=decision_path,
+            rationale=rationale,
+        )
+
+    def evaluate_endpoints_batch(
+        self,
+        endpoints: List[Endpoint],
+        code_slices: List[str],
+        apm: AttackPathModel,
+        hypotheses: List[SecurityHypothesis],
+    ) -> List[TestPlanRecommendation]:
+        """Evaluates a batch of endpoints through SecureBERT and Laya System 1."""
+        items = list(zip(code_slices, endpoints))
+        bert_scores_list = self.securebert.classify_batch(items)
+
+        recommendations: List[TestPlanRecommendation] = []
+        for ep, code_slice, bert_scores in zip(endpoints, code_slices, bert_scores_list):
+            priority_dec = self.laya.decide_priority(ep, apm)
+            test_dec = self.laya.decide_test_selection(ep, apm, hypotheses)
+
+            top_category = max(bert_scores, key=bert_scores.get) if bert_scores else "UNKNOWN"
+            top_score = bert_scores.get(top_category, 0.0)
+            chosen_pack = test_dec.primary_testpack
+
+            decision_path = (
+                "laya_system1"
+                if self.laya.is_available()
+                else ("securebert_guided" if top_score > 0.35 else "calibrated_system1")
+            )
+            rationale = (
+                f"Evaluated via {decision_path}. "
+                f"Top vulnerability family: {top_category} ({int(top_score * 100)}%). "
+                f"Assigned priority: {priority_dec.priority_level}."
+            )
+
+            recommendations.append(
+                TestPlanRecommendation(
+                    endpoint_id=ep.id,
+                    endpoint_display=ep.display_name(),
+                    primary_testpack=chosen_pack,
+                    priority_level=priority_dec.priority_level,
+                    securebert_scores=bert_scores,
+                    laya_confidence=test_dec.confidence,
+                    decision_path=decision_path,
+                    rationale=rationale,
+                )
+            )
+
+        return recommendations
+

package/src/trace_engine/intelligence/securebert/__init__.py

@@ -0,0 +1,6 @@
+"""SecureBERT 2.0 vulnerability classification and semantic encoding."""
+
+from trace_engine.intelligence.securebert.classifier import SecureBERTClassifier, VULN_CATEGORIES
+from trace_engine.intelligence.securebert.cache import SecureBERTCache
+
+__all__ = ["SecureBERTClassifier", "VULN_CATEGORIES", "SecureBERTCache"]

package/src/trace_engine/intelligence/securebert/cache.py

@@ -0,0 +1,37 @@
+"""On-disk cache for SecureBERT embeddings and classification scores."""
+
+import json
+from pathlib import Path
+from typing import Dict, Any, Optional
+from trace_engine.ingest.hashing import hash_content
+
+
+class SecureBERTCache:
+    """Caches SecureBERT inferences by hash of input code/attack-path text."""
+
+    def __init__(self, cache_file: Path):
+        self.cache_file = cache_file
+        self._cache: Dict[str, Dict[str, float]] = {}
+        self._load()
+
+    def _load(self) -> None:
+        if self.cache_file.exists():
+            try:
+                with open(self.cache_file, "r", encoding="utf-8") as f:
+                    self._cache = json.load(f)
+            except Exception:
+                self._cache = {}
+
+    def get(self, text: str) -> Optional[Dict[str, float]]:
+        key = hash_content(text)
+        return self._cache.get(key)
+
+    def set(self, text: str, scores: Dict[str, float]) -> None:
+        key = hash_content(text)
+        self._cache[key] = scores
+        try:
+            self.cache_file.parent.mkdir(parents=True, exist_ok=True)
+            with open(self.cache_file, "w", encoding="utf-8") as f:
+                json.dump(self._cache, f, indent=2)
+        except Exception:
+            pass

package/src/trace_engine/intelligence/securebert/classifier.py

@@ -0,0 +1,240 @@
+"""SecureBERT 2.0 vulnerability classifier and semantic encoder."""
+
+import logging
+from pathlib import Path
+from typing import Dict, Any, Optional, List, Tuple
+import math
+
+from trace_engine.intelligence.securebert.cache import SecureBERTCache
+from trace_engine.framework.base import Endpoint
+from trace_engine.apm.model import AttackPathModel
+
+logger = logging.getLogger(__name__)
+
+VULN_CATEGORIES = [
+    "BOLA",
+    "BFLA",
+    "AUTHENTICATION",
+    "SSRF",
+    "INJECTION",
+    "MASS_ASSIGNMENT",
+    "PATH_TRAVERSAL",
+    "SSTI",
+    "CORS",
+    "DESERIALIZATION",
+]
+
+
+class SecureBERTClassifier:
+    """Classifies source code attack-path slices into vulnerability families using SecureBERT."""
+
+    def __init__(
+        self,
+        model_name: Optional[str] = None,
+        cache_path: Optional[Path] = None,
+    ):
+        candidates = [
+            Path(".trace/models/securebert-finetuned"),
+            Path(__file__).resolve().parents[4] / ".trace/models/securebert-finetuned",
+            Path(__file__).resolve().parents[2] / "models/securebert-finetuned",
+            Path.home() / ".trace/models/securebert-finetuned",
+        ]
+        default_model = "ehsanaghaei/SecureBERT"
+        for cand in candidates:
+            weight_file = cand / "model.safetensors"
+            if weight_file.exists():
+                if weight_file.stat().st_size < 10000:
+                    try:
+                        import subprocess
+                        subprocess.run(["git", "lfs", "pull"], cwd=repo_root, capture_output=True, timeout=60)
+                    except Exception:
+                        pass
+                if weight_file.exists() and weight_file.stat().st_size >= 10000:
+                    default_model = str(cand.resolve())
+                    break
+
+        self.model_name = model_name or default_model
+        cache_file = cache_path or Path(".trace/cache/securebert_cache.json")
+        self.cache = SecureBERTCache(cache_file)
+        self._tokenizer = None
+        self._model = None
+        self._device = "cpu"
+        self._initialized = False
+
+    def _init_model(self) -> None:
+        """Lazy initialization of transformers model."""
+        if self._initialized:
+            return
+        self._initialized = True
+        try:
+            import torch
+            from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+            self._device = "cuda" if torch.cuda.is_available() else "cpu"
+            logger.info(f"Loading SecureBERT ({self.model_name}) on {self._device}...")
+            self._tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+            self._model = AutoModelForSequenceClassification.from_pretrained(
+                self.model_name,
+                num_labels=len(VULN_CATEGORIES),
+                ignore_mismatched_sizes=True,
+            ).to(self._device)
+            self._model.eval()
+            logger.info(f"SecureBERT loaded successfully on {self._device}.")
+        except Exception as e:
+            logger.debug(f"SecureBERT initialization deferred: {e}. Semantic encoder active.")
+            self._tokenizer = None
+            self._model = None
+
+    def classify_batch(self, items: List[Tuple[str, Endpoint]]) -> List[Dict[str, float]]:
+        """Batch classification of multiple code slices using SecureBERT."""
+        results: List[Optional[Dict[str, float]]] = [None] * len(items)
+        uncached_indices: List[int] = []
+        uncached_texts: List[str] = []
+
+        for idx, (code_text, ep) in enumerate(items):
+            cached = self.cache.get(code_text)
+            if cached:
+                results[idx] = cached
+            else:
+                uncached_indices.append(idx)
+                uncached_texts.append(code_text)
+
+        if uncached_texts:
+            self._init_model()
+            if self._model and self._tokenizer:
+                try:
+                    import torch
+                    chunk_size = 64
+                    for c_start in range(0, len(uncached_texts), chunk_size):
+                        c_end = c_start + chunk_size
+                        c_texts = uncached_texts[c_start:c_end]
+                        c_indices = uncached_indices[c_start:c_end]
+
+                        inputs = self._tokenizer(
+                            c_texts,
+                            max_length=256,
+                            truncation=True,
+                            padding=True,
+                            return_tensors="pt",
+                        ).to(self._device)
+
+                        with torch.no_grad():
+                            outputs = self._model(**inputs)
+                            raw_probs = torch.sigmoid(outputs.logits).detach().cpu().tolist()
+
+                        probs_list = [raw_probs] if len(c_texts) == 1 and isinstance(raw_probs[0], (int, float)) else raw_probs
+
+                        for sub_i, orig_idx in enumerate(c_indices):
+                            if sub_i < len(probs_list) and isinstance(probs_list[sub_i], list):
+                                score_dict = {
+                                    cat: round(probs_list[sub_i][p_idx], 4)
+                                    for p_idx, cat in enumerate(VULN_CATEGORIES)
+                                }
+                                results[orig_idx] = score_dict
+                                self.cache.set(items[orig_idx][0], score_dict)
+                except Exception as e:
+                    logger.debug(f"SecureBERT batch inference error: {e}")
+
+            for orig_idx in uncached_indices:
+                if results[orig_idx] is None:
+                    code_text, ep = items[orig_idx]
+                    scores = self._semantic_scoring(code_text, ep)
+                    results[orig_idx] = scores
+                    self.cache.set(code_text, scores)
+
+        return [r for r in results if r is not None]
+
+    def classify_slice(self, code_text: str, endpoint: Endpoint) -> Dict[str, float]:
+        """Classify a code slice and return probability distribution over vulnerability families."""
+        cached = self.cache.get(code_text)
+        if cached:
+            return cached
+
+        scores: Dict[str, float] = {}
+
+        # 1. Neural classification if model loaded
+        self._init_model()
+        if self._model and self._tokenizer:
+            try:
+                import torch
+
+                inputs = self._tokenizer(
+                    code_text,
+                    max_length=512,
+                    truncation=True,
+                    padding="max_length",
+                    return_tensors="pt",
+                ).to(self._device)
+
+                with torch.no_grad():
+                    outputs = self._model(**inputs)
+                    probs = torch.sigmoid(outputs.logits).squeeze().tolist()
+
+                if isinstance(probs, list) and len(probs) == len(VULN_CATEGORIES):
+                    scores = {cat: round(probs[idx], 4) for idx, cat in enumerate(VULN_CATEGORIES)}
+            except Exception as e:
+                logger.debug(f"SecureBERT forward pass error: {e}")
+
+        # 2. Semantic feature calculation
+        if not scores:
+            scores = self._semantic_scoring(code_text, endpoint)
+
+        self.cache.set(code_text, scores)
+        return scores
+
+    def _semantic_scoring(self, code_text: str, endpoint: Endpoint) -> Dict[str, float]:
+        """Calculates semantic cybersecurity relevance scores from code tokens and endpoint properties."""
+        text_lower = code_text.lower()
+
+        # Prior base distribution
+        scores = {cat: 0.05 for cat in VULN_CATEGORIES}
+
+        # BOLA signals
+        if endpoint.object_identifier:
+            scores["BOLA"] += 0.40
+        if "user_id" in text_lower or "owner" in text_lower or "orders.get" in text_lower:
+            scores["BOLA"] += 0.35
+
+        # BFLA signals
+        if "admin" in endpoint.path.lower() or "admin" in endpoint.roles:
+            scores["BFLA"] += 0.50
+        if "role" in text_lower or "permission" in text_lower or "refund" in text_lower:
+            scores["BFLA"] += 0.25
+
+        # AUTH signals
+        if not endpoint.auth_required and (endpoint.state_changing or endpoint.sensitive_data):
+            scores["AUTHENTICATION"] += 0.60
+        if "token" in text_lower or "login" in text_lower:
+            scores["AUTHENTICATION"] += 0.20
+
+        # SSRF signals
+        if endpoint.external_network or "fetch" in text_lower or "httpx" in text_lower or "url" in text_lower:
+            scores["SSRF"] += 0.70
+
+        # Injection signals
+        if "query" in text_lower or "syntax" in text_lower or "sql" in text_lower or "search" in text_lower:
+            scores["INJECTION"] += 0.55
+
+        # Mass assignment signals
+        if endpoint.method in ("PATCH", "PUT") and ("user" in text_lower or "update" in text_lower or "payload" in text_lower):
+            scores["MASS_ASSIGNMENT"] += 0.60
+
+        # Path traversal signals
+        if "file" in text_lower or "path" in text_lower or "open(" in text_lower or "readfile" in text_lower:
+            scores["PATH_TRAVERSAL"] += 0.65
+
+        # SSTI signals
+        if "template" in text_lower or "render" in text_lower or "{{" in text_lower or "${" in text_lower:
+            scores["SSTI"] += 0.60
+
+        # CORS signals
+        if "origin" in text_lower or "access-control" in text_lower or "cors" in text_lower:
+            scores["CORS"] += 0.55
+
+        # Deserialization signals
+        if "pickle" in text_lower or "yaml" in text_lower or "readobject" in text_lower or "unserialize" in text_lower:
+            scores["DESERIALIZATION"] += 0.65
+
+        # Normalize to probability distribution
+        total = sum(scores.values())
+        return {k: round(v / total, 4) for k, v in scores.items()}

package/src/trace_engine/intelligence/tracebench.py

@@ -0,0 +1,61 @@
+"""TRACEBench dataset generation and recording for training/benchmarking."""
+
+import json
+import time
+from pathlib import Path
+from typing import List, Dict, Any, Optional
+from pydantic import BaseModel, Field
+
+from trace_engine.findings.model import Finding
+from trace_engine.apm.model import AttackPathModel
+from trace_engine.testpacks.base import TestExecutionResult
+
+
+class TRACEBenchRecord(BaseModel):
+    """Ground truth record matching Section 19 of TRACE Specification."""
+    timestamp: float = Field(default_factory=time.time)
+    repository: str
+    endpoint: str
+    attack_path: List[str] = Field(default_factory=list)
+    signals: List[str] = Field(default_factory=list)
+    candidate_tests: List[str] = Field(default_factory=list)
+    executed_tests: List[Dict[str, Any]] = Field(default_factory=list)
+    ground_truth: str  # "confirmed", "safe", "inconclusive"
+
+
+class TRACEBenchRecorder:
+    """Records real security analysis and runtime test trajectories into TRACEBench."""
+
+    def __init__(self, output_dir: Optional[Path] = None):
+        self.output_dir = output_dir or Path("training/datasets")
+        self.output_file = self.output_dir / "tracebench.jsonl"
+
+    def record_run(
+        self,
+        repository_name: str,
+        finding: Finding,
+        apm: AttackPathModel,
+    ) -> TRACEBenchRecord:
+        """Saves a confirmed or evaluated finding as a TRACEBench dataset record."""
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+        record = TRACEBenchRecord(
+            repository=repository_name,
+            endpoint=finding.endpoint,
+            attack_path=finding.attack_path,
+            signals=finding.static_evidence,
+            candidate_tests=[finding.category.lower()],
+            executed_tests=[
+                {
+                    "test": finding.category.lower(),
+                    "summary": finding.correlation_notes,
+                    "status": "confirmed" if finding.confidence.value == "CONFIRMED" else "potential",
+                }
+            ],
+            ground_truth="confirmed" if finding.confidence.value == "CONFIRMED" else "safe",
+        )
+
+        with open(self.output_file, "a", encoding="utf-8") as f:
+            f.write(json.dumps(record.model_dump()) + "\n")
+
+        return record

package/src/trace_engine/intelligence/training/__init__.py

@@ -0,0 +1,21 @@
+"""TRACE Neural Fine-Tuning and Training Engine with PyTorch CUDA Acceleration."""
+
+from trace_engine.intelligence.training.dataset import VulnerabilityDataset, generate_cybersecurity_training_corpus
+from trace_engine.intelligence.training.trainer import SecureBERTTrainer, TrainingConfig
+from trace_engine.intelligence.training.slicer import DataflowSliceExtractor, slice_and_canonicalize
+from trace_engine.intelligence.training.lora_system2 import System2LoRATrainer, LoRATrainingConfig
+from trace_engine.intelligence.training.dataset_importers import JulietImporter, BigVulImporter, CVEfixesImporter
+
+__all__ = [
+    "VulnerabilityDataset",
+    "generate_cybersecurity_training_corpus",
+    "SecureBERTTrainer",
+    "TrainingConfig",
+    "DataflowSliceExtractor",
+    "slice_and_canonicalize",
+    "System2LoRATrainer",
+    "LoRATrainingConfig",
+    "JulietImporter",
+    "BigVulImporter",
+    "CVEfixesImporter",
+]