smukx-RustyPacker/build.rs

//! Walks `src/techniques/{encryption,injection,evasion}/*/technique.toml`,
//! parses each manifest, and emits `$OUT_DIR/registry.rs` with:
//!   - one `pub static <ID>_META: TechniqueMeta = ...` per technique
//!   - a `REGISTRY: &[&'static dyn Technique]` array
//!
//! Each technique folder must contain a `mod.rs` that defines a unit struct
//! `pub struct <PascalId>;` implementing `Technique` and that uses
//! `crate::techniques::<category>::<id>::<PascalId>_META` for its meta.

use serde::Deserialize;
use std::collections::HashSet;
use std::fs;
use std::path::PathBuf;

#[derive(Debug, Deserialize)]
struct Manifest {
    id: String,
    display_name: String,
    description: String,
    category: String, // "encryption" | "injection" | "evasion"
    #[serde(default)]
    tags: Vec<String>,
    #[serde(default)]
    requires: Vec<String>, // "target_process" | "self_injection" | "format:exe" | "format:dll"
    #[serde(default)]
    incompatible_with: Vec<String>,
    #[serde(default)]
    template_dir: Option<String>, // for injection only
    #[serde(default)]
    params: Vec<ParamManifest>,
}

#[derive(Debug, Deserialize)]
struct ParamManifest {
    name: String,
    kind: String, // "text" | "bool" | "choice"
    label: String,
    #[serde(default)]
    default: Option<toml::Value>,
    #[serde(default)]
    options: Vec<String>,
}

fn main() {
    let crate_root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    let techniques_dir = crate_root.join("src/techniques");
    let out_dir: PathBuf = std::env::var_os("OUT_DIR").expect("OUT_DIR not set by cargo").into();

    println!("cargo:rerun-if-changed=build.rs");
    println!("cargo:rerun-if-changed=src/techniques");

    let mut found: Vec<(Manifest, String)> = Vec::new(); // (manifest, "<cat>/<id>")
    let mut seen_ids: HashSet<String> = HashSet::new();

    for cat in ["encryption", "injection", "evasion"] {
        let cat_dir = techniques_dir.join(cat);
        if !cat_dir.exists() {
            continue;
        }
        for entry in fs::read_dir(&cat_dir).unwrap_or_else(|e| panic!("read_dir {}: {e}", cat_dir.display())) {
            let entry = entry.unwrap_or_else(|e| panic!("read_dir entry in {}: {e}", cat_dir.display()));
            let path = entry.path();
            if !path.is_dir() {
                continue;
            }
            let manifest_path = path.join("technique.toml");
            if !manifest_path.exists() {
                continue;
            }
            println!("cargo:rerun-if-changed={}", manifest_path.display());

            let text = fs::read_to_string(&manifest_path)
                .unwrap_or_else(|e| panic!("read {}: {e}", manifest_path.display()));
            let manifest: Manifest = toml::from_str(&text)
                .unwrap_or_else(|e| panic!("parse {}: {e}", manifest_path.display()));

            if manifest.category != cat {
                panic!(
                    "manifest {} category={} but folder is under {}",
                    manifest_path.display(), manifest.category, cat
                );
            }
            if !is_valid_id(&manifest.id) {
                panic!("invalid technique id {:?} in {}: must match ^[a-z][a-z0-9_]*$",
                       manifest.id, manifest_path.display());
            }
            if manifest.template_dir.is_some() && manifest.category != "injection" {
                panic!("manifest {} declares template_dir but category={} (only 'injection' may use template_dir)",
                       manifest_path.display(), manifest.category);
            }
            if !seen_ids.insert(manifest.id.clone()) {
                panic!("duplicate technique id: {}", manifest.id);
            }
            found.push((manifest, format!("{}/{}", cat, path.file_name().unwrap().to_string_lossy())));
        }
    }

    found.sort_by(|a, b| a.0.id.cmp(&b.0.id));

    let mut out = String::new();
    out.push_str("// Auto-generated by build.rs. Do not edit.\n\n");

    // Emit module trees grouped by category.
    // We must use #[path = "..."] attributes because this file is emitted into
    // $OUT_DIR and included via include!() — Rust resolves bare `pub mod`
    // declarations relative to the including source file's directory, but when
    // the mod declaration lives inside an include!() expansion, rustc uses the
    // OUT_DIR path for resolution.  Absolute #[path] attributes sidestep this.
    for cat in ["encryption", "injection", "evasion"] {
        let items: Vec<&(Manifest, String)> = found.iter().filter(|(m, _)| m.category == cat).collect();
        if items.is_empty() { continue; }
        out.push_str(&format!("pub mod {cat} {{\n"));
        for (m, _) in &items {
            let mod_path = techniques_dir
                .join(cat)
                .join(&m.id)
                .join("mod.rs");
            // Use forward slashes so the path attribute works cross-platform inside the string literal.
            let mod_path_str = mod_path.to_string_lossy().replace('\\', "/");
            out.push_str(&format!("    #[path = {:?}]\n    pub mod {};\n", mod_path_str, m.id));
        }
        out.push_str("}\n");
    }
    out.push('\n');

    // Emit one `static META` per technique with the parsed values.
    for (m, _) in &found {
        emit_meta(&mut out, m);
    }

    // Emit the REGISTRY array.
    out.push_str("static REGISTRY: &[&'static dyn crate::techniques::Technique] = &[\n");
    for (m, _) in &found {
        let pascal = to_pascal(&m.id);
        out.push_str(&format!("    &crate::techniques::{}::{}::{} {{}},\n", m.category, m.id, pascal));
    }
    out.push_str("];\n");

    let dest = out_dir.join("registry.rs");
    fs::write(&dest, out).unwrap_or_else(|e| panic!("write {}: {e}", dest.display()));
}

fn is_valid_id(id: &str) -> bool {
    let mut chars = id.chars();
    let Some(first) = chars.next() else { return false; };
    if !first.is_ascii_lowercase() { return false; }
    chars.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_')
}

fn to_pascal(snake: &str) -> String {
    snake.split('_').map(|part| {
        let mut chars = part.chars();
        match chars.next() {
            Some(c) => c.to_ascii_uppercase().to_string() + chars.as_str(),
            None => String::new(),
        }
    }).collect()
}

fn emit_meta(out: &mut String, m: &Manifest) {
    let cat_variant = match m.category.as_str() {
        "encryption" => "Encryption",
        "injection" => "Injection",
        "evasion" => "Evasion",
        other => panic!("unknown category: {other}"),
    };
    let id_upper = m.id.to_uppercase();

    out.push_str(&format!(
        "pub static {ID}_META: crate::techniques::TechniqueMeta = crate::techniques::TechniqueMeta {{\n",
        ID = id_upper
    ));
    out.push_str(&format!("    id: {:?},\n", m.id));
    out.push_str(&format!("    display_name: {:?},\n", m.display_name));
    out.push_str(&format!("    description: {:?},\n", m.description));
    out.push_str(&format!("    category: crate::techniques::Category::{},\n", cat_variant));

    out.push_str("    tags: &[");
    for t in &m.tags { out.push_str(&format!("{:?}, ", t)); }
    out.push_str("],\n");

    out.push_str("    requires: &[");
    for r in &m.requires { out.push_str(&emit_requirement(r)); out.push_str(", "); }
    out.push_str("],\n");

    out.push_str("    incompatible_with: &[");
    for w in &m.incompatible_with { out.push_str(&format!("{:?}, ", w)); }
    out.push_str("],\n");

    out.push_str("    params: &[");
    for p in &m.params { out.push_str(&emit_param(p)); out.push_str(", "); }
    out.push_str("],\n");

    out.push_str("};\n\n");

    if let Some(t) = &m.template_dir {
        out.push_str(&format!(
            "pub static {ID}_TEMPLATE_DIR: &str = {t:?};\n\n",
            ID = id_upper
        ));
    }
}

fn emit_requirement(r: &str) -> String {
    match r {
        "target_process" => "crate::techniques::Requirement::TargetProcess".to_string(),
        "self_injection" => "crate::techniques::Requirement::SelfInjection".to_string(),
        "format:exe" => "crate::techniques::Requirement::Format(crate::techniques::Format::Exe)".to_string(),
        "format:dll" => "crate::techniques::Requirement::Format(crate::techniques::Format::Dll)".to_string(),
        other => panic!("unknown requirement: {other}"),
    }
}

fn emit_param(p: &ParamManifest) -> String {
    match p.kind.as_str() {
        "text" => {
            let default = p.default.as_ref().and_then(|v| v.as_str()).unwrap_or("");
            format!(
                "crate::techniques::ParamSpec::Text {{ name: {:?}, label: {:?}, default: {:?} }}",
                p.name, p.label, default
            )
        }
        "bool" => {
            let default = p.default.as_ref().and_then(|v| v.as_bool()).unwrap_or(false);
            format!(
                "crate::techniques::ParamSpec::Bool {{ name: {:?}, label: {:?}, default: {} }}",
                p.name, p.label, default
            )
        }
        "choice" => {
            if p.options.is_empty() {
                panic!("choice param {:?} has no options", p.name);
            }
            let opts: Vec<String> = p.options.iter().map(|o| format!("{:?}", o)).collect();
            format!(
                "crate::techniques::ParamSpec::Choice {{ name: {:?}, label: {:?}, options: &[{}] }}",
                p.name, p.label, opts.join(", ")
            )
        }
        other => panic!("unknown param kind: {other}"),
    }
}