mintrt 0.1.2

//     <<*>> Revenge of snowflake.
//     -^v-  SASWE
//     @ This source code is licensed under a dual license model:
//       1. MPL-2.0 for non-commercial use only
//       2. Commercial License for commercial use
//       See LICENSE and LICENSE_COMMERCIAL files for details.

//! Lua 字节码安全检查模块
//!
//! 参考 Lua 5.5 源码：src/security/lua_src/lundump.c
//! 
//! 极简流程：
//! 1. 使用 mlua 加载字节码文件（不执行）
//! 2. 通过调试接口获取 main 函数的指令数组
//! 3. 扫描 CALL (68) / TAILCALL (69) 指令
//! 4. 找到返回 false，没找到返回 true

use std::process::Command;
use std::path::Path;
use std::fs::File;
use std::io::{Read, BufReader};
use log::{info, trace, warn};
use crate::config::LUAC_PATH;

/// 检查 Lua 字节码文件是否可执行（不包含 CALL/TAILCALL）
/// 
/// 流程：
/// 1. 调用 luac.exe -l 反编译字节码
/// 2. 解析反汇编输出查找 CALL/TAILCALL 指令
/// 3. 找到返回 false，没找到返回 true
pub fn check_executable_lua_file(filename: &str) -> bool {
    //TODO  Release 模式使用纯 Rust 实现，Debug 模式使用 luac
    #[cfg(not(debug_assertions))]
    {
        todo!();
        check_bytecode_pure_rust(filename).unwrap_or_else(|e| {
            warn!("检查失败：{}", e);
            false
        })
    }
    
    #[cfg(debug_assertions)]
    {
        check_bytecode_with_luac(filename).unwrap_or_else(|e| {
            warn!("检查失败：{}", e);
            false
        })
    }
}

/// 使用 luac.exe 反编译并检查字节码
#[cfg(debug_assertions)]
fn check_bytecode_with_luac<P: AsRef<Path>>(file_path: P) -> Result<bool, String> {
    // 获取 luac.exe 的路径（使用绝对路径）
    let manifest_dir = std::env!("CARGO_MANIFEST_DIR");
    let luac_path = Path::new(manifest_dir).join(LUAC_PATH);
    
    if !luac_path.exists() {
        return Err(format!("luac.exe 不存在：{:?}", luac_path));
    }
    
    // 调用 luac -l 反编译字节码
    let output = Command::new(&luac_path)
        .arg("-l")
        .arg(file_path.as_ref())
        .output()
        .map_err(|e| format!("执行 luac 失败：{}", e))?;
    
    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        return Err(format!("luac 执行失败：{}", stderr));
    }
    
    // 解析反汇编输出
    let stdout = String::from_utf8_lossy(&output.stdout);
    
    // 查找 CALL 和 TAILCALL 指令（仅检查 main 函数）
    let mut in_main_function = false;
    let mut found_main_header = false;
    
    for (line_num, line) in stdout.lines().enumerate() {
        // 检测 main 函数开始
        if line.starts_with("main <") && line.contains("instructions") {
            in_main_function = true;
            found_main_header = true;
            trace!("开始检查 main 函数...");
            continue;
        }
        
        // 如果已经在 main 函数中
        if in_main_function {
            // 检测嵌套函数开始（遇到空行后的 function 行）
            if line.trim().is_empty() || line.starts_with("function <") {
                trace!("✓  main 函数检查完成，未检测到危险指令");
                return Ok(true); // 安全
            }
            
            // 匹配指令：行中包含大写的 CALL 或 TAILCALL，且前后有空白字符分隔
            if line.contains(" CALL ") || line.contains("\tCALL\t") || 
               line.contains("\tCALL ") || line.contains(" CALL\t") ||
               line.contains(" TAILCALL ") || line.contains("\tTAILCALL\t") ||
               line.contains("\tTAILCALL ") || line.contains(" TAILCALL\t") {
                warn!("检测到危险指令 @行 {}: {}", line_num + 1, line.trim());
                return Ok(false); // 不安全
            }
        }
    }
    
    // 如果没有找到 main 函数头
    if !found_main_header {
        return Err("luac 输出格式错误：未找到 main 函数".to_string());
    }
    
    trace!("✓ 字节码安全，未检测到 CALL/TAILCALL 指令");
    Ok(true) // 安全
}

/// 使用纯 Rust 解析 Lua 5.5 字节码（备选方案）
fn check_bytecode_pure_rust<P: AsRef<Path>>(file_path: P) -> Result<bool, String> {
    let file = File::open(file_path.as_ref())
        .map_err(|e| format!("无法打开文件：{}", e))?;
    let mut reader = BufReader::new(file);
    
    // 读取整个文件到内存
    let mut buffer = Vec::new();
    reader.read_to_end(&mut buffer)
        .map_err(|e| format!("读取文件失败：{}", e))?;
    
    // 验证魔数
    if buffer.len() < 4 || &buffer[0..4] != b"\x1B\x4C\x75\x61" {
        return Err("无效的 Lua 字节码魔数".to_string());
    }
    
    // 跳过头部 (魔数 4 + 版本 1 + 格式 1 + LUAC_DATA 6 + 格式检查约 24 字节)
    // Lua 5.5 头部总大小通常是 36-40 字节
    let mut pos = 4; // 跳过魔数
    
    // version (1) + format (1) + luac_data (6) = 8
    pos += 8;
    
    // 格式检查部分（动态大小）
    if pos >= buffer.len() {
        return Err("文件太短，缺少头部信息".to_string());
    }
    
    // int size (1) + int value (sz bytes)
    let int_sz = buffer[pos] as usize;
    pos += 1 + int_sz;
    
    // instruction size (1) + instruction value (sz bytes)
    if pos >= buffer.len() { return Err("格式错误".into()); }
    let inst_sz = buffer[pos] as usize;
    pos += 1 + inst_sz;
    
    // lua_Integer size (1) + value (sz bytes)
    if pos >= buffer.len() { return Err("格式错误".into()); }
    let intg_sz = buffer[pos] as usize;
    pos += 1 + intg_sz;
    
    // lua_Number size (1) + value (sz bytes)
    if pos >= buffer.len() { return Err("格式错误".into()); }
    let num_sz = buffer[pos] as usize;
    pos += 1 + num_sz;
    
    // 现在到达 main 函数
    // linedefined (varint)
    let (linedefined, n) = read_varint_at(&buffer, pos)?;
    pos += n;
    let _ = linedefined;
    
    // lastlinedefined (varint)
    let (lastlinedefined, n) = read_varint_at(&buffer, pos)?;
    pos += n;
    let _ = lastlinedefined;
    
    // numparams (byte)
    if pos >= buffer.len() { return Err("格式错误".into()); }
    let _numparams = buffer[pos];
    pos += 1;
    
    // flag (byte)
    if pos >= buffer.len() { return Err("格式错误".into()); }
    let _flag = buffer[pos];
    pos += 1;
    
    // maxstacksize (byte)
    if pos >= buffer.len() { return Err("格式错误".into()); }
    let _maxstacksize = buffer[pos];
    pos += 1;
    
    // numupvalues (varint) ← 新增！Lua 5.5 有这个字段
    let (numupvalues, n) = read_varint_at(&buffer, pos)?;
    pos += n;
    trace!("ℹnumupvalues = {}", numupvalues);
    
    // code_size (varint) ← 关键！
    let (code_size, n) = read_varint_at(&buffer, pos)?;
    pos += n;
    
    trace!("解析到 code_size = {}", code_size);
    
    // 检查指令
    for i in 0..code_size {
        if pos + 4 > buffer.len() {
            return Err(format!("指令 {} 超出文件范围", i));
        }
        
        let instruction = u32::from_le_bytes([
            buffer[pos],
            buffer[pos + 1],
            buffer[pos + 2],
            buffer[pos + 3]
        ]);
        pos += 4;
        
        let opcode = (instruction & 0x7F) as u8;
        
        // Lua 5.5 OP_CALL = 68, OP_TAILCALL = 69
        if opcode == 68 || opcode == 69 {
            warn!("检测到危险指令 #{}: opcode={} (0x{:02X})", i, opcode, opcode);
            return Ok(false); // 不安全
        }
    }
    
    info!("字节码安全，检查了 {} 条指令", code_size);
    Ok(true) // 安全
}

fn read_varint_at(buffer: &[u8], pos: usize) -> Result<(usize, usize), String> {
    let mut result: usize = 0;
    let mut shift = 0;
    let mut bytes_read = 0;
    
    loop {
        if pos + bytes_read >= buffer.len() {
            return Err("varint 超出缓冲区".to_string());
        }
        
        let byte = buffer[pos + bytes_read];
        result |= ((byte & 0x7F) as usize) << shift;
        shift += 7;
        bytes_read += 1;
        
        if byte < 0x80 {
            break;
        }
    }
    
    Ok((result, bytes_read))
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;
    use std::path::PathBuf;

    /// 获取 samples 目录的绝对路径
    fn get_samples_dir() -> PathBuf {
        let manifest_dir = env!("CARGO_MANIFEST_DIR");
        PathBuf::from(manifest_dir).join("src/security/samples")
    }

    /// 从 .wing 文件名提取不带后缀的名称
    fn extract_name_without_suffix(wing_path: &Path) -> String {
        wing_path
            .file_stem()
            .and_then(|s| s.to_str())
            .map(|s| s.to_string())
            .unwrap_or_default()
    }

    /// 去重：收集所有唯一的名称（不含后缀）
    fn collect_unique_names(samples_dir: &Path) -> Vec<String> {
        let mut names = std::collections::HashSet::new();
        
        if let Ok(entries) = fs::read_dir(samples_dir) {
            for entry in entries.flatten() {
                if let Some(name) = entry.path().file_stem().and_then(|s| s.to_str()) {
                    names.insert(name.to_string());
                }
            }
        }
        
        let mut unique_names: Vec<String> = names.into_iter().collect();
        unique_names.sort(); // 排序以保证测试顺序稳定
        unique_names
    }

    #[test]
    fn test_all_samples_automated() {
        let samples_dir = get_samples_dir();
        let unique_names = collect_unique_names(&samples_dir);
        
        println!("\n📁 测试目录：{:?}", samples_dir);
        println!("📋 发现 {} 个唯一名称", unique_names.len());
        
        let mut passed = 0;
        let mut failed = 0;
        
        for name in &unique_names {
            let wing_path = samples_dir.join(format!("{}.wing", name));
            let txt_path = samples_dir.join(format!("{}.txt", name));
            
            // 只测试同时有 .wing 和 .txt 的文件
            if !wing_path.exists() || !txt_path.exists() {
                println!("⏭️  跳过 {}: 缺少必要文件", name);
                continue;
            }
            
            // 读取预期结果
            let expected_str = match fs::read_to_string(&txt_path) {
                Ok(content) => content.trim().to_lowercase(),
                Err(e) => {
                    println!("❌ {}: 无法读取预期结果 - {}", name, e);
                    failed += 1;
                    continue;
                }
            };
            
            let expected_bool = match expected_str.as_str() {
                "true" | "1" | "yes" | "safe" => true,
                "false" | "0" | "no" | "unsafe" => false,
                _ => {
                    println!("⚠️  {}: 无效的预期结果 '{}'", name, expected_str);
                    failed += 1;
                    continue;
                }
            };
            
            // 执行检查
            let result = check_executable_lua_file(&wing_path.to_string_lossy());
            
            // 比对结果
            if result == expected_bool {
                println!("✅ {}: 通过 (expected={}, got={})", 
                    name, expected_bool, result);
                passed += 1;
            } else {
                println!("❌ {}: 失败 (expected={}, got={})", 
                    name, expected_bool, result);
                failed += 1;
            }
        }
        
        println!("\n📊 测试结果汇总：{} 通过，{} 失败", passed, failed);
        assert_eq!(failed, 0, "有 {} 个测试失败", failed);
    }

    // 保留原有的简单测试作为快速验证
    #[test]
    fn test_quick_validation() {
        let samples_dir = get_samples_dir();
        
        // 快速验证一个安全样本
        let safe_path = samples_dir.join("01_safe_def.wing");
        if safe_path.exists() {
            let result = check_executable_lua_file(&safe_path.to_string_lossy());
            assert!(result, "01_safe_def.wing 应该返回 true");
        }
        
        // 快速验证一个危险样本
        let unsafe_path = samples_dir.join("02_unsafe_call.wing");
        if unsafe_path.exists() {
            let result = check_executable_lua_file(&unsafe_path.to_string_lossy());
            assert!(!result, "02_unsafe_call.wing 应该返回 false");
        }
    }
}