#!/usr/bin/env python3
# Copyright 2014 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import glob
import optparse
import os
import shutil
import sys

sys.path.insert(
    0,
    os.path.join(
        os.path.dirname(__file__), '..', '..', 'third_party', 'pefile_py3'
    ),
)
import pefile


def reorder_imports(input_dir, output_dir, architecture):
    """Swap chrome_elf.dll to be the first import of chrome.exe.
    Also copy over any related files that might be needed
    (pdbs, manifests etc.).
    """
    # TODO(thakis): See if there is a reliable way to write the
    # correct executable in the first place, so that this script
    # only needs to verify that and not write a whole new exe.

    input_image = os.path.join(input_dir, 'chrome.exe')
    output_image = os.path.join(output_dir, 'chrome.exe')

    # pefile mmap()s the whole executable, and then parses parts of
    # it into python data structures for ease of processing.
    # To write the file again, only the mmap'd data is written back,
    # so modifying the parsed python objects generally has no effect.
    # However, parsed raw data ends up in pe.Structure instances,
    # and these all get serialized back when the file gets written.
    # So things that are in a Structure must have their data set
    # through the Structure, while other data must bet set through
    # the set_bytes_*() methods.
    pe = pefile.PE(input_image, fast_load=True)
    if architecture == 'x64' or architecture == 'arm64':
        assert pe.PE_TYPE == pefile.OPTIONAL_HEADER_MAGIC_PE_PLUS
    else:
        assert pe.PE_TYPE == pefile.OPTIONAL_HEADER_MAGIC_PE

    pe.parse_data_directories(
        directories=[pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_IMPORT']]
    )

    found_elf = False
    for i, peimport in enumerate(pe.DIRECTORY_ENTRY_IMPORT):
        if peimport.dll.lower() == b'chrome_elf.dll':
            assert not found_elf, 'only one chrome_elf.dll import expected'
            found_elf = True
            if i > 0:
                swap = pe.DIRECTORY_ENTRY_IMPORT[0]

                # Morally we want to swap peimport.struct and swap.struct here,
                # but the pe module doesn't expose a public method on Structure
                # to get all data of a Structure without explicitly listing all
                # field names.
                # NB: OriginalFirstThunk and Characteristics are an union both at
                # offset 0, handling just one of them is enough.
                (
                    peimport.struct.OriginalFirstThunk,
                    swap.struct.OriginalFirstThunk,
                ) = (
                    swap.struct.OriginalFirstThunk,
                    peimport.struct.OriginalFirstThunk,
                )
                peimport.struct.TimeDateStamp, swap.struct.TimeDateStamp = (
                    swap.struct.TimeDateStamp,
                    peimport.struct.TimeDateStamp,
                )
                peimport.struct.ForwarderChain, swap.struct.ForwarderChain = (
                    swap.struct.ForwarderChain,
                    peimport.struct.ForwarderChain,
                )
                peimport.struct.Name, swap.struct.Name = (
                    swap.struct.Name,
                    peimport.struct.Name,
                )
                peimport.struct.FirstThunk, swap.struct.FirstThunk = (
                    swap.struct.FirstThunk,
                    peimport.struct.FirstThunk,
                )
    assert found_elf, 'chrome_elf.dll import not found'

    pe.write(filename=output_image)

    for fname in glob.iglob(os.path.join(input_dir, 'chrome.exe.*')):
        shutil.copy(fname, os.path.join(output_dir, os.path.basename(fname)))
    return 0


def main(argv):
    usage = 'reorder_imports.py -i <input_dir> -o <output_dir> -a <target_arch>'
    parser = optparse.OptionParser(usage=usage)
    parser.add_option(
        '-i', '--input', help='reorder chrome.exe in DIR', metavar='DIR'
    )
    parser.add_option(
        '-o', '--output', help='write new chrome.exe to DIR', metavar='DIR'
    )
    parser.add_option(
        '-a', '--arch', help='architecture of build (optional)', default='ia32'
    )
    opts, args = parser.parse_args()

    if not opts.input or not opts.output:
        parser.error('Please provide and input and output directory')
    return reorder_imports(opts.input, opts.output, opts.arch)


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))