reorder-imports.py (4014B)
1 #!/usr/bin/env python3 2 # Copyright 2014 The Chromium Authors 3 # Use of this source code is governed by a BSD-style license that can be 4 # found in the LICENSE file. 5 6 import glob 7 import optparse 8 import os 9 import shutil 10 import sys 11 12 sys.path.insert( 13 0, 14 os.path.join(os.path.dirname(__file__), '..', '..', 'third_party', 15 'pefile_py3')) 16 import pefile 17 18 def reorder_imports(input_dir, output_dir, architecture): 19 """Swap chrome_elf.dll to be the first import of chrome.exe. 20 Also copy over any related files that might be needed 21 (pdbs, manifests etc.). 22 """ 23 # TODO(thakis): See if there is a reliable way to write the 24 # correct executable in the first place, so that this script 25 # only needs to verify that and not write a whole new exe. 26 27 input_image = os.path.join(input_dir, 'chrome.exe') 28 output_image = os.path.join(output_dir, 'chrome.exe') 29 30 # pefile mmap()s the whole executable, and then parses parts of 31 # it into python data structures for ease of processing. 32 # To write the file again, only the mmap'd data is written back, 33 # so modifying the parsed python objects generally has no effect. 34 # However, parsed raw data ends up in pe.Structure instances, 35 # and these all get serialized back when the file gets written. 36 # So things that are in a Structure must have their data set 37 # through the Structure, while other data must bet set through 38 # the set_bytes_*() methods. 39 pe = pefile.PE(input_image, fast_load=True) 40 if architecture == 'x64' or architecture == 'arm64': 41 assert pe.PE_TYPE == pefile.OPTIONAL_HEADER_MAGIC_PE_PLUS 42 else: 43 assert pe.PE_TYPE == pefile.OPTIONAL_HEADER_MAGIC_PE 44 45 pe.parse_data_directories(directories=[ 46 pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_IMPORT']]) 47 48 found_elf = False 49 for i, peimport in enumerate(pe.DIRECTORY_ENTRY_IMPORT): 50 if peimport.dll.lower() == b'chrome_elf.dll': 51 assert not found_elf, 'only one chrome_elf.dll import expected' 52 found_elf = True 53 if i > 0: 54 swap = pe.DIRECTORY_ENTRY_IMPORT[0] 55 56 # Morally we want to swap peimport.struct and swap.struct here, 57 # but the pe module doesn't expose a public method on Structure 58 # to get all data of a Structure without explicitly listing all 59 # field names. 60 # NB: OriginalFirstThunk and Characteristics are an union both at 61 # offset 0, handling just one of them is enough. 62 peimport.struct.OriginalFirstThunk, swap.struct.OriginalFirstThunk = \ 63 swap.struct.OriginalFirstThunk, peimport.struct.OriginalFirstThunk 64 peimport.struct.TimeDateStamp, swap.struct.TimeDateStamp = \ 65 swap.struct.TimeDateStamp, peimport.struct.TimeDateStamp 66 peimport.struct.ForwarderChain, swap.struct.ForwarderChain = \ 67 swap.struct.ForwarderChain, peimport.struct.ForwarderChain 68 peimport.struct.Name, swap.struct.Name = \ 69 swap.struct.Name, peimport.struct.Name 70 peimport.struct.FirstThunk, swap.struct.FirstThunk = \ 71 swap.struct.FirstThunk, peimport.struct.FirstThunk 72 assert found_elf, 'chrome_elf.dll import not found' 73 74 pe.write(filename=output_image) 75 76 for fname in glob.iglob(os.path.join(input_dir, 'chrome.exe.*')): 77 shutil.copy(fname, os.path.join(output_dir, os.path.basename(fname))) 78 return 0 79 80 81 def main(argv): 82 usage = 'reorder_imports.py -i <input_dir> -o <output_dir> -a <target_arch>' 83 parser = optparse.OptionParser(usage=usage) 84 parser.add_option('-i', '--input', help='reorder chrome.exe in DIR', 85 metavar='DIR') 86 parser.add_option('-o', '--output', help='write new chrome.exe to DIR', 87 metavar='DIR') 88 parser.add_option('-a', '--arch', help='architecture of build (optional)', 89 default='ia32') 90 opts, args = parser.parse_args() 91 92 if not opts.input or not opts.output: 93 parser.error('Please provide and input and output directory') 94 return reorder_imports(opts.input, opts.output, opts.arch) 95 96 if __name__ == "__main__": 97 sys.exit(main(sys.argv[1:]))