cs2-dumper/src/remote/process.rs

use std::ffi::{c_void, CStr};
use std::mem;
use std::ptr;

use windows::Win32::Foundation::*;
use windows::Win32::System::Diagnostics::Debug::*;
use windows::Win32::System::Diagnostics::ToolHelp::*;
use windows::Win32::System::Threading::*;

use crate::error::{Error, Result};

use super::Module;

pub struct Process {
    process_id: u32,
    process_handle: HANDLE,
}

impl Process {
    pub fn new(process_name: &str) -> Result<Self> {
        let process_id = Self::get_process_id_by_name(process_name)?;

        let process_handle = unsafe { OpenProcess(PROCESS_ALL_ACCESS, false, process_id) }?;

        Ok(Self {
            process_id,
            process_handle,
        })
    }

    pub fn find_pattern(&self, module_name: &str, pattern: &str) -> Result<usize> {
        let module = self.get_module_by_name(module_name)?;

        let mut module_data: Vec<u8> = vec![0; module.size() as usize];

        self.read_memory_raw(
            module.address(),
            module_data.as_mut_ptr() as *mut _,
            module_data.len(),
        )?;

        let pattern_bytes = Self::pattern_to_bytes(pattern);

        for i in 0..module.size() as usize - pattern.len() {
            let mut found = true;

            for j in 0..pattern_bytes.len() {
                if module_data[i + j] != pattern_bytes[j] as u8 && pattern_bytes[j] != -1 {
                    found = false;

                    break;
                }
            }

            if found {
                return Ok(module.address() + i);
            }
        }

        Err(Error::PatternNotFound)
    }

    pub fn get_loaded_modules(&self) -> Result<Vec<String>> {
        let snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, self.process_id) }?;

        let mut entry = MODULEENTRY32 {
            dwSize: mem::size_of::<MODULEENTRY32>() as u32,
            ..Default::default()
        };

        let mut modules = Vec::new();

        unsafe {
            Module32First(snapshot, &mut entry)?;

            while Module32Next(snapshot, &mut entry).is_ok() {
                let name = CStr::from_ptr(&entry.szModule as *const _ as *const _)
                    .to_string_lossy()
                    .into_owned();

                modules.push(name);
            }
        }

        Ok(modules)
    }

    pub fn get_module_by_name(&self, module_name: &str) -> Result<Module> {
        let snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, self.process_id) }?;

        let mut entry = MODULEENTRY32 {
            dwSize: mem::size_of::<MODULEENTRY32>() as u32,
            ..Default::default()
        };

        unsafe {
            Module32First(snapshot, &mut entry)?;

            while Module32Next(snapshot, &mut entry).is_ok() {
                let name = CStr::from_ptr(&entry.szModule as *const _ as *const _)
                    .to_string_lossy()
                    .into_owned();

                if name == module_name {
                    return Module::new(self, entry.modBaseAddr as usize);
                }
            }
        }

        Err(Error::ModuleNotFound)
    }

    pub fn read_memory_raw(&self, address: usize, buffer: *mut c_void, size: usize) -> Result<()> {
        unsafe {
            ReadProcessMemory(
                self.process_handle,
                address as *const _,
                buffer,
                size,
                Some(ptr::null_mut()),
            )
        }
        .map_err(Into::into)
    }

    pub fn write_memory_raw(
        &self,
        address: usize,
        buffer: *const c_void,
        size: usize,
    ) -> Result<()> {
        unsafe {
            WriteProcessMemory(
                self.process_handle,
                address as *const _,
                buffer,
                size,
                Some(ptr::null_mut()),
            )
        }
        .map_err(Into::into)
    }

    pub fn read_memory<T>(&self, address: usize) -> Result<T> {
        let mut buffer: T = unsafe { mem::zeroed() };

        self.read_memory_raw(
            address,
            &mut buffer as *const _ as *mut _,
            mem::size_of::<T>(),
        )?;

        Ok(buffer)
    }

    pub fn write_memory<T>(&self, address: usize, value: T) -> Result<()> {
        self.write_memory_raw(address, &value as *const _ as *const _, mem::size_of::<T>())
    }

    pub fn read_string(&self, address: usize, length: usize) -> Result<String> {
        let mut buffer: Vec<u8> = vec![0; length];

        self.read_memory_raw(address, buffer.as_mut_ptr() as *mut _, length)?;

        if let Some(end) = buffer.iter().position(|&x| x == 0) {
            buffer.truncate(end);
        }

        Ok(String::from_utf8(buffer)?)
    }

    pub fn resolve_jmp(&self, address: usize) -> Result<usize> {
        let displacement = self.read_memory::<i32>(address + 0x1)?;

        Ok((address + 0x5) + displacement as usize)
    }

    pub fn resolve_rip(&self, address: usize) -> Result<usize> {
        let displacement = self.read_memory::<i32>(address + 0x3)?;

        Ok((address + 0x7) + displacement as usize)
    }

    fn get_process_id_by_name(process_name: &str) -> Result<u32> {
        let snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) }?;

        let mut entry = PROCESSENTRY32 {
            dwSize: mem::size_of::<PROCESSENTRY32>() as u32,
            ..Default::default()
        };

        unsafe {
            Process32First(snapshot, &mut entry)?;

            while Process32Next(snapshot, &mut entry).is_ok() {
                let name = CStr::from_ptr(&entry.szExeFile as *const _ as *const _)
                    .to_string_lossy()
                    .into_owned();

                if name == process_name {
                    return Ok(entry.th32ProcessID);
                }
            }
        }

        Err(Error::ProcessNotFound)
    }

    fn pattern_to_bytes(pattern: &str) -> Vec<i32> {
        let mut bytes = Vec::new();

        let chars: Vec<char> = pattern.chars().collect();

        let mut i = 0;

        while i < chars.len() {
            if chars[i] == ' ' {
                i += 1;

                continue;
            }

            if chars[i] == '?' {
                bytes.push(-1);

                i += 1;

                continue;
            }

            if i + 1 < chars.len() {
                let value = i32::from_str_radix(&pattern[i..i + 2], 16);

                match value {
                    Ok(v) => bytes.push(v),
                    Err(_) => {}
                }

                i += 1;
            }

            i += 1;
        }

        bytes
    }
}

impl Drop for Process {
    fn drop(&mut self) {
        if !self.process_handle.is_invalid() {
            unsafe { CloseHandle(self.process_handle).unwrap() }
        }
    }
}