aegis_vm_macro 0.2.51

//! Compile-time cryptography for bytecode encryption
//!
//! Uses the same algorithms as anticheat-vm/crypto.rs

use hmac::{Hmac, Mac};
use sha2::Sha256;

type HmacSha256 = Hmac<Sha256>;

pub const KEY_SIZE: usize = 32;
pub const NONCE_SIZE: usize = 12;
pub const TAG_SIZE: usize = 16;

// Domain separation strings (same as runtime)
const NONCE_DOMAIN: &[u8] = b"anticheat-vm-nonce-v1";
const BUILDID_DOMAIN: &[u8] = b"anticheat-vm-build-id-v1";
const OPCODE_SHUFFLE_DOMAIN: &[u8] = b"opcode-shuffle-v1";

/// Get build seed from shared file
/// This is called at proc-macro expansion time (compile time)
pub fn get_build_seed() -> [u8; 32] {
    // ALWAYS read shared seed file written by vm build.rs
    // We do NOT calculate seed from env var here because HMAC implementation
    // might differ from build.rs (different crates/versions), leading to opcode mismatch.
    if let Some(seed) = read_shared_seed() {
        return seed;
    }

    // PANIC if seed not found - do NOT silently generate different seed!
    // This would cause opcode mismatch between vm-macro and vm runtime.
    panic!(
        "vm-macro: Could not find shared build seed file (.anticheat_build_seed). \
         Make sure anticheat-vm is built before using vm_protect macro. \
         Or set ANTICHEAT_BUILD_KEY environment variable for reproducible builds."
    );
}

/// Read shared seed from target directory
fn read_shared_seed() -> Option<[u8; 32]> {
    use std::fs;
    use std::path::PathBuf;

    let mut all_candidates: Vec<PathBuf> = Vec::new();

    // Try CARGO_TARGET_DIR if set (highest priority)
    if let Ok(target_dir) = std::env::var("CARGO_TARGET_DIR") {
        all_candidates.push(PathBuf::from(&target_dir).join(".anticheat_build_seed"));
    }

    // Try OUT_DIR based path (works during build)
    if let Ok(out_dir) = std::env::var("OUT_DIR") {
        let path = PathBuf::from(&out_dir);
        // Walk up to find target directory
        for ancestor in path.ancestors() {
            if ancestor.file_name().is_some_and(|n| n == "target") {
                all_candidates.push(ancestor.join(".anticheat_build_seed"));
                break;
            }
        }
    }

    // Try CARGO_MANIFEST_DIR based path (works for proc-macros)
    if let Ok(manifest_dir) = std::env::var("CARGO_MANIFEST_DIR") {
        let path = PathBuf::from(&manifest_dir);
        // Walk up to find project root with target directory
        for ancestor in path.ancestors() {
            let target_path = ancestor.join("target/.anticheat_build_seed");
            if target_path.exists() {
                all_candidates.push(target_path);
                break;
            }
        }
    }

    // Common relative paths as fallback
    all_candidates.extend([
        PathBuf::from("target/.anticheat_build_seed"),
        PathBuf::from("../target/.anticheat_build_seed"),
        PathBuf::from("../../target/.anticheat_build_seed"),
        PathBuf::from("../../../target/.anticheat_build_seed"),
        PathBuf::from("../../../../target/.anticheat_build_seed"),
    ]);

    for path in all_candidates {
        if let Ok(hex_str) = fs::read_to_string(&path) {
            if let Ok(bytes) = hex::decode(hex_str.trim()) {
                if bytes.len() == 32 {
                    let mut seed = [0u8; 32];
                    seed.copy_from_slice(&bytes);
                    return Some(seed);
                }
            }
        }
    }

    None
}

// NOTE: generate_random_seed() removed - not used in production.
// Build seed is always read from shared file written by aegis_vm build.rs
// NOTE: derive_key() removed - WBC key derivation is now used instead.

/// Derive nonce from build seed and counter
pub fn derive_nonce(build_seed: &[u8; 32], counter: u64) -> [u8; NONCE_SIZE] {
    let mut mac = <HmacSha256 as Mac>::new_from_slice(build_seed)
        .expect("HMAC can take any size key");
    mac.update(&counter.to_le_bytes());
    mac.update(NONCE_DOMAIN);

    let result = mac.finalize();
    let mut nonce = [0u8; NONCE_SIZE];
    nonce.copy_from_slice(&result.into_bytes()[..NONCE_SIZE]);
    nonce
}

/// Derive build ID from seed
pub fn derive_build_id(build_seed: &[u8; 32]) -> u64 {
    let mut mac = <HmacSha256 as Mac>::new_from_slice(build_seed)
        .expect("HMAC can take any size key");
    mac.update(BUILDID_DOMAIN);

    let result = mac.finalize();
    let bytes = result.into_bytes();
    u64::from_le_bytes([
        bytes[0], bytes[1], bytes[2], bytes[3],
        bytes[4], bytes[5], bytes[6], bytes[7],
    ])
}

/// Encrypt bytecode using AES-256-GCM
pub fn encrypt_bytecode(
    key: &[u8; KEY_SIZE],
    nonce: &[u8; NONCE_SIZE],
    plaintext: &[u8],
) -> Result<(Vec<u8>, [u8; TAG_SIZE]), String> {
    use aes_gcm::{Aes256Gcm, KeyInit, aead::Aead, Nonce};

    let cipher = Aes256Gcm::new_from_slice(key)
        .map_err(|e| format!("Failed to create cipher: {}", e))?;

    let nonce_obj = Nonce::from_slice(nonce);

    let ciphertext = cipher
        .encrypt(nonce_obj, plaintext)
        .map_err(|e| format!("Encryption failed: {}", e))?;

    // Extract tag from the end
    if ciphertext.len() < TAG_SIZE {
        return Err("Ciphertext too short".to_string());
    }

    let tag_start = ciphertext.len() - TAG_SIZE;
    let mut tag = [0u8; TAG_SIZE];
    tag.copy_from_slice(&ciphertext[tag_start..]);

    let encrypted_data = ciphertext[..tag_start].to_vec();

    Ok((encrypted_data, tag))
}

/// Encrypted bytecode package for embedding
pub struct EncryptedPackage {
    pub build_id: u64,
    pub nonce: [u8; NONCE_SIZE],
    pub tag: [u8; TAG_SIZE],
    pub ciphertext: Vec<u8>,
}

/// Encrypt bytecode with build seed
/// Uses WBC-derived key for maximum key-hiding protection
pub fn encrypt_with_seed(bytecode: &[u8], function_id: u64) -> Result<EncryptedPackage, String> {
    let seed = get_build_seed();

    // Use WBC to derive the encryption key (key-hiding)
    let key = crate::whitebox::derive_bytecode_key(&seed);

    // Nonce is still HMAC-derived (no need to hide nonce)
    let nonce = derive_nonce(&seed, function_id);
    let build_id = derive_build_id(&seed);

    let (ciphertext, tag) = encrypt_bytecode(&key, &nonce, bytecode)?;

    Ok(EncryptedPackage {
        build_id,
        nonce,
        tag,
        ciphertext,
    })
}

/// Base opcode definitions (canonical values)
/// Must match the definitions in vm/build.rs EXACTLY (same order!)
const BASE_OPCODES: &[u8] = &[
    // Stack operations
    0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09,
    // Register operations
    0x10, 0x11, 0x12, 0x13,
    // Arithmetic operations
    0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26, 0x27, 0x28, 0x29, 0x2A, 0x2B, 0x2C, 0x46, 0x47, 0x48, 0x49,
    // Control flow
    0x30, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39,
    // Special operations
    0x40, 0x41, 0x42, 0x43, 0x44, 0x45,
    // Type conversion
    0x50, 0x51, 0x52, 0x53, 0x54, 0x55,
    // Memory operations
    0x60, 0x61, 0x62, 0x63, 0x64, 0x65, 0x66, 0x67,
    // Heap operations (must match vm/build.rs)
    0x70, 0x71, 0x72, 0x73, 0x74, 0x75, 0x76, 0x77, 0x78, 0x79, 0x7A,
    // Vector operations (must match vm/build.rs)
    0x80, 0x81, 0x82, 0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89,
    // String operations (must match vm/build.rs)
    0x90, 0x91, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98,
    // Native calls
    0xF0, 0xF1, 0xF2, 0xF3,
    // Execution control (HALT/HALT_ERR are kept fixed)
];

/// Opcode encoding table (base -> shuffled)
/// Generated at compile time from build seed
#[derive(Clone)]
pub struct OpcodeTable {
    pub encode: [u8; 256],
    seed: u64, // For MBA determinism
}

impl OpcodeTable {
    /// Generate opcode table from build seed
    /// This must produce identical results to vm/build.rs
    pub fn generate(seed: &[u8; 32]) -> Self {
        // Derive shuffle key from seed
        let mut mac = <HmacSha256 as Mac>::new_from_slice(seed)
            .expect("HMAC can take any size key");
        mac.update(OPCODE_SHUFFLE_DOMAIN);
        let shuffle_key: [u8; 32] = mac.finalize().into_bytes().into();

        // Create list of available byte values (0x00-0xFD, excluding 0xFE and 0xFF for HALT)
        let mut available: Vec<u8> = (0x00..0xFE).collect();

        // Fisher-Yates shuffle using HMAC-derived randomness
        let mut rng_state = shuffle_key;
        for i in (1..available.len()).rev() {
            // Get next random index using HMAC
            let mut mac = <HmacSha256 as Mac>::new_from_slice(&rng_state)
                .expect("HMAC can take any size key");
            mac.update(&(i as u32).to_le_bytes());
            let rand_bytes: [u8; 32] = mac.finalize().into_bytes().into();
            rng_state = rand_bytes;

            let j = (u64::from_le_bytes([
                rand_bytes[0], rand_bytes[1], rand_bytes[2], rand_bytes[3],
                rand_bytes[4], rand_bytes[5], rand_bytes[6], rand_bytes[7],
            ]) as usize) % (i + 1);
            available.swap(i, j);
        }

        // Build the encoding table
        let mut encode = [0u8; 256];

        // Initialize as identity mapping
        for (i, val) in encode.iter_mut().enumerate() {
            *val = i as u8;
        }

        // Assign shuffled values to each base opcode (except HALT/HALT_ERR)
        let mut available_idx = 0;
        for &base_val in BASE_OPCODES {
            if base_val == 0xFF || base_val == 0xFE {
                continue;
            }
            let shuffled_val = available[available_idx];
            available_idx += 1;
            encode[base_val as usize] = shuffled_val;
        }

        // Derive seed for MBA from shuffle key
        let mba_seed = u64::from_le_bytes([
            rng_state[0], rng_state[1], rng_state[2], rng_state[3],
            rng_state[4], rng_state[5], rng_state[6], rng_state[7],
        ]);

        OpcodeTable { encode, seed: mba_seed }
    }

    /// Encode a base opcode to its shuffled value
    #[inline]
    pub fn encode(&self, base_opcode: u8) -> u8 {
        self.encode[base_opcode as usize]
    }

    /// Get seed for MBA transformations
    #[inline]
    pub fn get_seed(&self) -> u64 {
        self.seed
    }
}

/// Read opcode table from shared file written by build.rs
/// This is the Single Source of Truth - we read what build.rs generated
fn read_shared_opcode_table() -> Option<[u8; 256]> {
    use std::fs;
    use std::path::PathBuf;

    let mut all_candidates: Vec<PathBuf> = Vec::new();

    // Try CARGO_TARGET_DIR if set (highest priority)
    if let Ok(target_dir) = std::env::var("CARGO_TARGET_DIR") {
        all_candidates.push(PathBuf::from(&target_dir).join(".anticheat_opcode_table"));
    }

    // Try OUT_DIR based path (works during build)
    if let Ok(out_dir) = std::env::var("OUT_DIR") {
        let path = PathBuf::from(&out_dir);
        for ancestor in path.ancestors() {
            if ancestor.file_name().is_some_and(|n| n == "target") {
                all_candidates.push(ancestor.join(".anticheat_opcode_table"));
                break;
            }
        }
    }

    // Try CARGO_MANIFEST_DIR based path (works for proc-macros)
    if let Ok(manifest_dir) = std::env::var("CARGO_MANIFEST_DIR") {
        let path = PathBuf::from(&manifest_dir);
        for ancestor in path.ancestors() {
            let target_path = ancestor.join("target/.anticheat_opcode_table");
            if target_path.exists() {
                all_candidates.push(target_path);
                break;
            }
        }
    }

    // Common relative paths as fallback
    all_candidates.extend([
        PathBuf::from("target/.anticheat_opcode_table"),
        PathBuf::from("../target/.anticheat_opcode_table"),
        PathBuf::from("../../target/.anticheat_opcode_table"),
        PathBuf::from("../../../target/.anticheat_opcode_table"),
        PathBuf::from("../../../../target/.anticheat_opcode_table"),
    ]);

    for path in all_candidates {
        if let Ok(hex_str) = fs::read_to_string(&path) {
            if let Ok(bytes) = hex::decode(hex_str.trim()) {
                if bytes.len() == 256 {
                    let mut table = [0u8; 256];
                    table.copy_from_slice(&bytes);
                    return Some(table);
                }
            }
        }
    }

    None
}

/// Get the opcode table for current build
/// CRITICAL: Reads from shared file written by build.rs (Single Source of Truth)
/// Falls back to generation only if file not found (should not happen in normal builds)
pub fn get_opcode_table() -> OpcodeTable {
    // First try to read from shared file (the correct way)
    if let Some(encode) = read_shared_opcode_table() {
        // Derive MBA seed from build seed (this part is still needed)
        let seed = get_build_seed();
        let mut mac = <HmacSha256 as Mac>::new_from_slice(&seed)
            .expect("HMAC can take any size key");
        mac.update(OPCODE_SHUFFLE_DOMAIN);
        let shuffle_key: [u8; 32] = mac.finalize().into_bytes().into();

        // Derive MBA seed the same way as OpcodeTable::generate
        let mut rng_state = shuffle_key;
        for i in (1..254).rev() {
            let mut mac = <HmacSha256 as Mac>::new_from_slice(&rng_state)
                .expect("HMAC can take any size key");
            mac.update(&(i as u32).to_le_bytes());
            rng_state = mac.finalize().into_bytes().into();
        }

        let mba_seed = u64::from_le_bytes([
            rng_state[0], rng_state[1], rng_state[2], rng_state[3],
            rng_state[4], rng_state[5], rng_state[6], rng_state[7],
        ]);

        return OpcodeTable { encode, seed: mba_seed };
    }

    // Fallback: generate ourselves (should not happen if build order is correct)
    // This maintains backwards compatibility but is not the preferred path
    eprintln!("WARNING: vm-macro: Opcode table file not found, generating locally. \
               This may cause opcode mismatch if build.rs uses different algorithm!");
    let seed = get_build_seed();
    OpcodeTable::generate(&seed)
}