#!/usr/bin/env python
#
# Copyright(C) 2016 Simon Howard
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#

#
# This script copies a Windows .exe file from a source to a destination
# (like cp); however, it also uses the binutils objdump command to
# recursively determine all its DLL dependencies and copy those too.
#

import argparse
import os
import re
import shlex
import shutil
import sys
import subprocess

DLL_NAME_RE = re.compile('\s+DLL Name: (.*\.dll)')

# DLLs that are bundled with Windows, and we don't need to copy.
# Thanks to Martin Preisler; this is from mingw-bundledlls.
WIN32_DLLS = {
	"advapi32.dll", "kernel32.dll", "msvcrt.dll", "ole32.dll",
        "user32.dll", "ws2_32.dll", "comdlg32.dll", "gdi32.dll", "imm32.dll",
        "oleaut32.dll", "shell32.dll", "winmm.dll", "winspool.drv",
        "wldap32.dll", "ntdll.dll", "d3d9.dll", "mpr.dll", "crypt32.dll",
        "dnsapi.dll", "shlwapi.dll", "version.dll", "iphlpapi.dll",
        "msimg32.dll", "setupapi.dll",
}

parser = argparse.ArgumentParser(description='Copy EXE with DLLs.')
parser.add_argument('--objdump', type=str, default='objdump',
                    help='name of objdump binary to invoke')
parser.add_argument('--dll_path', type=str, nargs='*', default=(),
                    help='list of paths to check for DLLs')
parser.add_argument('--ldflags', type=str, default="",
                    help='linker flags, which can be used to automatically '
                         'determine the DLL path')
parser.add_argument('source', type=str,
                    help='path to binary to copy')
parser.add_argument('destination', type=str,
                    help='destination to copy binary')

def file_dependencies(filename, objdump):
	"""Get the direct DLL dependencies of the given file.

	The DLLs depended on by the given file are determined by invoking
	the provided objdump binary. DLLs which are part of the standard
	Win32 API are excluded from the result.

	Args:
	  filename: Path to a file to query.
	  objdump: Name of objdump binary to invoke.
	Returns:
	  List of filenames (not fully qualified paths).
	"""
	cmd = subprocess.Popen([objdump, '-p', filename],
	                       stdout=subprocess.PIPE)
	try:
		result = []
		for line in cmd.stdout:
			m = DLL_NAME_RE.match(line.decode())
			if m:
				dll = m.group(1)
				if dll.lower() not in WIN32_DLLS:
					result.append(dll)

		return result
	finally:
		cmd.wait()
		assert cmd.returncode == 0, (
			'%s invocation for %s exited with %d' % (
				objdump, filename, cmd.returncode))

def find_dll(filename, dll_paths):
	"""Search the given list of paths for a DLL with the given name."""
	for path in dll_paths:
		full_path = os.path.join(path, filename)
		if os.path.exists(full_path):
			return full_path

	raise IOError('DLL file %s not found in path: %s' % (
		filename, dll_paths))

def all_dependencies(filename, objdump, dll_paths):
	"""Recursively find all dependencies of the given executable file.

	Args:
	  filename: Executable file to examine.
	  objdump: Command to invoke to find dependencies.
	  dll_paths: List of directories to search for DLL files.
	Returns:
	  Tuple containing:
	    Set containing paths to all DLL dependencies of the file.
	    Set of filenames of missing DLLs.
	"""
	result, missing = set(), set()
	to_process = {filename}
	while len(to_process) > 0:
		filename = to_process.pop()
		for dll in file_dependencies(filename, objdump):
			try:
				dll = find_dll(dll, dll_paths)
				if dll not in result:
					result |= {dll}
					to_process |= {dll}
			except IOError as e:
				missing |= {dll}

	return result, missing

def get_dll_path():
	"""Examine command line arguments and determine the DLL search path.

	If the --path argument is provided, paths from this are added.
	Furthermore, if --ldflags is provided, this is interpreted as a list of
	linker flags and the -L paths are used to find associated paths that are
	likely to contain DLLs, with the assumption that autotools usually
	installs DLLs to ${prefix}/bin when installing Unix-style libraries into
	${prefix}/lib.

	Returns:
	  List of filesystem paths to check for DLLs.
	"""
	result = set(args.dll_path)

	if args.ldflags != '':
		for arg in shlex.split(args.ldflags):
			if arg.startswith("-L"):
				prefix, libdir = os.path.split(arg[2:])
				if libdir != "lib":
					continue
				bindir = os.path.join(prefix, "bin")
				if os.path.exists(bindir):
					result |= {bindir}

	return list(result)

args = parser.parse_args()

dll_path = get_dll_path()
dll_files, missing = all_dependencies(args.source, args.objdump, dll_path)

# Exit with failure if DLLs are missing.
if missing:
	sys.stderr.write("Missing DLLs not found in path %s:\n" % (dll_path,))
	for filename in missing:
		sys.stderr.write("\t%s\n" % filename)
	sys.exit(1)

# Destination may be a full path (rename) or may be a directory to copy into:
#   cp foo.exe bar/baz.exe
#   cp foo.exe bar/
if os.path.isdir(args.destination):
	dest_dir = args.destination
else:
	dest_dir = os.path.dirname(args.destination)

# Copy .exe and DLLs.
shutil.copy(args.source, args.destination)
for filename in dll_files:
	shutil.copy(filename, dest_dir)

