1#!/usr/bin/env python
2# Copyright (c) 2012 The Chromium Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""A utility script that can extract and edit resources in a Windows binary.
7
8For detailed help, see the script's usage by invoking it with --help."""
9
10import ctypes
11import ctypes.wintypes
12import logging
13import optparse
14import os
15import shutil
16import sys
17import tempfile
18import win32api
19import win32con
20
21
22_LOGGER = logging.getLogger(__name__)
23
24
25# The win32api-supplied UpdateResource wrapper unfortunately does not allow
26# one to remove resources due to overzealous parameter verification.
27# For that case we're forced to go straight to the native API implementation.
28UpdateResource = ctypes.windll.kernel32.UpdateResourceW
29UpdateResource.argtypes = [
30    ctypes.wintypes.HANDLE,  # HANDLE hUpdate
31    ctypes.c_wchar_p,  # LPCTSTR lpType
32    ctypes.c_wchar_p,  # LPCTSTR lpName
33    ctypes.c_short,  # WORD wLanguage
34    ctypes.c_void_p,  # LPVOID lpData
35    ctypes.c_ulong,  # DWORD cbData
36    ]
37UpdateResource.restype = ctypes.c_short
38
39
40def _ResIdToString(res_id):
41  # Convert integral res types/ids to a string.
42  if isinstance(res_id, int):
43    return "#%d" % res_id
44
45  return res_id
46
47
48class _ResourceEditor(object):
49  """A utility class to make it easy to extract and manipulate resources in a
50  Windows binary."""
51
52  def __init__(self, input_file, output_file):
53    """Create a new editor.
54
55    Args:
56        input_file: path to the input file.
57        output_file: (optional) path to the output file.
58    """
59    self._input_file = input_file
60    self._output_file = output_file
61    self._modified = False
62    self._module = None
63    self._temp_dir = None
64    self._temp_file = None
65    self._update_handle = None
66
67  def __del__(self):
68    if self._module:
69      win32api.FreeLibrary(self._module)
70      self._module = None
71
72    if self._update_handle:
73      _LOGGER.info('Canceling edits to "%s".', self.input_file)
74      win32api.EndUpdateResource(self._update_handle, False)
75      self._update_handle = None
76
77    if self._temp_dir:
78      _LOGGER.info('Removing temporary directory "%s".', self._temp_dir)
79      shutil.rmtree(self._temp_dir)
80      self._temp_dir = None
81
82  def _GetModule(self):
83    if not self._module:
84      # Specify a full path to LoadLibraryEx to prevent
85      # it from searching the path.
86      input_file = os.path.abspath(self.input_file)
87      _LOGGER.info('Loading input_file from "%s"', input_file)
88      self._module = win32api.LoadLibraryEx(
89          input_file, None, win32con.LOAD_LIBRARY_AS_DATAFILE)
90    return self._module
91
92  def _GetTempDir(self):
93    if not self._temp_dir:
94      self._temp_dir = tempfile.mkdtemp()
95      _LOGGER.info('Created temporary directory "%s".', self._temp_dir)
96
97    return self._temp_dir
98
99  def _GetUpdateHandle(self):
100    if not self._update_handle:
101      # Make a copy of the input file in the temp dir.
102      self._temp_file = os.path.join(self.temp_dir,
103                                     os.path.basename(self._input_file))
104      shutil.copyfile(self._input_file, self._temp_file)
105      # Open a resource update handle on the copy.
106      _LOGGER.info('Opening temp file "%s".', self._temp_file)
107      self._update_handle = win32api.BeginUpdateResource(self._temp_file, False)
108
109    return self._update_handle
110
111  modified = property(lambda self: self._modified)
112  input_file = property(lambda self: self._input_file)
113  module = property(_GetModule)
114  temp_dir = property(_GetTempDir)
115  update_handle = property(_GetUpdateHandle)
116
117  def ExtractAllToDir(self, extract_to):
118    """Extracts all resources from our input file to a directory hierarchy
119    in the directory named extract_to.
120
121    The generated directory hierarchy is three-level, and looks like:
122      resource-type/
123        resource-name/
124          lang-id.
125
126    Args:
127      extract_to: path to the folder to output to. This folder will be erased
128          and recreated if it already exists.
129    """
130    _LOGGER.info('Extracting all resources from "%s" to directory "%s".',
131        self.input_file, extract_to)
132
133    if os.path.exists(extract_to):
134      _LOGGER.info('Destination directory "%s" exists, deleting', extract_to)
135      shutil.rmtree(extract_to)
136
137    # Make sure the destination dir exists.
138    os.makedirs(extract_to)
139
140    # Now enumerate the resource types.
141    for res_type in win32api.EnumResourceTypes(self.module):
142      res_type_str = _ResIdToString(res_type)
143
144      # And the resource names.
145      for res_name in win32api.EnumResourceNames(self.module, res_type):
146        res_name_str = _ResIdToString(res_name)
147
148        # Then the languages.
149        for res_lang in win32api.EnumResourceLanguages(self.module,
150            res_type, res_name):
151          res_lang_str = _ResIdToString(res_lang)
152
153          dest_dir = os.path.join(extract_to, res_type_str, res_lang_str)
154          dest_file = os.path.join(dest_dir, res_name_str)
155          _LOGGER.info('Extracting resource "%s", lang "%d" name "%s" '
156                       'to file "%s".',
157                       res_type_str, res_lang, res_name_str, dest_file)
158
159          # Extract each resource to a file in the output dir.
160          os.makedirs(dest_dir)
161          self.ExtractResource(res_type, res_lang, res_name, dest_file)
162
163  def ExtractResource(self, res_type, res_lang, res_name, dest_file):
164    """Extracts a given resource, specified by type, language id and name,
165    to a given file.
166
167    Args:
168      res_type: the type of the resource, e.g. "B7".
169      res_lang: the language id of the resource e.g. 1033.
170      res_name: the name of the resource, e.g. "SETUP.EXE".
171      dest_file: path to the file where the resource data will be written.
172    """
173    _LOGGER.info('Extracting resource "%s", lang "%d" name "%s" '
174                 'to file "%s".', res_type, res_lang, res_name, dest_file)
175
176    data = win32api.LoadResource(self.module, res_type, res_name, res_lang)
177    with open(dest_file, 'wb') as f:
178      f.write(data)
179
180  def RemoveResource(self, res_type, res_lang, res_name):
181    """Removes a given resource, specified by type, language id and name.
182
183    Args:
184      res_type: the type of the resource, e.g. "B7".
185      res_lang: the language id of the resource, e.g. 1033.
186      res_name: the name of the resource, e.g. "SETUP.EXE".
187    """
188    _LOGGER.info('Removing resource "%s:%s".', res_type, res_name)
189    # We have to go native to perform a removal.
190    ret = UpdateResource(self.update_handle,
191                         res_type,
192                         res_name,
193                         res_lang,
194                         None,
195                         0)
196    # Raise an error on failure.
197    if ret == 0:
198      error = win32api.GetLastError()
199      print "error", error
200      raise RuntimeError(error)
201    self._modified = True
202
203  def UpdateResource(self, res_type, res_lang, res_name, file_path):
204    """Inserts or updates a given resource with the contents of a file.
205
206    Args:
207      res_type: the type of the resource, e.g. "B7".
208      res_lang: the language id of the resource, e.g. 1033.
209      res_name: the name of the resource, e.g. "SETUP.EXE".
210      file_path: path to the file containing the new resource data.
211    """
212    _LOGGER.info('Writing resource "%s:%s" from file.',
213        res_type, res_name, file_path)
214
215    with open(file_path, 'rb') as f:
216      win32api.UpdateResource(self.update_handle,
217                              res_type,
218                              res_name,
219                              f.read(),
220                              res_lang);
221
222    self._modified = True
223
224  def Commit(self):
225    """Commit any successful resource edits this editor has performed.
226
227    This has the effect of writing the output file.
228    """
229    if self._update_handle:
230      update_handle = self._update_handle
231      self._update_handle = None
232      win32api.EndUpdateResource(update_handle, False)
233
234    _LOGGER.info('Writing edited file to "%s".', self._output_file)
235    shutil.copyfile(self._temp_file, self._output_file)
236
237
238_USAGE = """\
239usage: %prog [options] input_file
240
241A utility script to extract and edit the resources in a Windows executable.
242
243EXAMPLE USAGE:
244# Extract from mini_installer.exe, the resource type "B7", langid 1033 and
245# name "CHROME.PACKED.7Z" to a file named chrome.7z.
246# Note that 1033 corresponds to English (United States).
247%prog mini_installer.exe --extract B7 1033 CHROME.PACKED.7Z chrome.7z
248
249# Update mini_installer.exe by removing the resouce type "BL", langid 1033 and
250# name "SETUP.EXE". Add the resource type "B7", langid 1033 and name
251# "SETUP.EXE.packed.7z" from the file setup.packed.7z.
252# Write the edited file to mini_installer_packed.exe.
253%prog mini_installer.exe \\
254    --remove BL 1033 SETUP.EXE \\
255    --update B7 1033 SETUP.EXE.packed.7z setup.packed.7z \\
256    --output-file mini_installer_packed.exe
257"""
258
259def _ParseArgs():
260  parser = optparse.OptionParser(_USAGE)
261  parser.add_option('', '--verbose', action='store_true',
262      help='Enable verbose logging.')
263  parser.add_option('', '--extract_all',
264      help='Path to a folder which will be created, in which all resources '
265           'from the input_file will be stored, each in a file named '
266           '"res_type/lang_id/res_name".')
267  parser.add_option('', '--extract', action='append', default=[], nargs=4,
268      help='Extract the resource with the given type, language id and name '
269           'to the given file.',
270      metavar='type langid name file_path')
271  parser.add_option('', '--remove', action='append', default=[], nargs=3,
272      help='Remove the resource with the given type, langid and name.',
273      metavar='type langid name')
274  parser.add_option('', '--update', action='append', default=[], nargs=4,
275      help='Insert or update the resource with the given type, langid and '
276           'name with the contents of the file given.',
277      metavar='type langid name file_path')
278  parser.add_option('', '--output_file',
279    help='On success, OUTPUT_FILE will be written with a copy of the '
280         'input file with the edits specified by any remove or update '
281         'options.')
282
283  options, args = parser.parse_args()
284
285  if len(args) != 1:
286    parser.error('You have to specify an input file to work on.')
287
288  modify = options.remove or options.update
289  if modify and not options.output_file:
290    parser.error('You have to specify an output file with edit options.')
291
292  return options, args
293
294
295def main(options, args):
296  """Main program for the script."""
297  if options.verbose:
298    logging.basicConfig(level=logging.INFO)
299
300  # Create the editor for our input file.
301  editor = _ResourceEditor(args[0], options.output_file)
302
303  if options.extract_all:
304    editor.ExtractAllToDir(options.extract_all)
305
306  for res_type, res_lang, res_name, dest_file in options.extract:
307    editor.ExtractResource(res_type, int(res_lang), res_name, dest_file)
308
309  for res_type, res_lang, res_name in options.remove:
310    editor.RemoveResource(res_type, int(res_lang), res_name)
311
312  for res_type, res_lang, res_name, src_file in options.update:
313    editor.UpdateResource(res_type, int(res_lang), res_name, src_file)
314
315  if editor.modified:
316    editor.Commit()
317
318
319if __name__ == '__main__':
320  sys.exit(main(*_ParseArgs()))
321