diff --git a/crates/ruvector-decompiler/Cargo.toml b/crates/ruvector-decompiler/Cargo.toml index 1b6b9a25c..273ceb23a 100644 --- a/crates/ruvector-decompiler/Cargo.toml +++ b/crates/ruvector-decompiler/Cargo.toml @@ -20,12 +20,14 @@ thiserror = { workspace = true } once_cell = "1" rayon = { workspace = true } memchr = "2" +ort = { version = "=2.0.0-rc.10", optional = true, default-features = false, features = ["ndarray", "std"] } +ndarray = { version = "0.16", optional = true } [features] default = [] -# Enable neural name inference using a trained GGUF model. -# Adds ~2MB to binary size for model loading and validation. -neural = [] +# Enable neural name inference using ONNX Runtime (via `ort` crate) +# or a GGUF/RVF model file. Adds model loading + inference capability. +neural = ["ort", "ndarray"] [dev-dependencies] criterion = { version = "0.5", features = ["html_reports"] } diff --git a/crates/ruvector-decompiler/src/inferrer.rs b/crates/ruvector-decompiler/src/inferrer.rs index 9dbab08a7..666f58b5d 100644 --- a/crates/ruvector-decompiler/src/inferrer.rs +++ b/crates/ruvector-decompiler/src/inferrer.rs @@ -7,9 +7,6 @@ //! 4. Property correlation //! 5. Structural heuristics -#[cfg(feature = "neural")] -use std::path::{Path, PathBuf}; - use crate::training::TrainingCorpus; use crate::types::{Declaration, InferredName, Module}; @@ -122,7 +119,7 @@ pub fn infer_names_with_corpus( /// 3. Property access correlation (MEDIUM confidence) /// 4. Multiple string literal heuristic (MEDIUM confidence) /// 5. Structural heuristics (LOW confidence) -fn infer_declaration_name( +pub(crate) fn infer_declaration_name( decl: &Declaration, corpus: &TrainingCorpus, ) -> Option { @@ -302,7 +299,7 @@ pub struct LearnedPattern { } // --------------------------------------------------------------------------- -// Neural name inference (behind `neural` feature flag) +// Neural name inference context (shared with `neural` module) // --------------------------------------------------------------------------- /// Context signals passed to the neural inferrer for a single declaration. @@ -327,136 +324,6 @@ impl InferenceContext { } } -/// Neural name inference using a trained deobfuscation model. -/// -/// Falls back to pattern-based inference if the model is not available. -/// Only compiled when the `neural` feature is enabled. -#[cfg(feature = "neural")] -pub struct NeuralInferrer { - /// Path to the GGUF model file (or RVF with OVERLAY). - model_path: PathBuf, - /// Whether the model is loaded and ready for inference. - active: bool, - // In a full implementation, the loaded GGUF weights and inference - // runtime would be stored here. For now we keep the structure ready - // for RuvLLM integration. -} - -#[cfg(feature = "neural")] -impl NeuralInferrer { - /// Attempt to load a neural deobfuscation model from a GGUF or RVF file. - /// - /// Returns `Ok(Self)` if the file exists and appears valid; inference - /// may still fall back to `None` if the runtime is not compiled in. - pub fn load(path: &Path) -> Result { - if !path.exists() { - return Err(crate::error::DecompilerError::ModelError(format!( - "model file not found: {}", - path.display() - ))); - } - - // Validate magic bytes: GGUF (0x46475547) or RVF (0x52564601). - let data = std::fs::read(path).map_err(|e| { - crate::error::DecompilerError::ModelError(format!( - "failed to read model file: {}", - e - )) - })?; - - if data.len() < 4 { - return Err(crate::error::DecompilerError::ModelError( - "model file too small".to_string(), - )); - } - - let magic = u32::from_le_bytes([data[0], data[1], data[2], data[3]]); - let is_gguf = magic == 0x46475547; - let is_rvf = &data[..4] == b"RVF\x01"; - - if !is_gguf && !is_rvf { - return Err(crate::error::DecompilerError::ModelError( - "unrecognized model format (expected GGUF or RVF)".to_string(), - )); - } - - Ok(Self { - model_path: path.to_path_buf(), - active: true, - }) - } - - /// Predict the original name for a minified identifier using the - /// neural model. - /// - /// Returns `None` if the model is not active or confidence is too low. - pub fn predict_name( - &self, - minified: &str, - context: &InferenceContext, - ) -> Option { - if !self.active { - return None; - } - - // TODO: integrate with RuvLLM GGUF runtime for real inference. - // For now, return None so the pipeline falls through to - // pattern-based strategies. This stub ensures the API is - // stable and the integration points are well-defined. - let _ = (minified, context); - None - } - - /// Whether the neural model is loaded and ready. - pub fn is_active(&self) -> bool { - self.active - } - - /// Path to the loaded model file. - pub fn model_path(&self) -> &Path { - &self.model_path - } -} - -/// Infer names with optional neural model support. -/// -/// When the `neural` feature is enabled and a model path is provided, -/// neural inference is attempted first for each declaration. Results -/// with confidence > 0.8 are accepted directly; otherwise the pipeline -/// falls through to corpus-based and heuristic strategies. -#[cfg(feature = "neural")] -pub fn infer_names_neural( - modules: &[Module], - model_path: Option<&Path>, -) -> Vec { - let corpus = TrainingCorpus::builtin(); - let neural = model_path.and_then(|p| NeuralInferrer::load(p).ok()); - - let mut inferred = Vec::new(); - - for module in modules { - for decl in &module.declarations { - // 1. Try neural inference (highest accuracy). - if let Some(ref model) = neural { - let ctx = InferenceContext::from_declaration(decl); - if let Some(name) = model.predict_name(&decl.name, &ctx) { - if name.confidence > 0.8 { - inferred.push(name); - continue; - } - } - } - - // 2. Fall back to corpus + heuristic strategies. - if let Some(inf) = infer_declaration_name(decl, &corpus) { - inferred.push(inf); - } - } - } - - inferred -} - #[cfg(test)] mod tests { use super::*; diff --git a/crates/ruvector-decompiler/src/lib.rs b/crates/ruvector-decompiler/src/lib.rs index 1ba1607b0..2d4255629 100644 --- a/crates/ruvector-decompiler/src/lib.rs +++ b/crates/ruvector-decompiler/src/lib.rs @@ -30,6 +30,8 @@ pub mod beautifier; pub mod error; pub mod graph; pub mod inferrer; +#[cfg(feature = "neural")] +pub mod neural; pub mod parser; pub mod partitioner; pub mod sourcemap; diff --git a/crates/ruvector-decompiler/src/neural.rs b/crates/ruvector-decompiler/src/neural.rs new file mode 100644 index 000000000..f20b104c2 --- /dev/null +++ b/crates/ruvector-decompiler/src/neural.rs @@ -0,0 +1,246 @@ +//! Neural name inference via ONNX Runtime (behind `neural` feature flag). +//! +//! Loads a trained deobfuscation model in ONNX, GGUF, or RVF format and +//! predicts human-readable names for minified JS identifiers. + +use std::path::{Path, PathBuf}; + +use crate::inferrer::{infer_declaration_name, InferenceContext}; +use crate::training::TrainingCorpus; +use crate::types::{InferredName, Module}; + +/// Neural name inference using a trained deobfuscation model. +/// +/// When an ONNX model is loaded, inference runs through ONNX Runtime. +/// GGUF and RVF formats are validated but inference is a stub pending +/// RuvLLM integration. +pub struct NeuralInferrer { + model_path: PathBuf, + /// Uses `RefCell` so `predict_name` can keep `&self` for the caller. + session: Option>, + active: bool, +} + +impl NeuralInferrer { + const MAX_CONTEXT_LEN: usize = 256; + const MAX_NAME_LEN: usize = 32; + const MAX_OUTPUT_LEN: usize = 64; + + /// Load a deobfuscation model from `path`. + /// + /// Supports `.onnx` (ONNX Runtime), GGUF (`0x46475547`), and + /// RVF (`RVF\x01`) formats. + pub fn load(path: &Path) -> Result { + if !path.exists() { + return Err(crate::error::DecompilerError::ModelError(format!( + "model file not found: {}", + path.display() + ))); + } + + let is_onnx = path + .extension() + .map_or(false, |ext| ext.eq_ignore_ascii_case("onnx")); + + if is_onnx { + return Self::load_onnx(path); + } + + Self::load_legacy(path) + } + + fn load_onnx(path: &Path) -> Result { + let session = ort::session::Session::builder() + .and_then(|b| b.commit_from_file(path)) + .map_err(|e| { + crate::error::DecompilerError::ModelError(format!( + "failed to load ONNX model: {e}" + )) + })?; + + Ok(Self { + model_path: path.to_path_buf(), + session: Some(std::cell::RefCell::new(session)), + active: true, + }) + } + + fn load_legacy(path: &Path) -> Result { + let data = std::fs::read(path).map_err(|e| { + crate::error::DecompilerError::ModelError(format!( + "failed to read model file: {e}" + )) + })?; + + if data.len() < 4 { + return Err(crate::error::DecompilerError::ModelError( + "model file too small".to_string(), + )); + } + + let magic = u32::from_le_bytes([data[0], data[1], data[2], data[3]]); + let is_gguf = magic == 0x46475547; + let is_rvf = &data[..4] == b"RVF\x01"; + + if !is_gguf && !is_rvf { + return Err(crate::error::DecompilerError::ModelError( + "unrecognized model format (expected ONNX, GGUF, or RVF)".to_string(), + )); + } + + Ok(Self { + model_path: path.to_path_buf(), + session: None, + active: true, + }) + } + + /// Predict the original name for a minified identifier. + pub fn predict_name( + &self, + minified: &str, + context: &InferenceContext, + ) -> Option { + if !self.active { + return None; + } + + let cell = self.session.as_ref()?; + let mut session = cell.borrow_mut(); + Self::run_onnx_inference(&mut session, minified, context) + } + + fn run_onnx_inference( + session: &mut ort::session::Session, + minified: &str, + context: &InferenceContext, + ) -> Option { + use ort::value::Tensor; + + let name_bytes: Vec = minified + .bytes() + .take(Self::MAX_NAME_LEN) + .map(|b| b as f32) + .chain(std::iter::repeat(0.0f32)) + .take(Self::MAX_NAME_LEN) + .collect(); + + let ctx_joined = [ + context.kind.as_str(), + " ", + &context.string_literals.join(" "), + " ", + &context.property_accesses.join(" "), + ] + .concat(); + let ctx_bytes: Vec = ctx_joined + .bytes() + .take(Self::MAX_CONTEXT_LEN) + .map(|b| b as f32) + .chain(std::iter::repeat(0.0f32)) + .take(Self::MAX_CONTEXT_LEN) + .collect(); + + let name_tensor = Tensor::from_array(( + vec![1i64, Self::MAX_NAME_LEN as i64], + name_bytes, + )) + .ok()?; + let ctx_tensor = Tensor::from_array(( + vec![1i64, Self::MAX_CONTEXT_LEN as i64], + ctx_bytes, + )) + .ok()?; + + let outputs = session + .run(ort::inputs![name_tensor, ctx_tensor]) + .ok()?; + + if outputs.len() < 2 { + return None; + } + + let (_shape, out_data) = outputs[0] + .try_extract_tensor::() + .ok()?; + let (_cshape, conf_data) = outputs[1] + .try_extract_tensor::() + .ok()?; + + let confidence = *conf_data.first()? as f64; + if confidence < 0.5 { + return None; + } + + let decoded: String = out_data + .iter() + .take(Self::MAX_OUTPUT_LEN) + .map(|&v| v.round() as u8) + .take_while(|&b| b > 0) + .filter(|b| b.is_ascii_alphanumeric() || *b == b'_') + .map(|b| b as char) + .collect(); + + if decoded.is_empty() { + return None; + } + + Some(InferredName { + original: minified.to_string(), + inferred: decoded, + confidence, + evidence: vec![format!( + "neural model prediction (ONNX, confidence: {confidence:.3})" + )], + }) + } + + /// Whether the neural model is loaded and ready. + pub fn is_active(&self) -> bool { + self.active + } + + /// Path to the loaded model file. + pub fn model_path(&self) -> &Path { + &self.model_path + } + + /// Whether the inferrer has a live ONNX session for real inference. + pub fn has_onnx_session(&self) -> bool { + self.session.is_some() + } +} + +/// Infer names with optional neural model support. +/// +/// Neural inference is attempted first; results with confidence > 0.8 +/// are accepted directly. Otherwise falls through to corpus + heuristics. +pub fn infer_names_neural( + modules: &[Module], + model_path: Option<&Path>, +) -> Vec { + let corpus = TrainingCorpus::builtin(); + let neural = model_path.and_then(|p| NeuralInferrer::load(p).ok()); + + let mut inferred = Vec::new(); + + for module in modules { + for decl in &module.declarations { + if let Some(ref model) = neural { + let ctx = InferenceContext::from_declaration(decl); + if let Some(name) = model.predict_name(&decl.name, &ctx) { + if name.confidence > 0.8 { + inferred.push(name); + continue; + } + } + } + + if let Some(inf) = infer_declaration_name(decl, &corpus) { + inferred.push(inf); + } + } + } + + inferred +} diff --git a/scripts/training/evaluate-model.py b/scripts/training/evaluate-model.py new file mode 100644 index 000000000..b69998115 --- /dev/null +++ b/scripts/training/evaluate-model.py @@ -0,0 +1,292 @@ +#!/usr/bin/env python3 +""" +Evaluate a trained deobfuscation model on held-out test data. + +Metrics: + - Exact match accuracy (full name correct) + - Prefix match accuracy (first N chars correct) + - Character-level accuracy (per-character correct rate) + - Top-K accuracy (correct answer in top K predictions) + +Usage: + python evaluate-model.py --model model-v2/best_model.pt --test training-data-v2.jsonl + python evaluate-model.py --model model-v2/best_model.pt --test training-data-v2.jsonl --split 0.1 +""" + +import argparse +import json +import os +import sys +from collections import Counter, defaultdict +from pathlib import Path + +import torch +import torch.nn as nn + +# Import model definition from training script +sys.path.insert(0, str(Path(__file__).parent)) +from importlib import import_module + +# Inline the constants and model to avoid import issues +VOCAB_SIZE = 256 +PAD_TOKEN = 0 +SOS_TOKEN = 1 +EOS_TOKEN = 2 +MAX_CONTEXT = 64 +MAX_NAME = 32 +EMBED_DIM = 128 +NUM_HEADS = 4 +NUM_LAYERS = 3 +FFN_DIM = 512 + + +class DeobfuscationModel(nn.Module): + """Mirror of the training model for inference.""" + + def __init__(self): + super().__init__() + total_seq = MAX_CONTEXT + MAX_NAME + self.max_context = MAX_CONTEXT + self.max_name = MAX_NAME + + self.char_embed = nn.Embedding(VOCAB_SIZE, EMBED_DIM, padding_idx=PAD_TOKEN) + self.pos_embed = nn.Embedding(total_seq, EMBED_DIM) + + encoder_layer = nn.TransformerEncoderLayer( + d_model=EMBED_DIM, + nhead=NUM_HEADS, + dim_feedforward=FFN_DIM, + batch_first=True, + dropout=0.1, + activation="gelu", + ) + self.encoder = nn.TransformerEncoder(encoder_layer, NUM_LAYERS) + self.layer_norm = nn.LayerNorm(EMBED_DIM) + self.output = nn.Linear(EMBED_DIM, VOCAB_SIZE) + + def forward(self, input_tokens): + batch_size, seq_len = input_tokens.shape + device = input_tokens.device + positions = torch.arange(seq_len, device=device).unsqueeze(0).expand(batch_size, -1) + x = self.char_embed(input_tokens) + self.pos_embed(positions) + pad_mask = input_tokens == PAD_TOKEN + x = self.encoder(x, src_key_padding_mask=pad_mask) + x = self.layer_norm(x) + name_out = x[:, -self.max_name:, :] + logits = self.output(name_out) + return logits + + +def encode(text, max_len): + """Encode text to byte-level tensor.""" + encoded = [min(b, VOCAB_SIZE - 1) for b in text.encode("utf-8")[:max_len]] + padded = encoded + [PAD_TOKEN] * (max_len - len(encoded)) + return torch.tensor(padded, dtype=torch.long) + + +def encode_target(text, max_len): + """Encode target with SOS/EOS.""" + encoded = [min(b, VOCAB_SIZE - 1) for b in text.encode("utf-8")[:max_len - 2]] + tokens = [SOS_TOKEN] + encoded + [EOS_TOKEN] + padded = tokens + [PAD_TOKEN] * (max_len - len(tokens)) + return torch.tensor(padded, dtype=torch.long) + + +def decode_tokens(tokens): + """Decode byte-level tokens back to string.""" + chars = [] + for t in tokens: + t = t.item() if hasattr(t, "item") else t + if t == PAD_TOKEN or t == EOS_TOKEN: + break + if t == SOS_TOKEN: + continue + if 32 <= t < 127: + chars.append(chr(t)) + return "".join(chars) + + +def prepare_input(sample): + """Prepare a single sample for inference.""" + minified = sample["minified"] + context_strings = sample.get("context_strings", []) + properties = sample.get("properties", []) + + context_text = " ".join(context_strings[:8]) + " | " + " ".join(properties[:8]) + context_tokens = encode(context_text, MAX_CONTEXT) + minified_tokens = encode(minified, MAX_NAME) + input_tokens = torch.cat([context_tokens, minified_tokens]) + return input_tokens + + +def evaluate(model_path, test_path, split_ratio=0.1, top_k=5, device_str="cpu"): + """Run evaluation.""" + device = torch.device(device_str) + + # Load model + print(f"Loading model from {model_path}") + model = DeobfuscationModel().to(device) + + checkpoint = torch.load(model_path, map_location=device, weights_only=False) + if "model_state_dict" in checkpoint: + model.load_state_dict(checkpoint["model_state_dict"]) + print(f" Checkpoint epoch: {checkpoint.get('epoch', '?')}") + print(f" Checkpoint val_loss: {checkpoint.get('val_loss', '?'):.4f}") + print(f" Checkpoint val_acc: {checkpoint.get('val_acc', '?'):.4f}") + else: + model.load_state_dict(checkpoint) + + model.eval() + param_count = sum(p.numel() for p in model.parameters()) + print(f" Model parameters: {param_count:,}") + + # Load test data + print(f"\nLoading test data from {test_path}") + samples = [] + with open(test_path) as f: + for line in f: + line = line.strip() + if line: + samples.append(json.loads(line)) + + # Use last N% as test set (same split as training) + total = len(samples) + test_size = max(100, int(total * split_ratio)) + test_samples = samples[-test_size:] + print(f" Total samples: {total}, test samples: {test_size}") + + # Evaluate + exact_match = 0 + prefix_3_match = 0 + prefix_5_match = 0 + char_correct = 0 + char_total = 0 + top_k_match = 0 + + kind_correct = defaultdict(int) + kind_total = defaultdict(int) + failures = [] + + with torch.no_grad(): + for i, sample in enumerate(test_samples): + original = sample["original"] + kind = sample.get("kind", "var") + + input_tokens = prepare_input(sample).unsqueeze(0).to(device) + logits = model(input_tokens) # (1, MAX_NAME, VOCAB_SIZE) + + # Greedy decode + preds = logits.argmax(dim=-1)[0] # (MAX_NAME,) + predicted = decode_tokens(preds) + + # Exact match + kind_total[kind] += 1 + if predicted == original: + exact_match += 1 + kind_correct[kind] += 1 + else: + failures.append({ + "minified": sample["minified"], + "original": original, + "predicted": predicted, + "kind": kind, + "context": sample.get("context_strings", [])[:3], + }) + + # Prefix matches + if predicted[:3] == original[:3]: + prefix_3_match += 1 + if predicted[:5] == original[:5]: + prefix_5_match += 1 + + # Character-level accuracy + for j in range(min(len(predicted), len(original))): + if predicted[j] == original[j]: + char_correct += 1 + char_total += 1 + # Count missing or extra chars as errors + char_total += abs(len(predicted) - len(original)) + + # Top-K: check if correct name is in top K at each position + target_tokens = encode_target(original, MAX_NAME) + match_in_topk = True + for j in range(min(len(original) + 2, MAX_NAME)): + if target_tokens[j] == PAD_TOKEN: + break + topk_vals = torch.topk(logits[0, j], top_k).indices + if target_tokens[j] not in topk_vals: + match_in_topk = False + break + if match_in_topk: + top_k_match += 1 + + # Print results + n = len(test_samples) + print("\n" + "=" * 60) + print("EVALUATION RESULTS") + print("=" * 60) + print(f" Test samples: {n}") + print(f" Exact match: {exact_match}/{n} = {100*exact_match/n:.1f}%") + print(f" Prefix-3 match: {prefix_3_match}/{n} = {100*prefix_3_match/n:.1f}%") + print(f" Prefix-5 match: {prefix_5_match}/{n} = {100*prefix_5_match/n:.1f}%") + print(f" Char-level accuracy: {100*char_correct/max(char_total,1):.1f}%") + print(f" Top-{top_k} accuracy: {top_k_match}/{n} = {100*top_k_match/n:.1f}%") + + print("\nAccuracy by kind:") + for kind in sorted(kind_total.keys()): + total_k = kind_total[kind] + correct_k = kind_correct[kind] + print(f" {kind:10s}: {correct_k}/{total_k} = {100*correct_k/max(total_k,1):.1f}%") + + # Show some failures + print(f"\nSample failures (showing first 15):") + for f in failures[:15]: + ctx_str = ", ".join(f["context"][:3]) if f["context"] else "none" + print(f" {f['minified']:8s} -> predicted '{f['predicted']:20s}' " + f"expected '{f['original']:20s}' [{f['kind']}] ctx=[{ctx_str}]") + + # Failure analysis + print("\nFailure analysis:") + len_buckets = defaultdict(lambda: [0, 0]) # [correct, total] + for sample in test_samples: + orig_len = len(sample["original"]) + bucket = "short(1-5)" if orig_len <= 5 else "medium(6-12)" if orig_len <= 12 else "long(13+)" + input_tokens = prepare_input(sample).unsqueeze(0).to(device) + with torch.no_grad(): + logits = model(input_tokens) + predicted = decode_tokens(logits.argmax(dim=-1)[0]) + len_buckets[bucket][1] += 1 + if predicted == sample["original"]: + len_buckets[bucket][0] += 1 + + for bucket in ["short(1-5)", "medium(6-12)", "long(13+)"]: + if bucket in len_buckets: + c, t = len_buckets[bucket] + print(f" {bucket:15s}: {c}/{t} = {100*c/max(t,1):.1f}%") + + return exact_match / n + + +def main(): + parser = argparse.ArgumentParser(description="Evaluate deobfuscation model") + parser.add_argument("--model", required=True, help="Path to model checkpoint (.pt)") + parser.add_argument("--test", required=True, help="Path to test data JSONL") + parser.add_argument("--split", type=float, default=0.1, help="Test split ratio") + parser.add_argument("--top-k", type=int, default=5, help="Top-K accuracy K") + parser.add_argument("--device", default="cpu", help="Device") + args = parser.parse_args() + + accuracy = evaluate( + model_path=args.model, + test_path=args.test, + split_ratio=args.split, + top_k=args.top_k, + device_str=args.device, + ) + + # Exit code: 0 if accuracy >= 80%, 1 otherwise + sys.exit(0 if accuracy >= 0.80 else 1) + + +if __name__ == "__main__": + main() diff --git a/scripts/training/filter-and-augment.py b/scripts/training/filter-and-augment.py new file mode 100644 index 000000000..d306891f6 --- /dev/null +++ b/scripts/training/filter-and-augment.py @@ -0,0 +1,233 @@ +#!/usr/bin/env python3 +""" +Filter training data for quality and add augmentation. +Aims for ~30-50K high-quality, diverse pairs that train well on CPU. + +Steps: + 1. Deduplicate by (original, context_hash) to remove near-duplicates + 2. Filter out low-quality pairs (no context, too-short names) + 3. Balance by kind (function/class/var) + 4. Augment with context shuffling, partial context, case variants + 5. Output filtered+augmented JSONL +""" + +import json +import hashlib +import random +import sys +from collections import defaultdict + +INPUT = sys.argv[1] if len(sys.argv) > 1 else "training-data-v2.jsonl" +OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "training-data-v2-filtered.jsonl" +TARGET_SIZE = 40000 + +random.seed(42) + +# Load all pairs +print(f"Loading {INPUT}...") +pairs = [] +with open(INPUT) as f: + for line in f: + line = line.strip() + if not line: + continue + pairs.append(json.loads(line)) + +print(f"Loaded {len(pairs)} pairs") + +# Step 1: Quality filter +quality_pairs = [] +for p in pairs: + original = p["original"] + ctx = p.get("context_strings", []) + + # Skip very short original names (not useful for learning) + if len(original) < 3: + continue + + # Skip if no context at all + if len(ctx) == 0: + continue + + # Skip names that are likely not real identifiers + if not original[0].isalpha() and original[0] not in ("_", "$"): + continue + + # Skip all-uppercase (likely constants, less interesting) + if original.isupper() and len(original) > 5: + continue + + quality_pairs.append(p) + +print(f"After quality filter: {len(quality_pairs)}") + +# Step 2: Deduplicate by original name - keep max 8 variants per name +by_original = defaultdict(list) +for p in quality_pairs: + by_original[p["original"]].append(p) + +deduped = [] +for original, variants in by_original.items(): + # Keep up to 8 most diverse variants (by context) + random.shuffle(variants) + deduped.extend(variants[:8]) + +print(f"After dedup (max 8 per name): {len(deduped)}") + +# Step 3: Balance by kind +by_kind = defaultdict(list) +for p in deduped: + by_kind[p["kind"]].append(p) + +print("By kind before balancing:") +for k, v in sorted(by_kind.items()): + print(f" {k}: {len(v)}") + +# Cap each kind to prevent overwhelming dominance +max_per_kind = TARGET_SIZE // 2 # Allow some imbalance +balanced = [] +for kind, items in by_kind.items(): + random.shuffle(items) + balanced.extend(items[:max_per_kind]) + +random.shuffle(balanced) +print(f"After balancing: {len(balanced)}") + +# Step 4: Augmentation +augmented = list(balanced) + + +def shuffle_context(p): + """Shuffle context strings order.""" + ctx = list(p["context_strings"]) + random.shuffle(ctx) + return {**p, "context_strings": ctx, + "minified": random_minified()} + + +def partial_context(p): + """Drop some context strings.""" + ctx = p["context_strings"] + if len(ctx) <= 1: + return None + # Keep 50-80% of context + keep = max(1, int(len(ctx) * random.uniform(0.5, 0.8))) + new_ctx = random.sample(ctx, keep) + return {**p, "context_strings": new_ctx, + "minified": random_minified()} + + +def case_variant(p): + """Generate case variant of the original name.""" + original = p["original"] + variants = [] + + # camelCase -> snake_case + snake = "" + for i, c in enumerate(original): + if c.isupper() and i > 0: + snake += "_" + c.lower() + else: + snake += c.lower() + if snake != original and len(snake) > 3: + variants.append(snake) + + # camelCase -> PascalCase + pascal = original[0].upper() + original[1:] + if pascal != original: + variants.append(pascal) + + if not variants: + return None + + chosen = random.choice(variants) + return {**p, "original": chosen, "minified": random_minified()} + + +MINIFIED_CHARS = "abcdefghijklmnopqrstuvwxyz" + +def random_minified(): + style = random.randint(0, 7) + i = random.randint(0, 200) + if style == 0: + return MINIFIED_CHARS[i % 26] + elif style == 1: + return MINIFIED_CHARS[i % 26] + str(i % 10) + elif style == 2: + return "_" + MINIFIED_CHARS[i % 26] + elif style == 3: + return "_0x" + hex(0x1a2b + i)[2:] + elif style == 4: + return "$" + MINIFIED_CHARS[i % 26] + elif style == 5: + return "t" + str(i) + elif style == 6: + a = MINIFIED_CHARS[i % 26] + b = MINIFIED_CHARS[(i + 1) % 26] + return a + b + else: + return "n" + str(i) + + +# Generate augmented pairs +aug_count = 0 +for p in balanced: + # 30% chance of context shuffle augmentation + if random.random() < 0.3: + aug = shuffle_context(p) + augmented.append(aug) + aug_count += 1 + + # 20% chance of partial context + if random.random() < 0.2: + aug = partial_context(p) + if aug: + augmented.append(aug) + aug_count += 1 + + # 10% chance of case variant + if random.random() < 0.1: + aug = case_variant(p) + if aug: + augmented.append(aug) + aug_count += 1 + +print(f"Augmented pairs added: {aug_count}") + +# Final dedup +seen = set() +final = [] +for p in augmented: + key = f"{p['minified']}|{p['original']}" + if key not in seen: + seen.add(key) + final.append(p) + +random.shuffle(final) + +# Trim to target size if too large +if len(final) > TARGET_SIZE * 1.5: + final = final[:int(TARGET_SIZE * 1.5)] + +print(f"\nFinal dataset: {len(final)} pairs") + +# Write output +with open(OUTPUT, "w") as f: + for p in final: + f.write(json.dumps(p) + "\n") + +print(f"Wrote to {OUTPUT}") + +# Stats +kinds = defaultdict(int) +for p in final: + kinds[p["kind"]] += 1 + +print("\nFinal breakdown:") +for k, v in sorted(kinds.items()): + print(f" {k}: {v}") + +avg_ctx = sum(len(p["context_strings"]) for p in final) / len(final) +avg_orig_len = sum(len(p["original"]) for p in final) / len(final) +print(f"Avg context strings: {avg_ctx:.1f}") +print(f"Avg original name length: {avg_orig_len:.1f}") diff --git a/scripts/training/generate-data-v2.mjs b/scripts/training/generate-data-v2.mjs new file mode 100644 index 000000000..e1b45a613 --- /dev/null +++ b/scripts/training/generate-data-v2.mjs @@ -0,0 +1,571 @@ +#!/usr/bin/env node +/** + * Generate expanded training data for JS deobfuscation model (v2). + * + * Sources: + * 1. Existing training-data.jsonl (merge) + * 2. Real JS files from node_modules (identifier extraction) + * 3. Synthetic augmentation with context diversity + * + * Targets 15,000+ unique pairs for SOTA training. + * + * Usage: + * node scripts/training/generate-data-v2.mjs [--output training-data-v2.jsonl] + */ + +import { readFileSync, writeFileSync, readdirSync, statSync, existsSync } from "fs"; +import { join, resolve, basename } from "path"; +import { parseArgs } from "util"; + +const { values: args } = parseArgs({ + options: { + output: { type: "string", default: "training-data-v2.jsonl" }, + help: { type: "boolean", short: "h", default: false }, + }, +}); + +if (args.help) { + console.log("Usage: generate-data-v2.mjs [--output FILE]"); + process.exit(0); +} + +const OUTPUT_PATH = resolve(args.output); +const ROOT = resolve(import.meta.dirname, "../.."); + +/** @type {Map} key -> pair object, for dedup */ +const pairMap = new Map(); + +function addPair(minified, original, contextStrings, properties, kind) { + if (!minified || !original || original.length <= 1) return; + // Skip if original looks minified itself + if (original.length <= 2 && !/^[A-Z]/.test(original)) return; + const key = `${minified}|${original}`; + if (pairMap.has(key)) return; + pairMap.set(key, { + minified, + original, + context_strings: contextStrings.slice(0, 8), + properties: properties.slice(0, 8), + kind, + }); +} + +// --------------------------------------------------------------------------- +// Source 1: Merge existing training data +// --------------------------------------------------------------------------- + +function mergeExisting() { + const existingPath = join(ROOT, "training-data.jsonl"); + if (!existsSync(existingPath)) { + console.log(" [existing] no training-data.jsonl found, skipping"); + return 0; + } + const lines = readFileSync(existingPath, "utf8").trim().split("\n"); + let count = 0; + for (const line of lines) { + if (!line.trim()) continue; + try { + const obj = JSON.parse(line); + addPair( + obj.minified, + obj.original, + obj.context_strings || [], + obj.properties || [], + obj.kind || "var" + ); + count++; + } catch { /* skip bad lines */ } + } + console.log(` [existing] merged ${count} pairs`); + return count; +} + +// --------------------------------------------------------------------------- +// Source 2: Extract identifiers from real JS files in node_modules +// --------------------------------------------------------------------------- + +/** Walk directory tree, collect .js files up to maxDepth */ +function collectJsFiles(dir, maxDepth = 3, depth = 0) { + const files = []; + if (depth > maxDepth) return files; + let entries; + try { entries = readdirSync(dir); } catch { return files; } + for (const entry of entries) { + if (entry === "node_modules" && depth > 0) continue; + if (entry.startsWith(".")) continue; + const full = join(dir, entry); + let stat; + try { stat = statSync(full); } catch { continue; } + if (stat.isDirectory()) { + files.push(...collectJsFiles(full, maxDepth, depth + 1)); + } else if (entry.endsWith(".js") && stat.size > 1000 && stat.size < 200000) { + files.push(full); + } + } + return files; +} + +/** + * Extract identifiers from a JS source file using regex patterns. + * Returns array of { name, kind, nearbyTokens } + */ +function extractIdentifiers(source) { + const results = []; + const seen = new Set(); + + // Pattern: function declarations + const funcDeclRe = /\bfunction\s+([a-zA-Z_$][a-zA-Z0-9_$]{2,})\s*\(/g; + let m; + while ((m = funcDeclRe.exec(source)) !== null) { + if (!seen.has(m[1])) { + seen.add(m[1]); + const ctx = extractNearbyContext(source, m.index, 200); + results.push({ name: m[1], kind: "function", ctx }); + } + } + + // Pattern: const/let/var declarations with meaningful names + const varDeclRe = /\b(?:const|let|var)\s+([a-zA-Z_$][a-zA-Z0-9_$]{2,})\s*=/g; + while ((m = varDeclRe.exec(source)) !== null) { + if (!seen.has(m[1])) { + seen.add(m[1]); + const ctx = extractNearbyContext(source, m.index, 200); + results.push({ name: m[1], kind: "var", ctx }); + } + } + + // Pattern: class declarations + const classDeclRe = /\bclass\s+([a-zA-Z_$][a-zA-Z0-9_$]{2,})\b/g; + while ((m = classDeclRe.exec(source)) !== null) { + if (!seen.has(m[1])) { + seen.add(m[1]); + const ctx = extractNearbyContext(source, m.index, 200); + results.push({ name: m[1], kind: "class", ctx }); + } + } + + // Pattern: method definitions (object/class methods) + const methodRe = /\b([a-zA-Z_$][a-zA-Z0-9_$]{2,})\s*\([^)]*\)\s*\{/g; + while ((m = methodRe.exec(source)) !== null) { + const name = m[1]; + if (!seen.has(name) && !SKIP_NAMES.has(name)) { + seen.add(name); + const ctx = extractNearbyContext(source, m.index, 200); + results.push({ name, kind: "function", ctx }); + } + } + + // Pattern: exports.X = or module.exports.X = + const exportsRe = /(?:exports|module\.exports)\.([a-zA-Z_$][a-zA-Z0-9_$]{2,})\s*=/g; + while ((m = exportsRe.exec(source)) !== null) { + if (!seen.has(m[1])) { + seen.add(m[1]); + const ctx = extractNearbyContext(source, m.index, 200); + results.push({ name: m[1], kind: "var", ctx }); + } + } + + // Pattern: prototype methods + const protoRe = /\.prototype\.([a-zA-Z_$][a-zA-Z0-9_$]{2,})\s*=/g; + while ((m = protoRe.exec(source)) !== null) { + if (!seen.has(m[1])) { + seen.add(m[1]); + const ctx = extractNearbyContext(source, m.index, 200); + results.push({ name: m[1], kind: "function", ctx }); + } + } + + return results; +} + +const SKIP_NAMES = new Set([ + "if", "else", "for", "while", "do", "switch", "case", "break", + "continue", "return", "try", "catch", "finally", "throw", "new", + "delete", "typeof", "void", "instanceof", "in", "of", "with", + "this", "super", "true", "false", "null", "undefined", "NaN", + "Infinity", "arguments", "eval", "constructor", "prototype", + "use", "strict", "exports", "module", "require", +]); + +/** + * Extract nearby context tokens around a match position. + */ +function extractNearbyContext(source, pos, window) { + const start = Math.max(0, pos - window); + const end = Math.min(source.length, pos + window); + const snippet = source.slice(start, end); + + // Extract string literals as context + const strings = []; + const strRe = /["']([a-zA-Z][a-zA-Z0-9_.-]{2,})["']/g; + let m; + while ((m = strRe.exec(snippet)) !== null) { + if (!SKIP_NAMES.has(m[1]) && m[1].length < 30) { + strings.push(m[1]); + } + } + + // Extract property accesses as context + const propRe = /\.([a-zA-Z_$][a-zA-Z0-9_$]{2,})/g; + while ((m = propRe.exec(snippet)) !== null) { + if (!SKIP_NAMES.has(m[1]) && m[1].length < 25) { + strings.push(m[1]); + } + } + + // Deduplicate and limit + return [...new Set(strings)].slice(0, 10); +} + +/** + * Extract property accesses for a given identifier from source. + */ +function extractProperties(source, name) { + const props = new Set(); + // Look for name.property patterns + const re = new RegExp(`\\b${escapeRegex(name)}\\.([a-zA-Z_$][a-zA-Z0-9_$]{1,})`, "g"); + let m; + while ((m = re.exec(source)) !== null) { + if (m[1].length < 25) props.add(m[1]); + } + return [...props].slice(0, 8); +} + +function escapeRegex(s) { + return s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); +} + +// Minifier name generators +const MINIFIER_STYLES = [ + (i) => String.fromCharCode(97 + (i % 26)), + (i) => String.fromCharCode(97 + (i % 26)) + "$", + (i) => "_" + String.fromCharCode(97 + (i % 26)), + (i) => "_0x" + (0x1a2b + i).toString(16), + (i) => String.fromCharCode(97 + (i % 26)) + (i % 10).toString(), + (i) => "__" + String.fromCharCode(97 + (i % 26)), + (i) => "$" + String.fromCharCode(97 + (i % 26)), + (i) => String.fromCharCode(65 + (i % 26)), + (i) => { + const a = String.fromCharCode(97 + (i % 26)); + const b = String.fromCharCode(97 + ((i + 1) % 26)); + return a + b; + }, + (i) => "$" + (i % 100).toString(), + (i) => "_" + (i % 100).toString(), + (i) => "t" + i, + (i) => "e$" + String.fromCharCode(97 + (i % 26)), + (i) => "n" + (i % 100), + (i) => "r" + String.fromCharCode(97 + (i % 26)), + (i) => String.fromCharCode(97 + (i % 26)) + String.fromCharCode(97 + ((i * 7) % 26)), +]; + +function extractFromNodeModules() { + const nmDir = join(ROOT, "node_modules"); + if (!existsSync(nmDir)) { + console.log(" [node_modules] directory not found"); + return 0; + } + + const jsFiles = collectJsFiles(nmDir, 4); + console.log(` [node_modules] found ${jsFiles.length} JS files to scan`); + + let totalExtracted = 0; + let fileIdx = 0; + + for (const file of jsFiles) { + let source; + try { source = readFileSync(file, "utf8"); } catch { continue; } + + // Skip minified files (low ratio of newlines to content) + const lineCount = source.split("\n").length; + if (lineCount < 10 && source.length > 5000) continue; + + const identifiers = extractIdentifiers(source); + if (identifiers.length === 0) continue; + + for (let i = 0; i < identifiers.length; i++) { + const { name, kind, ctx } = identifiers[i]; + if (name.length < 3 || SKIP_NAMES.has(name)) continue; + + const properties = extractProperties(source, name); + + // Generate multiple minified variants per identifier + const numVariants = Math.min(4, MINIFIER_STYLES.length); + for (let v = 0; v < numVariants; v++) { + const styleIdx = (fileIdx + i + v) % MINIFIER_STYLES.length; + const minified = MINIFIER_STYLES[styleIdx](fileIdx + i); + + // Vary context slightly for each variant + const contextVariant = varySyntheticContext(ctx, v); + addPair(minified, name, contextVariant, properties, kind); + totalExtracted++; + } + } + fileIdx++; + } + + console.log(` [node_modules] extracted ${totalExtracted} pairs`); + return totalExtracted; +} + +// --------------------------------------------------------------------------- +// Source 3: Augmentation -- camelCase splitting + semantic context +// --------------------------------------------------------------------------- + +/** Split camelCase/PascalCase into tokens */ +function splitCamelCase(name) { + return name + .replace(/([A-Z])/g, " $1") + .trim() + .toLowerCase() + .split(/\s+/) + .filter((t) => t.length > 1); +} + +/** Generate semantic context from the name itself */ +function generateSemanticContext(name) { + const tokens = splitCamelCase(name); + const semantic = []; + + // Add the camelCase tokens as context hints + semantic.push(...tokens.slice(0, 4)); + + // Add type hints based on common prefixes/suffixes + if (/^is[A-Z]/.test(name)) semantic.push("boolean", "check"); + if (/^has[A-Z]/.test(name)) semantic.push("boolean", "exists"); + if (/^get[A-Z]/.test(name)) semantic.push("getter", "return"); + if (/^set[A-Z]/.test(name)) semantic.push("setter", "assign"); + if (/^on[A-Z]/.test(name)) semantic.push("event", "handler"); + if (/^handle[A-Z]/.test(name)) semantic.push("event", "callback"); + if (/^create[A-Z]/.test(name)) semantic.push("factory", "new"); + if (/^parse[A-Z]/.test(name)) semantic.push("parse", "input"); + if (/^format[A-Z]/.test(name)) semantic.push("format", "output"); + if (/^validate[A-Z]/.test(name)) semantic.push("validate", "check"); + if (/^render[A-Z]/.test(name)) semantic.push("render", "display"); + if (/^fetch[A-Z]/.test(name)) semantic.push("async", "request"); + if (/^load[A-Z]/.test(name)) semantic.push("async", "data"); + if (/^save[A-Z]/.test(name)) semantic.push("persist", "store"); + if (/^delete[A-Z]/.test(name)) semantic.push("remove", "destroy"); + if (/^update[A-Z]/.test(name)) semantic.push("modify", "change"); + if (/^init/.test(name)) semantic.push("initialize", "setup"); + if (/^process/.test(name)) semantic.push("transform", "pipeline"); + + // Suffix-based hints + if (/Error$/.test(name)) semantic.push("error", "exception"); + if (/Handler$/.test(name)) semantic.push("handler", "callback"); + if (/Manager$/.test(name)) semantic.push("manager", "lifecycle"); + if (/Service$/.test(name)) semantic.push("service", "business"); + if (/Controller$/.test(name)) semantic.push("controller", "http"); + if (/Factory$/.test(name)) semantic.push("factory", "create"); + if (/Builder$/.test(name)) semantic.push("builder", "construct"); + if (/Adapter$/.test(name)) semantic.push("adapter", "convert"); + if (/Provider$/.test(name)) semantic.push("provider", "inject"); + if (/Listener$/.test(name)) semantic.push("listener", "event"); + if (/Config$/.test(name)) semantic.push("config", "settings"); + if (/Options$/.test(name)) semantic.push("options", "settings"); + if (/Result$/.test(name)) semantic.push("result", "output"); + if (/Callback$/.test(name)) semantic.push("callback", "async"); + + return [...new Set(semantic)].slice(0, 8); +} + +/** + * Vary context slightly for training diversity. + */ +function varySyntheticContext(ctx, variant) { + if (!ctx || ctx.length === 0) return ["unknown"]; + switch (variant % 5) { + case 0: return ctx; + case 1: return ctx.length > 2 ? [...ctx.slice(1), ctx[0]] : ctx; + case 2: return ctx.slice(0, Math.max(2, Math.ceil(ctx.length / 2))); + case 3: return [...ctx, "prototype", "constructor"].slice(0, 8); + case 4: return [...ctx.slice(0, 3), "undefined", "null"].slice(0, 8); + default: return ctx; + } +} + +/** + * Generate augmented pairs by cross-version simulation. + */ +function generateCrossVersionAugmentation() { + const originals = new Map(); + for (const [, pair] of pairMap) { + if (!originals.has(pair.original)) { + originals.set(pair.original, pair); + } + } + + let augmented = 0; + const allOriginals = [...originals.entries()]; + + for (const [originalName, basePair] of allOriginals) { + // Generate 2-3 extra "version" variants + const versions = 2 + Math.floor(Math.random() * 2); + for (let v = 0; v < versions; v++) { + const minified = randomMinifiedName(); + const key = `${minified}|${originalName}`; + if (pairMap.has(key)) continue; + + // Vary context + const ctx = varySyntheticContext(basePair.context_strings, v); + addPair(minified, originalName, ctx, basePair.properties, basePair.kind); + augmented++; + } + } + + console.log(` [cross-version] augmented ${augmented} pairs`); + return augmented; +} + +function randomMinifiedName() { + const styles = [ + () => String.fromCharCode(97 + rand(26)) + rand(100), + () => "_0x" + rand(0xffff).toString(16), + () => String.fromCharCode(97 + rand(26)) + String.fromCharCode(97 + rand(26)), + () => "$" + String.fromCharCode(97 + rand(26)), + () => "t" + rand(200), + () => "n" + rand(100), + () => "_" + rand(200), + () => String.fromCharCode(97 + rand(26)) + String.fromCharCode(97 + rand(26)) + rand(10), + ]; + return styles[rand(styles.length)](); +} + +function rand(max) { return Math.floor(Math.random() * max); } + +// --------------------------------------------------------------------------- +// Source 4: Additional synthetic names for coverage +// --------------------------------------------------------------------------- + +function generateAdditionalSynthetic() { + // Common web/Node.js identifiers not likely in node_modules source + const EXTRA_NAMES = { + function: [ + // Webpack/bundler internals + "__webpack_require__", "__webpack_modules__", "__webpack_exports__", + // React internals + "createElement", "cloneElement", "createRef", "forwardRef", + "memo", "lazy", "Suspense", "Fragment", + "useId", "useSyncExternalStore", "useInsertionEffect", + // Next.js patterns + "getServerSideProps", "getStaticProps", "getStaticPaths", + "generateMetadata", "generateStaticParams", + // Express patterns + "createApplication", "createMiddleware", "createRoute", + "useRouter", "useParams", "useSearchParams", + // Testing + "beforeEach", "afterEach", "beforeAll", "afterAll", + "spyOn", "mockImplementation", "mockReturnValue", + // Utilities + "cloneDeep", "mergeWith", "assignIn", "defaultsDeep", + "flattenDeep", "uniqBy", "groupBy", "sortBy", "orderBy", + "pickBy", "omitBy", "mapKeys", "mapValues", + // Crypto/Security + "createHash", "createCipher", "createDecipher", "createSign", + "randomBytes", "scrypt", "pbkdf2", + // Stream + "createReadStream", "createWriteStream", "pipeline", "finished", + "Transform", "Readable", "Writable", "Duplex", "PassThrough", + ], + class: [ + "AbortController", "AbortSignal", "TextEncoder", "TextDecoder", + "URLSearchParams", "FormData", "Headers", "ReadableStream", + "WritableStream", "TransformStream", "BroadcastChannel", + "IntersectionObserver", "MutationObserver", "ResizeObserver", + "PerformanceObserver", "MessageChannel", "MessagePort", + "WeakRef", "FinalizationRegistry", "SharedArrayBuffer", + // Framework classes + "EventTarget", "CustomEvent", "DOMParser", "XMLSerializer", + "WebSocket", "Worker", "ServiceWorker", "SharedWorker", + ], + var: [ + // Common config keys + "baseURL", "timeout", "maxRedirects", "maxContentLength", + "validateStatus", "transformRequest", "transformResponse", + "paramsSerializer", "withCredentials", "responseEncoding", + // State patterns + "initialState", "rootReducer", "rootSaga", "rootEpic", + "storeEnhancers", "middlewares", "devTools", + // Build tools + "webpackConfig", "rollupConfig", "viteConfig", "babelConfig", + "tsConfig", "eslintConfig", "prettierConfig", + // Environment + "NODE_ENV", "API_URL", "BASE_PATH", "PUBLIC_URL", + ], + }; + + let count = 0; + for (const [kind, names] of Object.entries(EXTRA_NAMES)) { + for (let i = 0; i < names.length; i++) { + const original = names[i]; + const semanticCtx = generateSemanticContext(original); + const props = kind === "function" + ? ["length", "name", "call", "apply", "bind"] + : kind === "class" + ? ["prototype", "constructor", "name"] + : ["toString", "valueOf"]; + + // 4 minified variants per name + for (let v = 0; v < 4; v++) { + const styleIdx = (i + v) % MINIFIER_STYLES.length; + const minified = MINIFIER_STYLES[styleIdx](i); + const ctx = varySyntheticContext(semanticCtx, v); + addPair(minified, original, ctx, props, kind); + count++; + } + } + } + + console.log(` [extra-synthetic] generated ${count} pairs`); + return count; +} + +// --------------------------------------------------------------------------- +// Main +// --------------------------------------------------------------------------- + +console.log("=== Generating expanded training data (v2) ===\n"); + +console.log("Step 1: Merging existing training data"); +mergeExisting(); + +console.log("\nStep 2: Extracting identifiers from node_modules"); +extractFromNodeModules(); + +console.log("\nStep 3: Additional synthetic identifiers"); +generateAdditionalSynthetic(); + +console.log("\nStep 4: Cross-version augmentation"); +generateCrossVersionAugmentation(); + +// Convert to array and shuffle +const allPairs = [...pairMap.values()]; + +// Fisher-Yates shuffle +for (let i = allPairs.length - 1; i > 0; i--) { + const j = Math.floor(Math.random() * (i + 1)); + [allPairs[i], allPairs[j]] = [allPairs[j], allPairs[i]]; +} + +console.log(`\n=== Total unique pairs: ${allPairs.length} ===`); + +// Write JSONL +const lines = allPairs.map((p) => JSON.stringify(p)).join("\n"); +writeFileSync(OUTPUT_PATH, lines + "\n", "utf8"); +console.log(`Wrote ${allPairs.length} pairs to ${OUTPUT_PATH}`); + +// Print stats +const kindCounts = {}; +for (const p of allPairs) { + kindCounts[p.kind] = (kindCounts[p.kind] || 0) + 1; +} +console.log("\nBreakdown by kind:"); +for (const [kind, count] of Object.entries(kindCounts).sort((a, b) => b[1] - a[1])) { + console.log(` ${kind}: ${count}`); +} + +// Print average context length +const avgCtx = allPairs.reduce((s, p) => s + p.context_strings.length, 0) / allPairs.length; +const avgProps = allPairs.reduce((s, p) => s + p.properties.length, 0) / allPairs.length; +console.log(`\nAverage context strings per pair: ${avgCtx.toFixed(1)}`); +console.log(`Average properties per pair: ${avgProps.toFixed(1)}`); diff --git a/scripts/training/generate-deobfuscation-data.mjs b/scripts/training/generate-deobfuscation-data.mjs index 297600f7c..8b86f249e 100644 --- a/scripts/training/generate-deobfuscation-data.mjs +++ b/scripts/training/generate-deobfuscation-data.mjs @@ -18,6 +18,7 @@ import { readFileSync, writeFileSync, readdirSync, statSync } from "fs"; import { join, resolve, extname } from "path"; import { execSync } from "child_process"; import { parseArgs } from "util"; +import { COMMON_NAMES, CONTEXT_MAP, PROPERTY_MAP } from "./data/identifier-dictionaries.mjs"; // --------------------------------------------------------------------------- // CLI @@ -130,144 +131,82 @@ function extractGroundTruthFixtures() { * This simulates what real minifiers produce. */ function generateSyntheticPairs() { - const COMMON_NAMES = { - function: [ - "createElement", "appendChild", "removeChild", "setAttribute", - "addEventListener", "removeEventListener", "querySelector", "querySelectorAll", - "getElementById", "getElementsByClassName", "preventDefault", "stopPropagation", - "dispatch", "subscribe", "unsubscribe", "connect", "disconnect", - "initialize", "configure", "validate", "serialize", "deserialize", - "transform", "normalize", "sanitize", "encode", "decode", - "encrypt", "decrypt", "compress", "decompress", - "fetchData", "postData", "getData", "setData", "deleteData", - "handleClick", "handleSubmit", "handleChange", "handleError", - "createRouter", "createStore", "createContext", "createRef", - "useEffect", "useState", "useCallback", "useMemo", "useReducer", - "parseJSON", "stringifyJSON", "parseURL", "formatDate", - "sortArray", "filterItems", "mapValues", "reduceTotal", - "debounce", "throttle", "memoize", "curry", - "deepClone", "deepMerge", "deepEqual", "shallowEqual", - "getToken", "setToken", "clearToken", "refreshToken", - "openModal", "closeModal", "toggleMenu", "scrollToTop", - "sendRequest", "cancelRequest", "retryRequest", - "renderComponent", "mountComponent", "unmountComponent", - "logMessage", "logError", "logWarning", "logInfo", - "readFile", "writeFile", "deleteFile", "listFiles", - "startServer", "stopServer", "restartServer", - "connectDatabase", "queryDatabase", "closeConnection", - "hashPassword", "verifyPassword", "generateSalt", - "createSession", "destroySession", "getSession", - "emitEvent", "onEvent", "offEvent", "broadcastEvent", - "parseTemplate", "renderTemplate", "compileTemplate", - "formatCurrency", "formatNumber", "formatPercentage", - "calculateTotal", "calculateTax", "calculateDiscount", - "validateEmail", "validatePhone", "validatePassword", - "uploadFile", "downloadFile", "processFile", - ], - class: [ - "Component", "Controller", "Service", "Factory", "Repository", - "Manager", "Handler", "Builder", "Parser", "Formatter", - "Validator", "Serializer", "Transformer", "Adapter", "Wrapper", - "EventEmitter", "Observable", "Iterator", "Generator", - "HttpClient", "WebSocketClient", "DatabaseClient", - "UserService", "AuthService", "DataService", "CacheService", - "Router", "Middleware", "Pipeline", "Queue", "Stack", - "Logger", "Monitor", "Tracker", "Analyzer", - "Config", "Settings", "Options", "Preferences", - "Request", "Response", "Context", "Session", - "Model", "View", "Presenter", "ViewModel", - "Store", "State", "Reducer", "Action", - "Plugin", "Extension", "Module", "Package", - ], - var: [ - "config", "options", "settings", "preferences", "defaults", - "state", "props", "context", "params", "args", - "result", "output", "response", "data", "payload", - "error", "message", "status", "code", "type", - "name", "label", "title", "description", "content", - "items", "list", "array", "collection", "set", - "map", "table", "index", "cache", "buffer", - "count", "total", "sum", "average", "max", "min", - "width", "height", "size", "length", "offset", - "timeout", "interval", "delay", "duration", - "callback", "handler", "listener", "observer", - "template", "pattern", "schema", "format", - "prefix", "suffix", "separator", "delimiter", - "source", "target", "origin", "destination", - "parent", "child", "root", "node", "element", - "key", "value", "pair", "entry", "record", - "token", "secret", "hash", "salt", "nonce", - "baseUrl", "endpoint", "apiKey", "apiVersion", - "currentUser", "currentPage", "currentIndex", - "isLoading", "isValid", "isActive", "isVisible", - "hasError", "hasChanges", "hasPermission", - ], - }; + // Dictionaries imported from ./data/identifier-dictionaries.mjs - // Context strings commonly found near specific identifier types. - const CONTEXT_MAP = { - createElement: ["div", "span", "button", "input", "innerHTML"], - addEventListener: ["click", "submit", "change", "keydown", "DOMContentLoaded"], - fetchData: ["GET", "POST", "Content-Type", "application/json", "Authorization"], - createRouter: ["GET", "POST", "route", "middleware", "path"], - useState: ["setState", "initialState", "render", "component"], - parseJSON: ["JSON", "parse", "stringify", "object", "string"], - connectDatabase: ["connection", "host", "port", "database", "query"], - hashPassword: ["bcrypt", "salt", "rounds", "hash", "verify"], - validateEmail: ["email", "regex", "pattern", "valid", "invalid"], - HttpClient: ["fetch", "XMLHttpRequest", "headers", "method", "body"], - Router: ["route", "path", "handler", "middleware", "GET"], - EventEmitter: ["emit", "on", "off", "once", "listeners"], - Logger: ["log", "error", "warn", "info", "debug"], - Store: ["state", "dispatch", "subscribe", "getState", "reducer"], - }; - - // Property access patterns. - const PROPERTY_MAP = { - createElement: ["tagName", "className", "id", "style"], - fetchData: ["method", "headers", "body", "status"], - createRouter: ["method", "path", "handler", "params"], - Router: ["routes", "middleware", "use", "get", "post"], - Component: ["props", "state", "render", "componentDidMount"], - Store: ["state", "dispatch", "subscribe", "getState"], - Logger: ["level", "message", "timestamp", "format"], - config: ["host", "port", "database", "username"], - }; - - // Minifier name generators. + // Minifier name generators -- expanded with more strategies. const minifierStyles = [ - (i) => String.fromCharCode(97 + (i % 26)), // a, b, c... - (i) => String.fromCharCode(97 + (i % 26)) + "$", // a$, b$... - (i) => "_" + String.fromCharCode(97 + (i % 26)), // _a, _b... - (i) => "_0x" + (0x1a2b + i).toString(16), // _0x1a2b... - (i) => String.fromCharCode(97 + (i % 26)) + (i % 10).toString(), // a0, b1... - (i) => "__" + String.fromCharCode(97 + (i % 26)), // __a, __b... - (i) => "$" + String.fromCharCode(97 + (i % 26)), // $a, $b... - (i) => String.fromCharCode(65 + (i % 26)), // A, B, C... + // Single letter: a, b, c ... z + (i) => String.fromCharCode(97 + (i % 26)), + // With dollar suffix: a$, b$... + (i) => String.fromCharCode(97 + (i % 26)) + "$", + // Underscore prefix: _a, _b... + (i) => "_" + String.fromCharCode(97 + (i % 26)), + // Hex obfuscation: _0x1a2b... + (i) => "_0x" + (0x1a2b + i).toString(16), + // Letter + digit: a0, b1... + (i) => String.fromCharCode(97 + (i % 26)) + (i % 10).toString(), + // Double underscore: __a, __b... + (i) => "__" + String.fromCharCode(97 + (i % 26)), + // Dollar prefix: $a, $b... + (i) => "$" + String.fromCharCode(97 + (i % 26)), + // Uppercase single: A, B, C... + (i) => String.fromCharCode(65 + (i % 26)), + // Double letter: aa, ab, ac... + (i) => String.fromCharCode(97 + (i % 26)) + String.fromCharCode(97 + ((i + 1) % 26)), + // Mixed case: aA, bB, cC... + (i) => String.fromCharCode(97 + (i % 26)) + String.fromCharCode(65 + (i % 26)), + // Dollar + digit: $0, $1... + (i) => "$" + (i % 100).toString(), + // Underscore + digit: _0, _1... + (i) => "_" + (i % 100).toString(), + // Two letters + digit: aa1, ab2... + (i) => String.fromCharCode(97 + (i % 26)) + String.fromCharCode(97 + ((i * 7) % 26)) + (i % 10), + // Webpack style: __WEBPACK_MODULE_a__ + (i) => "__W" + String.fromCharCode(97 + (i % 26)) + "__", + // Terser numbered: t0, t1, t2... + (i) => "t" + i, + // esbuild style: e$a, e$b... + (i) => "e$" + String.fromCharCode(97 + (i % 26)), + ]; + + // Context variation templates for richer training signal. + const CONTEXT_TEMPLATES = [ + (ctx) => ctx, // original + (ctx) => ctx.length > 2 ? [...ctx.slice(1), ctx[0]] : ctx, // rotated + (ctx) => ctx.slice(0, 3), // truncated + (ctx) => [...ctx, "prototype", "constructor"], // with prototype hints + (ctx) => [...ctx, "undefined", "null", "true", "false"], // with literals ]; let syntheticCount = 0; + let globalIdx = 0; for (const [kind, names] of Object.entries(COMMON_NAMES)) { for (let i = 0; i < names.length; i++) { const original = names[i]; - // Generate multiple minified variants per original name. - const numVariants = Math.min(3, minifierStyles.length); + const baseCtx = CONTEXT_MAP[original] || generateGenericContext(original); + const baseProps = PROPERTY_MAP[original] || generateGenericProperties(kind); + + // Generate 8 minified variants per original name using a global + // counter so names from different kinds do not collide. + const numVariants = 8; for (let v = 0; v < numVariants; v++) { - const styleIdx = (i + v) % minifierStyles.length; - const minified = minifierStyles[styleIdx](i); - const ctx = CONTEXT_MAP[original] || []; - const props = PROPERTY_MAP[original] || []; + const styleIdx = (globalIdx + v) % minifierStyles.length; + const minified = minifierStyles[styleIdx](globalIdx); + + const ctxVariant = CONTEXT_TEMPLATES[v % CONTEXT_TEMPLATES.length]; + const ctx = ctxVariant(baseCtx.length > 0 ? baseCtx : ["unknown"]); pairs.push({ minified, original, - context_strings: ctx.length > 0 ? ctx : generateGenericContext(original), - properties: props.length > 0 ? props : generateGenericProperties(kind), + context_strings: ctx, + properties: baseProps, kind, }); syntheticCount++; } + globalIdx++; } } @@ -320,8 +259,8 @@ function generateCrossVersionPairs() { const existing = pairs.find((p) => p.original === original); if (!existing) continue; - // Simulate 2-3 additional "versions" with different minified names. - const versions = 2 + Math.floor(Math.random() * 2); + // Simulate 3-5 additional "versions" with different minified names. + const versions = 3 + Math.floor(Math.random() * 3); for (let v = 0; v < versions; v++) { const minified = generateRandomMinifiedName(); if (pairs.some((p) => p.minified === minified && p.original === original)) continue; @@ -344,18 +283,22 @@ function generateCrossVersionPairs() { * Generate a random minified-style variable name. */ function generateRandomMinifiedName() { + const letter = () => String.fromCharCode(97 + Math.floor(Math.random() * 26)); + const LETTER = () => String.fromCharCode(65 + Math.floor(Math.random() * 26)); + const digit = () => Math.floor(Math.random() * 10).toString(); const styles = [ - () => { - const c = String.fromCharCode(97 + Math.floor(Math.random() * 26)); - return c + Math.floor(Math.random() * 100); - }, - () => "_0x" + Math.floor(Math.random() * 0xffff).toString(16), - () => { - const a = String.fromCharCode(97 + Math.floor(Math.random() * 26)); - const b = String.fromCharCode(97 + Math.floor(Math.random() * 26)); - return a + b; - }, - () => "$" + String.fromCharCode(97 + Math.floor(Math.random() * 26)), + () => letter() + Math.floor(Math.random() * 100), // a42 + () => "_0x" + Math.floor(Math.random() * 0xffff).toString(16), // _0x3f1a + () => letter() + letter(), // ab + () => "$" + letter(), // $a + () => "_" + letter(), // _a + () => letter() + LETTER(), // aB + () => letter() + letter() + digit(), // ab3 + () => "__" + letter() + letter(), // __ab + () => "$" + digit() + digit(), // $42 + () => letter() + "$" + digit(), // a$3 + () => "_" + digit() + letter(), // _3a + () => "t" + Math.floor(Math.random() * 1000), // t523 ]; return styles[Math.floor(Math.random() * styles.length)](); }