smukx-Rust-for-Malware-Development/Enumeration/extract_modules.rs

/*
    Extract Modules From Running Process.
    By @5mukx 
*/

macro_rules! log {
    ($level:expr, $msg:expr) => {{
        let color = match $level {
            "DEBUG" => "\x1b[94m", 
            "INFO" => "\x1b[32m",  
            "WARN" => "\x1b[33m",  
            "ERROR" => "\x1b[31m", 
            _ => "\x1b[0m",        
        };

        match $level {
            "DEBUG" => println!("{}[||] {}\x1b[0m", color, $msg),
            "INFO" => println!("{}[+] {}\x1b[0m", color, $msg),
            "WARN" => println!("{}[!] {}\x1b[0m", color, $msg),
            "ERROR" => println!("{}[-] {}\x1b[0m", color, $msg),
            _ => println!("{}{}\x1b[0m", color, $msg),
        };
    }};
}


use std::{ffi::CString, mem};
use winapi::um::{
    errhandlingapi::GetLastError, handleapi::CloseHandle, processthreadsapi::OpenProcess, tlhelp32::{CreateToolhelp32Snapshot, Module32First, Module32Next, Process32First, Process32Next, MODULEENTRY32, PROCESSENTRY32, TH32CS_SNAPPROCESS
}};


fn get_pid(process_name: &str) -> u32{
    unsafe{
        let mut pe: PROCESSENTRY32 = std::mem::zeroed();
        pe.dwSize = mem::size_of::<PROCESSENTRY32>() as u32;

        let snap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
        if snap.is_null(){
            log!("ERROR", format!("Error while snapshoting processes : Error : {}",GetLastError()));
            std::process::exit(0);
        }

        let mut pid = 0;

        let mut result = Process32First(snap, &mut pe) != 0;

        while result{

            let exe_file = CString::from_vec_unchecked(pe.szExeFile
                .iter()
                .map(|&file| file as u8)
                .take_while(|&c| c!=0)
                .collect::<Vec<u8>>(),
            );

            if exe_file.to_str().unwrap() == process_name {
                pid = pe.th32ProcessID;
                break;
            }
            result = Process32Next(snap, &mut pe) !=0;
        }

        if pid == 0{
            log!("ERROR", format!("Unable to get PID for {}: {}",process_name , "PROCESS DOESNT EXISTS"));           
            std::process::exit(0);
        }
    
        CloseHandle(snap);
        pid
    }
}


fn main(){
    let args: Vec<String> = std::env::args().collect();
    if args.len() != 2{
        log!("ERROR", format!("Usage: ./importer.exe {}","<PROCESS NAME>"));
        return;
    }

    let pid = get_pid(&args[1]);

    unsafe{
        let process_handle = OpenProcess(
            0x040 | 0x0010,
            0,
            pid,
        );

        if process_handle.is_null(){
            log!("ERROR", format!("Failed to open process: {}", GetLastError()));
            return;
        }

        let mut me32: MODULEENTRY32 = std::mem::zeroed();
        me32.dwSize = std::mem::size_of::<MODULEENTRY32>() as u32;

        let snapshot = CreateToolhelp32Snapshot(
            // TH32CS_SNAPMODULE | TH32CS_SNAPMODULE32, 
            0x00000008 | 0x00000010,
            pid
        );

        if snapshot.is_null(){
            log!("ERROR", format!("Failed to create SnapShot: {}",GetLastError()));
            CloseHandle(process_handle);
            return;
        }
        
        let module_info = Module32First(snapshot, &mut me32);

        if module_info != 0 {
            println!("{:<34} | {:<20}", "Module Name", "Base Address");
            println!("{:-<50}", "");

            loop {
                // printing module name along with its base address.
                let module_name = std::ffi::CStr::from_ptr(me32.szModule.as_ptr());
                log!(
                    "INFO",
                    format!("{:<30} | 0x{:016x}",
                    module_name.to_str().unwrap_or("Unknown Conversion"),
                    me32.modBaseAddr as usize
                ));

                if Module32Next(snapshot, &mut me32) == 0 {
                    break;
                }
            }
        } else {
            eprintln!("Failed to get first module: {}", GetLastError());
        }

        CloseHandle(snapshot);
        CloseHandle(process_handle);
    }
}