discovery.rs

use std::{borrow::Cow, ffi::CStr};

use iced_x86::{Code, Register};
use memchr::memmem;

use crate::exe_image::ExeImage;

/// It's actually almost same as the PE timestamp I've been using, but
/// they might have some more human-readable stuff here.
pub fn find_noita_build(image: &ExeImage) -> Option<Cow<str>> {
    let pos = memmem::find(image.rdata(), b"Noita - Build ")?;

    // + 8 to skip the "Noita - " part
    let prefix = image.rdata()[pos + 8..].split(|b| *b == 0).next()?;
    Some(String::from_utf8_lossy(prefix))
}

/// Assuming Lua API functions are set up like this..
/// ```c
///   lua_pushcclosure(L,function_pointer,0);
///   lua_setfield(L,LUA_GLOBALSINDEX,"UniqueString");
/// ```
/// ..we look for the `PUSH imm32` of the unique string given as `name`, and
/// then we look if there is a `PUSH imm32` at 8 bytes before that
/// (`CALL EDI => lua_pushcclosure` and `PUSH EBX` being 3 bytes, and
/// 5 bytes for the `PUSH imm32` image), and return it's argument.
///
/// Note that this completely breaks (already) with noita_dev.exe lol
fn find_lua_api_fn(image: &ExeImage, name: &CStr) -> Option<u32> {
    match image.text()[image.find_push_str_pos(name)? - 8..] {
        [0x68, a, b, c, d, ..] => {
            let addr = u32::from_le_bytes([a, b, c, d]);
            tracing::debug!("Found Lua API function {name:?} at 0x{addr:x}");
            Some(addr)
        }
        _ => {
            tracing::warn!("Did not find Lua API function {name:?}");
            None
        }
    }
}

/// We look for the `SetRandomSeed` Lua API function and then we look for
/// the `mov eax, [addr]` and `add eax, [addr]` instructions, which
/// correspond to WORLD_SEED + NEW_GAME_PLUS_COUNT being passed as a second
/// parameter of a (SetRandomSeedImpl) function call.
pub fn find_seed_pointers(image: &ExeImage) -> Option<(u32, u32)> {
    let mut state = None;
    for instr in image.decode_fn(find_lua_api_fn(image, c"SetRandomSeed")?) {
        state = match state {
            None if instr.code() == Code::Mov_EAX_moffs32 => Some(instr.memory_displacement32()),
            // allow the `add esp, 0x10` thing in between
            Some(addr) if instr.code() == Code::Add_rm32_imm8 => Some(addr),
            Some(addr)
                if instr.code() == Code::Add_r32_rm32 && instr.op0_register() == Register::EAX =>
            {
                return Some((addr, instr.memory_displacement32()));
            }
            _ => None,
        };
    }
    None
}

/// We look for the `AddFlagPersistent` Lua API function and then we look
/// for second-to-last `CALL rel32`, the last being some C++ exception
/// thing, and the second-to-last being a call to `AddFlagPersistentImpl`,
/// as I call it.
///
/// Then inside of that we look for `MOV ECX imm32` which is specifically
/// after `CALL rel32` which is after `MOV EDX, "progress_ending1"`.
/// The call being to a string equality check and our MOV being an
/// argument to a following call which is the global KEY_VALUE_STATS map
/// pointer.
pub fn find_stats_map_pointer(image: &ExeImage) -> Option<u32> {
    let mut before_last_call_rel = None;
    let mut last_call_rel = None;
    for instr in image.decode_fn(find_lua_api_fn(image, c"AddFlagPersistent")?) {
        if instr.code() == Code::Call_rel32_32 {
            before_last_call_rel = last_call_rel;
            last_call_rel = Some(instr.near_branch32());
        }
    }

    let end1_addr = image.find_string(c"progress_ending1")?;

    enum State {
        Init,
        FoundProgressEnding1,
        FoundStreqCall,
    }
    let mut state = State::Init;

    for instr in image.decode_fn(before_last_call_rel?) {
        match state {
            State::Init
                if instr.code() == Code::Mov_r32_imm32
                    && instr.op0_register() == Register::EDX
                    && instr.immediate32() == end1_addr =>
            {
                state = State::FoundProgressEnding1;
            }
            State::FoundProgressEnding1 if instr.code() == Code::Call_rel32_32 => {
                state = State::FoundStreqCall;
            }
            State::FoundStreqCall
                if instr.code() == Code::Mov_r32_imm32 && instr.op0_register() == Register::ECX =>
            {
                return Some(instr.immediate32());
            }
            _ => {}
        };
    }
    None
}

/// We look for the `EntityGetParent` Lua API function and there we look
/// for `MOV ECX, [addr]` which immediately follows a Lua call - that MOV
/// happens to be setting up an argument to a following relative call that
/// is the pointer to the entity manager global.
pub fn find_entity_manager_pointer(image: &ExeImage) -> Option<u32> {
    let mut state = false;

    for instr in image.decode_fn(find_lua_api_fn(image, c"EntityGetParent")?) {
        state = match state {
            false if instr.code() == Code::Call_rm32 => true,
            true if instr.code() == Code::Mov_r32_rm32 && instr.op0_register() == Register::ECX => {
                return Some(instr.memory_displacement32());
            }
            _ => false,
        };
    }
    None
}

/// Look for the `EntityHasTag` Lua API function and then look for the
/// second to last `CALL rel32` again, which is a call that accepts the
/// entity tag manager global in ECX
pub fn find_entity_tag_manager_pointer(image: &ExeImage) -> Option<u32> {
    let mut before_last_call_rel = None;
    let mut last_call_rel = None;

    let instrs = image
        .decode_fn(find_lua_api_fn(image, c"EntityHasTag")?)
        .enumerate()
        .map(|(i, instr)| {
            if instr.code() == Code::Call_rel32_32 {
                before_last_call_rel = last_call_rel;
                last_call_rel = Some(i);
            }
            instr
        })
        .collect::<Vec<_>>();

    instrs[..before_last_call_rel?]
        .iter()
        .rev()
        .find(|instr| instr.code() == Code::Mov_r32_rm32 && instr.op0_register() == Register::ECX)
        .map(|instr| instr.memory_displacement32())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::exe_image::PeHeader;

    use std::time::Instant;

    use anyhow::Context as _;
    use read_process_memory::Pid;
    use sysinfo::ProcessesToUpdate;
    use tracing::level_filters::LevelFilter;
    use tracing_subscriber::EnvFilter;

    #[test]
    fn test() -> anyhow::Result<()> {
        tracing_subscriber::fmt()
            .with_env_filter(
                EnvFilter::builder()
                    .with_default_directive(LevelFilter::DEBUG.into())
                    .from_env()?,
            )
            .init();

        let mut system = sysinfo::System::new();
        system.refresh_processes(ProcessesToUpdate::All);

        let Some(noita_pid) = system
            .processes_by_exact_name("noita.exe".as_ref())
            .find(|p| p.thread_kind().is_none())
            .map(|p| p.pid().as_u32() as Pid)
        else {
            eprintln!("Noita process not found");
            return Ok(());
        };

        let handle = noita_pid.try_into()?;
        let header = PeHeader::read(&handle)?;
        if header.timestamp() != 0x66ba59d6 {
            eprintln!("Timestamp mismatch: 0x{:x}", header.timestamp());
            return Ok(());
        }

        let instant = Instant::now();
        let image = header.read_image(&handle)?;
        println!("Image read in {:?}", instant.elapsed());

        let instant = Instant::now();

        let (world_seed, ng_plus) =
            find_seed_pointers(&image).context("Finding world seed and ng plus pointers")?;

        let kv_map_addr = find_stats_map_pointer(&image).context("Finding stats map pointer")?;

        // well we just pray they never touch GlobalStats
        // (which contains the kv_map as well as session/highest/global GameStats)
        // or GameStats layout lul
        let streak = kv_map_addr + 0x14; // global_stats.session.streaks
        let streak_pb = kv_map_addr + 0xdc; // global_stats.highest.streaks
        let deaths = kv_map_addr + 0x1a0; // global_stats.global.death_count

        let entity_manager =
            find_entity_manager_pointer(&image).context("Finding entity manager")?;

        let entity_tag_manager =
            find_entity_tag_manager_pointer(&image).context("Finding entity tag manager")?;

        println!("Pointers found in {:?}", instant.elapsed());

        println!("seed = 0x{world_seed:x}");
        println!("ng-plus-count = 0x{ng_plus:x}");

        println!();
        println!("stats-map = 0x{kv_map_addr:x}");
        println!("deaths = 0x{deaths:x}");
        println!("streak = 0x{streak:x}");
        println!("streak-pb = 0x{streak_pb:x}");

        println!();
        println!("entity-manager = 0x{entity_manager:x}");
        println!("entity-tag-manager = 0x{entity_tag_manager:x}");

        assert_eq!(kv_map_addr, 0x01206938);
        assert_eq!(entity_manager, 0x01202b78);
        assert_eq!(entity_tag_manager, 0x01204fbc);

        assert_eq!(world_seed, 0x01202fe4);
        assert_eq!(ng_plus, 0x01203004);
        assert_eq!(deaths, 0x01206ad8);
        assert_eq!(streak, 0x0120694c);
        assert_eq!(streak_pb, 0x01206a14);

        Ok(())
    }
}