use std::fs;

fn main() {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        eprintln!("Usage: dump_iat <exe_path>");
        return;
    }
    
    let data = fs::read(&args[1]).expect("Failed to read file");
    
    // 找到 Import Directory RVA (offset 0x88 in optional header for PE32+)
    let pe_offset = u32::from_le_bytes([data[0x3C], data[0x3D], data[0x3E], data[0x3F]]) as usize;
    let opt_header_offset = pe_offset + 24;
    
    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],
    ]);
    
    println!("Import Directory RVA: 0x{:X}", import_dir_rva);
    
    if import_dir_rva == 0 {
        println!("No import directory!");
        return;
    }
    
    // RVA to file offset
    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;
    
    fn rva_to_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_offset) = rva_to_offset(&data, import_dir_rva) {
        println!("Import Directory file offset: 0x{:X}", import_offset);
        
        // Read first IMPORT_DESCRIPTOR
        let orig_first_thunk = u32::from_le_bytes([
            data[import_offset],
            data[import_offset + 1],
            data[import_offset + 2],
            data[import_offset + 3],
        ]);
        
        let first_thunk = u32::from_le_bytes([
            data[import_offset + 16],
            data[import_offset + 17],
            data[import_offset + 18],
            data[import_offset + 19],
        ]);
        
        println!("\nIMPORT_DESCRIPTOR:");
        println!("  OriginalFirstThunk (INT): 0x{:X}", orig_first_thunk);
        println!("  FirstThunk (IAT): 0x{:X}", first_thunk);
        
        if let Some(int_offset) = rva_to_offset(&data, orig_first_thunk) {
            println!("\nINT (OriginalFirstThunk) at file offset 0x{:X}:", int_offset);
            for i in 0..3 {
                let offset = int_offset + i * 8;
                let value = u64::from_le_bytes([
                    data[offset],
                    data[offset + 1],
                    data[offset + 2],
                    data[offset + 3],
                    data[offset + 4],
                    data[offset + 5],
                    data[offset + 6],
                    data[offset + 7],
                ]);
                println!("  [{}] 0x{:016X}", i, value);
            }
        }
        
        if let Some(iat_offset) = rva_to_offset(&data, first_thunk) {
            println!("\nIAT (FirstThunk) at file offset 0x{:X}:", iat_offset);
            println!("  Raw bytes:");
            for i in 0..3 {
                let offset = iat_offset + i * 8;
                let bytes = &data[offset..offset + 8];
                let value = u64::from_le_bytes([
                    bytes[0], bytes[1], bytes[2], bytes[3],
                    bytes[4], bytes[5], bytes[6], bytes[7],
                ]);
                println!("  [{}] {:02X} {:02X} {:02X} {:02X} {:02X} {:02X} {:02X} {:02X} = 0x{:016X}",
                         i, bytes[0], bytes[1], bytes[2], bytes[3], 
                         bytes[4], bytes[5], bytes[6], bytes[7], value);
            }
        }
    }
}
