1 /* 2 * PE Fixup Utility 3 * Copyright (C) 2005 Filip Navara 4 * Copyright (C) 2020 Mark Jansen 5 * 6 * The purpose of this utility is fix PE binaries generated by binutils and 7 * to manipulate flags that can't be set by binutils. 8 * 9 * Currently one features is implemented: 10 * 11 * - Updating the PE header to use a LOAD_CONFIG, 12 * when the struct is exported with the name '_load_config_used' 13 */ 14 15 #include <stdio.h> 16 #include <stdlib.h> 17 #include <string.h> 18 19 // host_includes 20 #include <typedefs.h> 21 #include <pecoff.h> 22 #include "../../dll/win32/dbghelp/compat.h" 23 24 static const char* g_ApplicationName; 25 static const char* g_Target; 26 27 void *rva_to_ptr(unsigned char *buffer, PIMAGE_NT_HEADERS nt_header, DWORD rva) 28 { 29 unsigned int i; 30 PIMAGE_SECTION_HEADER section_header = IMAGE_FIRST_SECTION(nt_header); 31 32 for (i = 0; i < nt_header->FileHeader.NumberOfSections; i++, section_header++) 33 { 34 if (rva >= section_header->VirtualAddress && 35 rva < section_header->VirtualAddress + section_header->Misc.VirtualSize) 36 { 37 return buffer + rva - section_header->VirtualAddress + section_header->PointerToRawData; 38 } 39 } 40 41 return NULL; 42 } 43 44 static void error(const char* message, ...) 45 { 46 va_list args; 47 48 fprintf(stderr, "%s ERROR: '%s': ", g_ApplicationName, g_Target); 49 50 va_start(args, message); 51 fprintf(stderr, message, args); 52 va_end(args); 53 } 54 55 static void fix_checksum(unsigned char *buffer, long len, PIMAGE_NT_HEADERS nt_header) 56 { 57 unsigned int checksum = 0; 58 long n; 59 60 nt_header->OptionalHeader.CheckSum = 0; 61 62 for (n = 0; n < len; n += 2) 63 { 64 checksum += *(unsigned short *)(buffer + n); 65 checksum = (checksum + (checksum >> 16)) & 0xffff; 66 } 67 68 checksum += len; 69 nt_header->OptionalHeader.CheckSum = checksum; 70 } 71 72 static int add_loadconfig(unsigned char *buffer, PIMAGE_NT_HEADERS nt_header) 73 { 74 PIMAGE_DATA_DIRECTORY export_dir; 75 PIMAGE_EXPORT_DIRECTORY export_directory; 76 77 export_dir = &nt_header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT]; 78 if (export_dir->Size != 0) 79 { 80 export_directory = rva_to_ptr(buffer, nt_header, export_dir->VirtualAddress); 81 if (export_directory != NULL) 82 { 83 DWORD *name_ptr, *function_ptr, n; 84 WORD *ordinal_ptr; 85 86 name_ptr = rva_to_ptr(buffer, nt_header, export_directory->AddressOfNames); 87 ordinal_ptr = rva_to_ptr(buffer, nt_header, export_directory->AddressOfNameOrdinals); 88 function_ptr = rva_to_ptr(buffer, nt_header, export_directory->AddressOfFunctions); 89 90 for (n = 0; n < export_directory->NumberOfNames; n++) 91 { 92 const char* name = rva_to_ptr(buffer, nt_header, name_ptr[n]); 93 if (!strcmp(name, "_load_config_used")) 94 { 95 PIMAGE_DATA_DIRECTORY load_config_dir; 96 DWORD load_config_rva = function_ptr[ordinal_ptr[n]]; 97 DWORD* load_config_ptr = rva_to_ptr(buffer, nt_header, load_config_rva); 98 99 /* Update the DataDirectory pointer / size 100 The first entry of the LOAD_CONFIG struct is the size, use that as DataDirectory.Size */ 101 load_config_dir = &nt_header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG]; 102 load_config_dir->VirtualAddress = load_config_rva; 103 load_config_dir->Size = *load_config_ptr; 104 105 return 0; 106 } 107 } 108 109 error("Export '_load_config_used' not found\n"); 110 } 111 else 112 { 113 error("Invalid rva for export directory\n"); 114 } 115 } 116 else 117 { 118 error("No export directory\n"); 119 } 120 121 return 1; 122 } 123 124 int main(int argc, char **argv) 125 { 126 FILE* file; 127 long len; 128 unsigned char *buffer; 129 PIMAGE_DOS_HEADER dos_header; 130 int result = 1; 131 132 g_ApplicationName = argv[0]; 133 134 if (argc < 2) 135 { 136 printf("Usage: %s <filename>\n", g_ApplicationName); 137 return 1; 138 } 139 140 g_Target = argv[1]; 141 142 /* Read the whole file to memory. */ 143 file = fopen(argv[1], "r+b"); 144 if (!file) 145 { 146 fprintf(stderr, "%s ERROR: Can't open '%s'.\n", g_ApplicationName, g_Target); 147 return 1; 148 } 149 150 fseek(file, 0, SEEK_END); 151 len = ftell(file); 152 if (len < sizeof(IMAGE_DOS_HEADER)) 153 { 154 fclose(file); 155 error("Image size too small to be a PE image\n"); 156 return 1; 157 } 158 159 /* Add one byte extra for the case where the input file size is odd. 160 We rely on this in our crc calculation */ 161 buffer = calloc(len + 1, 1); 162 if (buffer == NULL) 163 { 164 fclose(file); 165 error("Not enough memory available: (Needed %u bytes).\n", len + 1); 166 return 1; 167 } 168 169 /* Read the whole input file into a buffer */ 170 fseek(file, 0, SEEK_SET); 171 fread(buffer, 1, len, file); 172 173 /* Check the headers and save pointers to them. */ 174 dos_header = (PIMAGE_DOS_HEADER)buffer; 175 if (dos_header->e_magic == IMAGE_DOS_SIGNATURE) 176 { 177 PIMAGE_NT_HEADERS nt_header; 178 179 nt_header = (PIMAGE_NT_HEADERS)(buffer + dos_header->e_lfanew); 180 181 if (nt_header->Signature == IMAGE_NT_SIGNATURE) 182 { 183 if (!add_loadconfig(buffer, nt_header)) 184 { 185 fix_checksum(buffer, len, nt_header); 186 187 /* We could 'optimize by only writing the changed parts, but keep it simple for now */ 188 fseek(file, 0, SEEK_SET); 189 fwrite(buffer, 1, len, file); 190 191 /* Success */ 192 result = 0; 193 } 194 else 195 { 196 /* Error already printed inside add_loadconfig */ 197 } 198 } 199 else 200 { 201 error("Invalid PE signature: %x\n", nt_header->Signature); 202 } 203 } 204 else 205 { 206 error("Invalid DOS signature: %x\n", dos_header->e_magic); 207 } 208 209 free(buffer); 210 fclose(file); 211 212 return result; 213 } 214