moon-struct 0.1.0

use std::collections::HashMap;
use std::env;
use std::fs;
use std::fs::File;
use std::fs::OpenOptions;
use std::io::BufReader;
use std::io::Write;
use std::path::Path;
use std::path::PathBuf;
use std::thread;
use std::time::Duration;

use scraper::Html;
use scraper::Selector;

extern crate regex;

use regex::Regex;
use serde::Deserialize;
use utils::calculate_md5;
use utils::create_file_with_directories;

pub mod offset_config;
pub mod utils;

#[macro_use]
extern crate lazy_static;

fn p(s: String) {
    let mut file = OpenOptions::new()
        .create(true)
        .append(true)
        .open("output.txt")
        .expect("Unable to open file");

    writeln!(file, "{}", s).expect("Unable to write to file");
}

fn collect_os_version(url: &str, version_desc: &str) -> Vec<HashMap<String, String>> {
    let mut os_versions = Vec::<HashMap<String, String>>::new();

    let response = reqwest::blocking::get(url);
    if response.is_err() {
        p("Error to get response".to_string());
        return os_versions;
    }
    let response = response.unwrap();
    let body = response.text();
    if body.is_err() {
        p("Failed to get response text".to_string());
        return os_versions;
    }
    let body = body.unwrap();

    let document = Html::parse_document(&body);
    let selector = Selector::parse(
        "html > body > div > div:nth-of-type(2) > div > div > div > div > div > div",
    )
    .unwrap();

    for element in document.select(&selector) {
        let text = element.text().collect::<Vec<_>>();

        for i in (0..text.len()).step_by(3) {
            let mut os_version = HashMap::new();
            os_version.insert("windows_version".to_string(), version_desc.to_string());
            os_version.insert("windows_version_desc".to_string(), text[i].to_string());
            os_version.insert(
                "windows_version_number".to_string(),
                text[i + 1].to_string(),
            );
            os_version.insert("windows_version_time".to_string(), text[i + 2].to_string());
            os_versions.push(os_version);
        }
    }

    os_versions
}

fn build_offset_struct_str() -> String {
    let mut r = "".to_string();
    r.push_str(
        "#[derive(Clone, Copy)]
    pub struct Offset {",
    );
    for ele in &*offset_config::NEED_STRUCT {
        r.push_str(&format!("pub {}: usize,\n", ele.variable_name).to_string());
    }
    r.push_str("}\n");
    r
}

fn build_os_version_info_struct_str() -> String {
    "#[derive(Clone, Copy)]
    pub struct OSVersionInfo {
        pub major_version: u32,
        pub minor_version: u32,
        pub build_number: u32,
        pub patch_number: u32,
        pub windows_version: &'static str,
        pub version_desc: &'static str,
        pub offset: Offset,
        pub func_offset: FunctionOffset,
    }"
    .to_string()
}

fn get_offset_by_url(
    windows_version: String,
    windows_version_desc: String,
    struct_name: String,
    struct_field_define: String,
) -> String {
    let s: Vec<&str> = windows_version_desc.split(" (").collect::<Vec<&str>>();
    let mut s = s[0].to_string();
    let is = s.split(" | ").collect::<Vec<&str>>();
    if is.len() > 1 {
        s = is[0].to_lowercase();
    }

    if s == "Insider Preview" {
        s = "insider-preview-jun-2021".to_string();
    }

    let windows_version_number = s.to_lowercase();
    let url = format!(
        "https://www.vergiliusproject.com/kernels/x64/{}/{}/{}",
        windows_version, windows_version_number, struct_name
    );

    let response = reqwest::blocking::get(url);
    if response.is_err() {
        p("Error to get response".to_string());
        panic!();
    }

    let response = response.unwrap();
    let body = response.text();
    if body.is_err() {
        p("Failed to get response text".to_string());
        panic!();
    }
    let body = body.unwrap();

    // struct_field_define
    let mut pattern = format!(
        r" {}[ \t]*//(0x[0-9a-fA-F]+)",
        regex::escape(&struct_field_define)
    );

    if struct_field_define == "sizeof" {
        pattern = r"//(0x[0-9a-fA-F]+) bytes \(sizeof\)".to_string();
    }

    let re = Regex::new(&pattern).unwrap();

    if let Some(captures) = re.captures(&body) {
        if let Some(value) = captures.get(1) {
            let r = value.as_str().to_string();
            // p(format!(
            //     "{},{} {} = {}",
            //     windows_version, windows_version_desc, struct_field_define, r
            // )
            // .to_string());
            return r;
        }
    }

    p(format!(
        "{},{} do not find field {}",
        windows_version, windows_version_desc, struct_field_define
    )
    .to_string());

    String::new()
}

fn build_spec_os_func_index_str(number: &str) -> String {
    p(number.to_string());
    let out_dir = env::var("CARGO_MANIFEST_DIR").unwrap();
    p(out_dir.clone());

    let mut input_path =
        PathBuf::from(out_dir.clone()).join(format!("srcipts/offset/{}.json", number));
    if !input_path.exists() {
        // p(format!("{} use default.json",number));
        input_path = PathBuf::from(out_dir.clone()).join(format!("srcipts/offset/default.json"));
        if !input_path.exists() {
            p("not find default.json".to_string());
            panic!();
        }
    } else {
        p(format!("{} use not default.json", number));
    }

    // p(format!("input_path={}", input_path.to_str().unwrap()).to_string());

    let content = fs::read_to_string(input_path);
    if content.is_err() {
        p("read input json fault".to_string());
        panic!();
    }
    let content = content.unwrap();

    let mut r = "".to_string();

    let root = serde_json::from_str(&content);
    if root.is_err() {
        p("Failed to parse JSON".to_string());
        return "".to_string();
    }

    let root: OutputJson = root.unwrap();

    r.push_str("FunctionOffset {");
    for item in root.result.into_iter() {
        r.push_str(&format!("{}: {},", item.variable, item.offset));
    }

    r.push_str("}");

    r
}

fn build_spec_os_offset_str(
    windows_version: String,
    windows_version_desc: String,
    number: &str,
) -> String {
    let current_dir = env::current_dir().expect("Failed to get current directory");
    let tmp = current_dir.join("offset_config.rs");
    let offset_config_path = tmp.to_str().expect("current_dir exepect");

    let md5_str = calculate_md5(offset_config_path).expect("calculate md5 fault");

    let cache_file = current_dir.join("temp").join(md5_str).join(number);

    let cache_file_path = cache_file.to_str().expect("cache_file to str fault");

    if cache_file.exists() {
        p(format!("cache file exist:{}", cache_file_path));
        let cache = fs::read_to_string(cache_file).expect("read cache file fault");
        return cache;
    }

    p(format!("cache file not exist:{}", cache_file_path));

    let mut r = "".to_string();
    r.push_str("Offset {");

    for item in &*offset_config::NEED_STRUCT {
        // https://www.vergiliusproject.com/kernels/x64/windows-10/22h2/DEBUG_MEMORY_REQUIREMENTS
        let off = get_offset_by_url(
            windows_version.to_string(),
            windows_version_desc.to_string(),
            item.struct_name.to_string(),
            item.struct_field_define.to_string(),
        );

        r.push_str(
            &format!(
                "{}: {},\n",
                item.variable_name,
                format!("{} {}", off, item.value_endfix)
            )
            .to_string(),
        );
    }

    r.push_str("}");

    create_file_with_directories(cache_file_path).expect("create_file_with_directories fault");

    fs::write(cache_file, &r).expect("write offset cache file fault");

    r
}

#[allow(dead_code)]
#[derive(Deserialize, Debug)]
struct Func {
    target_func: String,
    variable: String,
}

#[derive(Deserialize, Debug)]
struct Result {
    variable: String,
    offset: String,
}

#[derive(Deserialize, Debug)]
struct InputJson {
    func: Vec<Func>,
}

#[derive(Deserialize, Debug)]
struct OutputJson {
    result: Vec<Result>,
}

fn build_function_index_struct_str() -> String {
    let mut r = "".to_string();
    let file = File::open("d:\\input.json");
    if file.is_err() {
        p("parse input.json fault".to_string());
        return "".to_string();
    }

    let file = file.unwrap();

    let reader = BufReader::new(file);

    // 解析 JSON 数据
    let root = serde_json::from_reader(reader);
    if root.is_err() {
        p("parse root fault".to_string());
        return "".to_string();
    }

    let root: InputJson = root.unwrap();

    // 提取所有 variable
    let variables: Vec<String> = root.func.into_iter().map(|f| f.variable).collect();

    r.push_str("#[derive(Clone, Copy)]");
    r.push_str("pub struct FunctionOffset {");

    for v in variables {
        r.push_str(&format!("\t pub {}: u64,\n", v));
    }

    r.push_str("}");

    r
}

fn build_index_struct_str() -> String {
    return "#[derive(Clone, Copy)]
    pub struct Index {
        pub target_func: &'static str,
        pub offset: u64,
    }"
    .to_string();
}

fn main() {
    // file exist do nothing
    if Path::new("src").join("os_versions.rs").metadata().is_ok() {
        return;
    }

    let url_win10 = "https://www.vergiliusproject.com/kernels/x64/windows-10";
    let url_win11 = "https://www.vergiliusproject.com/kernels/x64/windows-11";

    let mut win10 = collect_os_version(&url_win10, "windows-10");
    if win10.len() == 0 {
        return;
    }
    let mut win11 = collect_os_version(&url_win11, "windows-11");
    if win11.len() == 0 {
        return;
    }

    let mut all = Vec::<HashMap<String, String>>::new();

    all.append(&mut win10);
    all.append(&mut win11);

    let mut os_version_map_str = format!(
        "pub static VERSION_MAP: [OSVersionInfo; {}] = [
        ",
        all.len()
    );

    for i in 0..all.len() {
        let item = &all[i];
        let number = &item["windows_version_number"];
        let number = number.replace("build: ", "");
        let split: Vec<&str> = number.split(".").collect::<Vec<&str>>();

        os_version_map_str.push_str("OSVersionInfo {");
        let s = format!(
            "
            major_version: {},
            minor_version: {},
            build_number: {},
            patch_number: {},
            windows_version: \"{}\",
            version_desc: \"{}\",
            offset: {},
            func_offset: {},
            \n",
            split[0],
            split[1],
            split[2],
            split[3],
            item["windows_version"],
            item["windows_version_desc"],
            build_spec_os_offset_str(
                item["windows_version"].clone(),
                item["windows_version_desc"].clone(),
                &number
            ),
            build_spec_os_func_index_str(&number[0..10]),
        );

        os_version_map_str.push_str(&s);
        os_version_map_str.push_str("\t\t},\n");

        thread::sleep(Duration::from_secs(1));
    }
    os_version_map_str.push_str("];");

    let offset_struct_str = build_offset_struct_str();
    let index_struct_str = build_index_struct_str();
    let function_index_struct_str = build_function_index_struct_str();
    let os_version_info_struct_str = build_os_version_info_struct_str();

    // p(function_index_struct_str.clone());

    let mut file = OpenOptions::new()
        .create(true)
        .append(true)
        .open(Path::new("src").join("os_versions.rs"))
        .expect("Unable to open file");

    writeln!(file, "{}", offset_struct_str).expect("Unable to write to file");
    writeln!(file, "{}", index_struct_str).expect("Unable to write to file");
    writeln!(file, "{}", function_index_struct_str).expect("Unable to write to file");
    writeln!(file, "{}", os_version_info_struct_str).expect("Unable to write to file");
    writeln!(file, "{}", os_version_map_str).expect("Unable to write to file");
}