topoexploit 0.1.0__py3-none-any.whl

topogpt3/exploitgym_loader.py

@@ -0,0 +1,416 @@
+#!/usr/bin/env python3
+"""
+ExploitGym Data Loader for TopoExploit
+
+Clones the ExploitGym repo and tokenizes task data into training memmaps.
+Each task contains vulnerability descriptions, patches, exploit docs, and PoC source.
+
+Tier structure:
+  Tier 1: vulnerability_analysis -- CVE description -> vulnerability analysis
+  Tier 2: patch_analysis        -- source + patch -> fix explanation
+  Tier 3: exploit_development   -- vulnerability -> exploit strategy + PoC code
+
+Author: Gris Iscomeback
+License: GPL v3
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import subprocess
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+
+from .model import BPETokenizer
+
+
+EXPLOITGYM_REPO = "https://github.com/sunblaze-ucb/exploitgym.git"
+EXPLOITGYM_BRANCH = "main"
+TASK_IDS_FILE = "data/task_ids/v1.txt"
+
+
+@dataclass
+class ExploitGymLoaderConfig:
+    data_cache: str = "data_exploitgym"
+    repo_cache: str = ".exploitgym_repo"
+    max_tokens_per_tier: Tuple[int, int, int] = (50_000_000, 80_000_000, 100_000_000)
+    val_fraction: float = 0.05
+    holdout_fraction: float = 0.05
+    split_seed: int = 2026
+    max_file_chars: int = 32_000
+    tier_names: Tuple[str, str, str] = (
+        "vulnerability_analysis",
+        "patch_analysis",
+        "exploit_development",
+    )
+
+
+def ensure_repo(repo_cache: str, logger: logging.Logger) -> Path:
+    repo_path = Path(repo_cache)
+    if repo_path.exists() and (repo_path / ".git").exists():
+        logger.info(f"ExploitGym repo found at {repo_path}, pulling latest...")
+        try:
+            subprocess.run(
+                ["git", "-C", str(repo_path), "pull", "--ff-only"],
+                capture_output=True, timeout=60,
+            )
+        except Exception:
+            logger.warning("git pull failed, using existing repo")
+    else:
+        logger.info(f"Cloning ExploitGym to {repo_path}...")
+        subprocess.run(
+            ["git", "clone", "--depth=1", "-b", EXPLOITGYM_BRANCH,
+             EXPLOITGYM_REPO, str(repo_path)],
+            check=True, timeout=300,
+        )
+    return repo_path
+
+
+def _read_file_safe(path: Path, max_chars: int = 32000) -> str:
+    try:
+        text = path.read_text(encoding="utf-8", errors="replace")
+        return text[:max_chars]
+    except (OSError, UnicodeDecodeError):
+        return ""
+
+
+def _parse_task_ids(repo_path: Path) -> List[str]:
+    ids_path = repo_path / TASK_IDS_FILE
+    if not ids_path.exists():
+        return []
+    return [
+        line.strip() for line in ids_path.read_text().splitlines()
+        if line.strip() and not line.startswith("#")
+    ]
+
+
+def _task_id_to_path(task_id: str) -> Optional[Path]:
+    parts = task_id.split(":", 1)
+    if len(parts) != 2:
+        return None
+    family, name = parts
+    return Path(family) / name
+
+
+def parse_task(task_dir: Path, task_id: str, max_chars: int) -> Optional[Dict[str, str]]:
+    family = task_id.split(":")[0] if ":" in task_id else "unknown"
+    rel = _task_id_to_path(task_id)
+    if rel is None:
+        return None
+    full_dir = task_dir / rel
+    if not full_dir.exists():
+        return None
+
+    result: Dict[str, str] = {"task_id": task_id, "family": family}
+
+    if family == "kernel":
+        result["description"] = _read_file_safe(full_dir / "docs" / "vulnerability.md", max_chars)
+        result["exploit_doc"] = _read_file_safe(full_dir / "docs" / "exploit.md", max_chars)
+        result["novel_techniques"] = _read_file_safe(full_dir / "docs" / "novel-techniques.md", max_chars)
+        result["reference_exploit"] = _read_file_safe(full_dir / "docs" / "reference-exploit.md", max_chars)
+        result["patch"] = _read_file_safe(full_dir / "patch.diff", max_chars)
+        result["pov_c"] = _read_file_safe(full_dir / "pov" / "pov.c", max_chars)
+    elif family == "user":
+        result["description"] = _read_file_safe(full_dir / "description.txt", max_chars)
+        result["patch"] = _read_file_safe(full_dir / "patch.diff", max_chars)
+        result["pov_c"] = _read_file_safe(full_dir / "pov" / "pov.c", max_chars)
+    elif family == "v8":
+        result["description"] = _read_file_safe(full_dir / "description.txt", max_chars)
+        result["patch"] = _read_file_safe(full_dir / "patch.diff", max_chars)
+        result["pov_c"] = _read_file_safe(full_dir / "pov" / "pov.c", max_chars)
+    else:
+        return None
+
+    has_any = any(result.get(k, "").strip() for k in
+                  ["description", "exploit_doc", "patch", "pov_c"])
+    return result if has_any else None
+
+
+
+
+def _format_vulnerability_analysis(task: Dict[str, str]) -> Optional[str]:
+    tid = task.get("task_id", "unknown")
+    family = task.get("family", "unknown")
+    desc = task.get("description", "").strip()
+    exploit_doc = task.get("exploit_doc", "").strip()
+    if not desc and not exploit_doc:
+        return None
+    parts = [f"### Task: {tid} ({family} vulnerability)"]
+    if desc:
+        parts.append(f"### Vulnerability Description\n{desc}")
+    if exploit_doc:
+        lines = exploit_doc.split("\n")[:100]
+        parts.append("### Vulnerability Analysis\n" + "\n".join(lines))
+    return "\n\n".join(parts) + "\n"
+
+
+def _format_patch_analysis(task: Dict[str, str]) -> Optional[str]:
+    patch = task.get("patch", "").strip()
+    desc = task.get("description", "").strip()
+    exploit_doc = task.get("exploit_doc", "").strip()
+    if not patch:
+        return None
+    parts = [f"### Task: {task.get('task_id', 'unknown')}"]
+    if desc:
+        parts.append(f"### Vulnerability Context\n{desc}")
+    parts.append(f"### Patch\n```diff\n{patch}\n```")
+    if exploit_doc:
+        parts.append(f"### Exploit Documentation\n{exploit_doc[:4000]}")
+    return "\n\n".join(parts) + "\n"
+
+
+def _format_exploit_development(task: Dict[str, str]) -> Optional[str]:
+    desc = task.get("description", "").strip()
+    exploit_doc = task.get("exploit_doc", "").strip()
+    novel_doc = task.get("novel_techniques", "").strip()
+    pov_c = task.get("pov_c", "").strip()
+    ref_exploit = task.get("reference_exploit", "").strip()
+    if not desc and not pov_c:
+        return None
+    parts = []
+    tid = task.get("task_id", "unknown")
+    family = task.get("family", "unknown")
+    parts.append(f"### Task: {tid} ({family} exploit development)")
+    if desc:
+        parts.append(f"### Vulnerability\n{desc}")
+    if exploit_doc:
+        parts.append(f"### Exploit Strategy\n{exploit_doc[:4000]}")
+    if novel_doc:
+        parts.append(f"### Novel Techniques\n{novel_doc[:2000]}")
+    if ref_exploit:
+        parts.append(f"### Reference Exploit\n{ref_exploit[:4000]}")
+    if pov_c:
+        parts.append(f"### Proof of Concept Code\n```c\n{pov_c[:8000]}\n```")
+    return "\n\n".join(parts) + "\n"
+
+
+FORMATTERS = {
+    "vulnerability_analysis": _format_vulnerability_analysis,
+    "patch_analysis": _format_patch_analysis,
+    "exploit_development": _format_exploit_development,
+}
+
+
+class ExploitGymDataLoader:
+    def __init__(self, config: ExploitGymLoaderConfig, tokenizer: BPETokenizer,
+                 logger: logging.Logger):
+        self.config = config
+        self.tokenizer = tokenizer
+        self.logger = logger
+        self.cache_dir = Path(config.data_cache)
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+
+    def _tier_paths(self, tier: str) -> Dict[str, Path]:
+        return {
+            split: self.cache_dir / f"tokens_{tier}_{split}.bin"
+            for split in ("train", "val", "holdout")
+        }
+
+    def _manifest_path(self, tier: str) -> Path:
+        return self.cache_dir / f"manifest_{tier}.json"
+
+    def _already_prepared(self, tier: str) -> bool:
+        paths = self._tier_paths(tier)
+        manifest = self._manifest_path(tier)
+        if not manifest.exists():
+            return False
+        if not all(p.exists() for p in paths.values()):
+            return False
+        try:
+            with open(manifest, "r") as f:
+                meta = json.load(f)
+            sizes = meta.get("sizes", {})
+        except (json.JSONDecodeError, OSError):
+            return False
+        for split, p in paths.items():
+            disk_n = p.stat().st_size // 4
+            if disk_n == 0:
+                return False
+            expected = sizes.get(split, 0)
+            if expected > 0 and abs(disk_n - expected) > 4:
+                return False
+        return True
+
+    def _load_all_tasks(self, repo_path: Path) -> List[Dict[str, str]]:
+        task_dir = repo_path / "data" / "tasks"
+        task_ids = _parse_task_ids(repo_path)
+        if not task_ids:
+            self.logger.warning("No task IDs found, scanning directories...")
+            for family_dir in ["kernel", "user", "v8"]:
+                fpath = task_dir / family_dir
+                if not fpath.exists():
+                    continue
+                for subset in fpath.iterdir():
+                    if subset.is_dir():
+                        for task in subset.iterdir():
+                            if task.is_dir():
+                                tid = f"{family_dir}:{subset.name}/{task.name}"
+                                task_ids.append(tid)
+        self.logger.info(f"Found {len(task_ids)} task IDs")
+        tasks = []
+        for tid in task_ids:
+            t = parse_task(task_dir, tid, self.config.max_file_chars)
+            if t is not None:
+                tasks.append(t)
+        self.logger.info(f"Parsed {len(tasks)} tasks with data")
+        return tasks
+
+    def prepare_tier(self, tier_index: int, force: bool = False) -> Dict[str, int]:
+        tier = self.config.tier_names[tier_index]
+        if not force and self._already_prepared(tier):
+            with open(self._manifest_path(tier), "r") as f:
+                return json.load(f)["sizes"]
+
+        repo_path = ensure_repo(self.config.repo_cache, self.logger)
+        all_tasks = self._load_all_tasks(repo_path)
+        formatter = FORMATTERS[tier]
+
+        # Split ESTRATIFICADO: reparte train/val/holdout con la MISMA
+        # distribucion de caracteristicas estructurales (pov_c, patch,
+        # familia, longitud) para evitar distribution shift entre splits.
+        #
+        # Antes el split era shuffle aleatorio -> el holdout podia quedarse
+        # sin tareas con pov_c mientras train si las tenia, lo que hacfa que
+        # el modelo aprendiera a generar PoC pero el holdout no las tuviera
+        # (loss alto por mismatch de formato, NO por overfitting).
+        rng = np.random.default_rng(self.config.split_seed + tier_index)
+        n = len(all_tasks)
+        n_hold = max(1, int(n * self.config.holdout_fraction))
+        n_val = max(1, int(n * self.config.val_fraction))
+
+        # Estrato = (familia, tiene_pov_c, tiene_patch, bucket_longitud)
+        def stratum_key(task: Dict[str, str]) -> Tuple:
+            f = formatter(task)
+            length = len(f) if f else 0
+            if length == 0:
+                len_bucket = 0
+            elif length < 500:
+                len_bucket = 1
+            elif length < 2000:
+                len_bucket = 2
+            else:
+                len_bucket = 3
+            return (
+                task.get("family", "unknown"),
+                1 if len(task.get("pov_c", "") or "") > 0 else 0,
+                1 if len(task.get("patch", "") or "") > 0 else 0,
+                len_bucket,
+            )
+
+        # Agrupar tareas por estrato
+        strata: Dict[Tuple, List[int]] = {}
+        for i, task in enumerate(all_tasks):
+            key = stratum_key(task)
+            strata.setdefault(key, []).append(i)
+
+        hold_idx: set = set()
+        val_idx: set = set()
+
+        # Dentro de cada estrato, repartir proporcionalmente
+        for key, members in strata.items():
+            rng.shuffle(members)
+            m = len(members)
+            # Asignar n_hold / n_val de forma proporcional al tamano del estrato
+            n_hold_s = max(1, int(round(m * self.config.holdout_fraction)))
+            n_val_s = max(1, int(round(m * self.config.val_fraction)))
+            # Asegurar que no exceda el tamano del estrato
+            n_hold_s = min(n_hold_s, m)
+            n_val_s = min(n_val_s, m - n_hold_s)
+            hold_idx.update(members[:n_hold_s])
+            val_idx.update(members[n_hold_s:n_hold_s + n_val_s])
+
+        self.logger.info(
+            f"[{tier}] stratified split: holdout={len(hold_idx)} val={len(val_idx)} "
+            f"train={n - len(hold_idx) - len(val_idx)} (estratos={len(strata)})"
+        )
+
+        eot = self.tokenizer.eot_token()
+        max_tokens = self.config.max_tokens_per_tier[tier_index]
+        sizes = {"train": 0, "val": 0, "holdout": 0}
+        files = {
+            split: open(self._tier_paths(tier)[split], "wb")
+            for split in ("train", "val", "holdout")
+        }
+        batch_tokens: Dict[str, List[int]] = {k: [] for k in files}
+        CHUNK = 1_000_000
+
+        def flush(split: str) -> None:
+            arr = np.array(batch_tokens[split], dtype=np.int32)
+            files[split].write(arr.tobytes())
+            sizes[split] += arr.size
+            batch_tokens[split] = []
+
+        train_cap_hit = False
+        # Dedup exacto del texto formateado en el split train (evita retrain
+        # sobre tareas casi identicas que pasan por el mismo template).
+        seen_train_text: set = set()
+        for i, task in enumerate(all_tasks):
+            if i in hold_idx:
+                split = "holdout"
+            elif i in val_idx:
+                split = "val"
+            else:
+                if train_cap_hit:
+                    continue
+                split = "train"
+
+            text = formatter(task)
+            if not text:
+                continue
+            # Dedup exacto solo en train
+            if split == "train":
+                tkey = text[:5000]
+                if tkey in seen_train_text:
+                    continue
+                seen_train_text.add(tkey)
+            try:
+                ids = self.tokenizer.encode(text)
+            except Exception:
+                continue
+            batch_tokens[split].extend(ids)
+            batch_tokens[split].append(eot)
+            if len(batch_tokens[split]) >= CHUNK:
+                flush(split)
+            if sizes["train"] >= max_tokens > 0:
+                train_cap_hit = True
+
+        for split in files:
+            if batch_tokens[split]:
+                flush(split)
+            files[split].close()
+
+        if sizes["train"] == 0:
+            for p in self._tier_paths(tier).values():
+                try:
+                    p.unlink()
+                except OSError:
+                    pass
+            raise RuntimeError(f"[{tier}] No training examples produced.")
+
+        manifest = {
+            "tier": tier,
+            "sizes": sizes,
+            "split_seed": self.config.split_seed + tier_index,
+            "saved_at": time.strftime("%Y-%m-%dT%H:%M:%S"),
+            "n_tasks": n,
+        }
+        with open(self._manifest_path(tier), "w") as f:
+            json.dump(manifest, f, indent=2)
+        self.logger.info(
+            f"[{tier}] train={sizes['train']:,} val={sizes['val']:,} "
+            f"holdout={sizes['holdout']:,} tokens from {n} tasks"
+        )
+        return sizes
+
+    def open_memmap(self, tier: str, split: str) -> np.memmap:
+        path = self._tier_paths(tier)[split]
+        if not path.exists():
+            raise FileNotFoundError(f"Missing {path}. Run --prepare-exploitgym first.")
+        n = path.stat().st_size // 4
+        return np.memmap(path, dtype=np.int32, mode="r", shape=(n,))

topogpt3/hodge_cm.py

@@ -0,0 +1,160 @@
+#!/usr/bin/env python3
+"""Fases 2-5: martillo Hodge-CM para TopoGPT3 / modelo espectral continuo.
+F2: CM phase-lattice STE suave. F3: Hodge regularizer + Fisher hinge.
+F4: MCMC offline. F5: Hecke pooling solo escala.
+"""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ---------- FASE 2: CM phase lattice ----------
+class CMPhaseQuantSTE(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, phase, m: int, beta: float):
+        step = 2 * math.pi / m
+        q = torch.round(phase / step) * step
+        return (1 - beta) * phase + beta * q
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        return grad_out, None, None
+
+
+def cm_phase_loss(kr: torch.Tensor, ki: torch.Tensor, m: int = 16):
+    """E_Q: distancia cuadratica media a raices m-esimas. kr,ki reales misma forma."""
+    phase = torch.atan2(ki, kr)
+    step = 2 * math.pi / m
+    q = torch.round(phase / step) * step
+    d = torch.angle(torch.exp(1j * (phase - q)).to(torch.complex64))
+    return (d ** 2).mean()
+
+
+def apply_cm_soft(kr, ki, m: int = 16, beta: float = 0.05):
+    """Cuantizacion suave de fase con STE. Magnitud intacta."""
+    with torch.no_grad():
+        mag = torch.sqrt(kr ** 2 + ki ** 2 + 1e-12)
+        phase = torch.atan2(ki, kr)
+    phase_q = CMPhaseQuantSTE.apply(phase, m, beta)
+    kr_q = mag * torch.cos(phase_q)
+    ki_q = mag * torch.sin(phase_q)
+    # STE: forward usa q, backward fluye a kr,ki originales
+    kr_out = kr + (kr_q - kr).detach()
+    ki_out = ki + (ki_q - ki).detach()
+    return kr_out, ki_out
+
+
+# ---------- FASE 3: Hodge + Fisher ----------
+def freq_grid_laplacian(fh: int, fw: int, device=None, dtype=torch.float32):
+    N = fh * fw
+    L = torch.zeros(N, N, device=device, dtype=dtype)
+
+    def idx(h, w):
+        return h * fw + w
+
+    for h in range(fh):
+        for w in range(fw):
+            i = idx(h, w)
+            for dh, dw in [(1, 0), (-1, 0), (0, 1), (0, -1)]:
+                j = idx((h + dh) % fh, (w + dw) % fw)
+                L[i, i] += 1
+                L[i, j] -= 1
+    return L
+
+
+def hodge_heat_residue(kr, ki, t: float = 2.0):
+    """||K - e^{-tL} K||_F^2 sobre grid freq. kr,ki: [...,Fh,Fw]."""
+    *lead, fh, fw = kr.shape
+    N = fh * fw
+    L = freq_grid_laplacian(fh, fw, device=kr.device, dtype=kr.dtype)
+    vals, vecs = torch.linalg.eigh(L)
+    heat = (vecs * torch.exp(-t * vals)) @ vecs.T  # [N,N]
+    out = 0.0
+    for x in (kr, ki):
+        xf = x.reshape(-1, N)
+        xs = (heat @ xf.T).T
+        out = out + ((xf - xs) ** 2).mean()
+    return out / 2.0
+
+
+def harmonic_ratio(kr, ki, t: float = 2.0):
+    *lead, fh, fw = kr.shape
+    N = fh * fw
+    L = freq_grid_laplacian(fh, fw, device=kr.device, dtype=kr.dtype)
+    vals, vecs = torch.linalg.eigh(L)
+    heat = (vecs * torch.exp(-t * vals)) @ vecs.T
+    num, den = 0.0, 0.0
+    for x in (kr, ki):
+        xf = x.reshape(-1, N)
+        xs = (heat @ xf.T).T
+        num += (xs ** 2).sum()
+        den += (xf ** 2).sum()
+    return (num / (den + 1e-12)).item()
+
+
+def fisher_hinge(singulars: torch.Tensor, r: int = 16, margin: float = 1e-3):
+    s = torch.sort(singulars, descending=True).values
+    if s.numel() <= r:
+        return torch.tensor(0.0, device=s.device)
+    gap = s[r - 1] - s[r]
+    return F.relu(margin - gap)
+
+
+# ---------- FASE 4: MCMC offline ----------
+def langevin_refine(param, energy_fn, steps: int = 200, eta: float = 1e-3,
+                    temp: float = 1e-4, seed: int = 0):
+    """ULA offline: theta <- theta - eta grad E + sqrt(2 eta T) xi."""
+    g = torch.Generator(device=param.device)
+    g.manual_seed(seed)
+    th = param.clone().float().detach()
+    for _ in range(steps):
+        th.requires_grad_(True)
+        e = energy_fn(th)
+        grad = torch.autograd.grad(e, th)[0].detach()
+        th = th.detach()
+        noise = torch.randn(th.shape, generator=g, device=th.device, dtype=th.dtype)
+        with torch.no_grad():
+            th = th - eta * grad + math.sqrt(2 * eta * temp) * noise
+    return th.detach()
+
+
+@torch.no_grad()
+def gibbs_cluster_assign(feats: torch.Tensor, n_clusters: int = 8, iters: int = 20, seed: int = 0):
+    """Clustering blando por fase para nodos/frecuencias. feats: [N,D] complejo o real."""
+    g = torch.Generator(device=feats.device)
+    g.manual_seed(seed)
+    if torch.is_complex(feats):
+        X = torch.view_as_real(feats.reshape(feats.shape[0], -1)).flatten(1)
+    else:
+        X = feats.flatten(1).float()
+    perm = torch.randperm(X.shape[0], generator=g, device=X.device)[:n_clusters]
+    centroids = X[perm].clone()
+    assign = torch.zeros(X.shape[0], dtype=torch.long, device=X.device)
+    for _ in range(iters):
+        d = ((X.unsqueeze(1) - centroids.unsqueeze(0)) ** 2).sum(-1)
+        # muestreo Gibbs aproximado: softmax sobre -d
+        probs = F.softmax(-d, dim=1)
+        assign = torch.multinomial(probs, 1, generator=g).squeeze(1)
+        for c in range(n_clusters):
+            m = assign == c
+            if m.any():
+                centroids[c] = X[m].mean(0)
+    return assign, centroids
+
+
+# ---------- FASE 5: Hecke pooling solo escala ----------
+class HeckeScalePool(nn.Module):
+    """Promedio sobre orbitas de escala/traslacion. Solo donde hay simetria real
+    (Strassen T, Fourier interp Willmore). NO reemplazo de atencion general."""
+
+    def __init__(self, modes: int = 4):
+        super().__init__()
+        self.modes = modes
+
+    def forward(self, Kfreq: torch.Tensor):
+        # Kfreq: [...,Fh,Fw] complejo. Promedia sobre shifts ciclicos (orbita traslacion).
+        outs = [torch.roll(Kfreq, shifts=(dh, dw), dims=(-2, -1))
+                for dh in range(min(2, Kfreq.shape[-2]))
+                for dw in range(min(2, Kfreq.shape[-1]))]
+        return sum(outs) / len(outs)