use std::fs;

fn main() {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        eprintln!("Usage: verify_pe <exe_path>");
        return;
    }
    
    let exe_path = &args[1];
    let data = match fs::read(exe_path) {
        Ok(d) => d,
        Err(e) => {
            eprintln!("Failed to read {}: {}", exe_path, e);
            return;
        }
    };
    
    // 解析 PE 头部
    let dos_sig = u16::from_le_bytes([data[0], data[1]]);
    if dos_sig != 0x5A4D {
        eprintln!("Not a valid PE file");
        return;
    }
    
    let pe_offset = u32::from_le_bytes([data[0x3C], data[0x3D], data[0x3E], data[0x3F]]) as usize;
    
    let pe_sig = u32::from_le_bytes([
        data[pe_offset],
        data[pe_offset + 1],
        data[pe_offset + 2],
        data[pe_offset + 3],
    ]);
    if pe_sig != 0x00004550 {
        eprintln!("Invalid PE signature");
        return;
    }
    
    // 可选头部偏移
    let opt_header_offset = pe_offset + 24;
    
    // Data Directory
    let import_dir_rva = u32::from_le_bytes([
        data[opt_header_offset + 120],
        data[opt_header_offset + 121],
        data[opt_header_offset + 122],
        data[opt_header_offset + 123],
    ]);
    
    let import_dir_size = u32::from_le_bytes([
        data[opt_header_offset + 124],
        data[opt_header_offset + 125],
        data[opt_header_offset + 126],
        data[opt_header_offset + 127],
    ]);
    
    println!("Import Directory RVA: 0x{:X}", import_dir_rva);
    println!("Import Directory Size: 0x{:X}", import_dir_size);
    
    if import_dir_rva == 0 {
        eprintln!("No import directory!");
        return;
    }
    
    // 将 RVA 转换为文件偏移
    fn rva_to_file_offset(data: &[u8], rva: u32) -> Option<usize> {
        let pe_offset = u32::from_le_bytes([data[0x3C], data[0x3D], data[0x3E], data[0x3F]]) as usize;
        let num_sections = u16::from_le_bytes([data[pe_offset + 6], data[pe_offset + 7]]) as usize;
        let opt_header_size = u16::from_le_bytes([data[pe_offset + 20], data[pe_offset + 21]]) as usize;
        let section_table_offset = pe_offset + 24 + opt_header_size;
        
        for i in 0..num_sections {
            let sec_offset = section_table_offset + i * 40;
            let sec_rva = u32::from_le_bytes([
                data[sec_offset + 12],
                data[sec_offset + 13],
                data[sec_offset + 14],
                data[sec_offset + 15],
            ]);
            let sec_size = u32::from_le_bytes([
                data[sec_offset + 8],
                data[sec_offset + 9],
                data[sec_offset + 10],
                data[sec_offset + 11],
            ]);
            let sec_ptr = u32::from_le_bytes([
                data[sec_offset + 20],
                data[sec_offset + 21],
                data[sec_offset + 22],
                data[sec_offset + 23],
            ]);
            
            if rva >= sec_rva && rva < sec_rva + sec_size {
                return Some((rva - sec_rva + sec_ptr) as usize);
            }
        }
        None
    }
    
    if let Some(import_file_offset) = rva_to_file_offset(&data, import_dir_rva) {
        println!("Import Directory file offset: 0x{:X}", import_file_offset);
        
        // 读取 IMPORT_DESCRIPTOR
        let mut offset = import_file_offset;
        let mut dll_count = 0;
        
        loop {
            let orig_first_thunk = u32::from_le_bytes([
                data[offset],
                data[offset + 1],
                data[offset + 2],
                data[offset + 3],
            ]);
            
            let time_stamp = u32::from_le_bytes([
                data[offset + 4],
                data[offset + 5],
                data[offset + 6],
                data[offset + 7],
            ]);
            
            let name_rva = u32::from_le_bytes([
                data[offset + 12],
                data[offset + 13],
                data[offset + 14],
                data[offset + 15],
            ]);
            
            let first_thunk = u32::from_le_bytes([
                data[offset + 16],
                data[offset + 17],
                data[offset + 18],
                data[offset + 19],
            ]);
            
            if orig_first_thunk == 0 && first_thunk == 0 {
                break;
            }
            
            let dll_name_offset = rva_to_file_offset(&data, name_rva).unwrap();
            let dll_name_end = data[dll_name_offset..].iter().position(|&b| b == 0).unwrap();
            let dll_name = String::from_utf8_lossy(&data[dll_name_offset..dll_name_offset + dll_name_end]);
            
            println!("\nDLL: {}", dll_name);
            println!("  OriginalFirstThunk (INT): 0x{:X}", orig_first_thunk);
            println!("  FirstThunk (IAT): 0x{:X}", first_thunk);
            
            // 读取 IAT 条目
            if let Some(iat_offset) = rva_to_file_offset(&data, first_thunk) {
                let mut entry_idx = 0;
                loop {
                    let entry_offset = iat_offset + entry_idx * 8;
                    let entry = u64::from_le_bytes([
                        data[entry_offset],
                        data[entry_offset + 1],
                        data[entry_offset + 2],
                        data[entry_offset + 3],
                        data[entry_offset + 4],
                        data[entry_offset + 5],
                        data[entry_offset + 6],
                        data[entry_offset + 7],
                    ]);
                    
                    if entry == 0 {
                        break;
                    }
                    
                    // 如果最高位是 0，这是指向 Hint/Name 的 RVA
                    if entry & 0x8000000000000000 == 0 {
                        let hint_name_rva = entry as u32;
                        if let Some(hn_offset) = rva_to_file_offset(&data, hint_name_rva) {
                            let hint = u16::from_le_bytes([data[hn_offset], data[hn_offset + 1]]);
                            let name_end = data[hn_offset + 2..].iter().position(|&b| b == 0).unwrap();
                            let func_name = String::from_utf8_lossy(&data[hn_offset + 2..hn_offset + 2 + name_end]);
                            println!("    [{}] IAT=0x{:016X} -> Hint/Name RVA=0x{:X} -> '{}'", 
                                     entry_idx, entry, hint_name_rva, func_name);
                        } else {
                            println!("    [{}] IAT=0x{:016X} -> Hint/Name RVA=0x{:X} (INVALID)", 
                                     entry_idx, entry, hint_name_rva);
                        }
                    } else {
                        println!("    [{}] IAT=0x{:016X} (ORDINAL)", entry_idx, entry);
                    }
                    
                    entry_idx += 1;
                }
            }
            
            dll_count += 1;
            offset += 20;
        }
        
        println!("\nTotal DLLs: {}", dll_count);
    } else {
        eprintln!("Import directory RVA not found in any section!");
    }
}
