1#!/usr/bin/env python
2#
3# Copyright (C) 2013 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17
18"""Command-line tool for checking and applying Chrome OS update payloads."""
19
20from __future__ import absolute_import
21from __future__ import print_function
22
23# pylint: disable=import-error
24import argparse
25import filecmp
26import os
27import sys
28import tempfile
29
30# pylint: disable=redefined-builtin
31from six.moves import zip
32from update_payload import error
33
34
35lib_dir = os.path.join(os.path.dirname(__file__), 'lib')
36if os.path.exists(lib_dir) and os.path.isdir(lib_dir):
37  sys.path.insert(1, lib_dir)
38import update_payload  # pylint: disable=wrong-import-position
39
40
41_TYPE_FULL = 'full'
42_TYPE_DELTA = 'delta'
43
44def CheckApplyPayload(args):
45  """Whether to check the result after applying the payload.
46
47  Args:
48    args: Parsed command arguments (the return value of
49          ArgumentParser.parse_args).
50
51  Returns:
52    Boolean value whether to check.
53  """
54  return args.dst_part_paths is not None
55
56def ApplyPayload(args):
57  """Whether to apply the payload.
58
59  Args:
60    args: Parsed command arguments (the return value of
61          ArgumentParser.parse_args).
62
63  Returns:
64    Boolean value whether to apply the payload.
65  """
66  return CheckApplyPayload(args) or args.out_dst_part_paths is not None
67
68def ParseArguments(argv):
69  """Parse and validate command-line arguments.
70
71  Args:
72    argv: command-line arguments to parse (excluding the program name)
73
74  Returns:
75    Returns the arguments returned by the argument parser.
76  """
77  parser = argparse.ArgumentParser(
78      description=('Applies a Chrome OS update PAYLOAD to src_part_paths'
79                   'emitting dst_part_paths, respectively. '
80                   'src_part_paths are only needed for delta payloads. '
81                   'When no partitions are provided, verifies the payload '
82                   'integrity.'),
83      epilog=('Note: a payload may verify correctly but fail to apply, and '
84              'vice versa; this is by design and can be thought of as static '
85              'vs dynamic correctness. A payload that both verifies and '
86              'applies correctly should be safe for use by the Chrome OS '
87              'Update Engine. Use --check to verify a payload prior to '
88              'applying it.'),
89      formatter_class=argparse.RawDescriptionHelpFormatter
90  )
91
92  check_args = parser.add_argument_group('Checking payload integrity')
93  check_args.add_argument('-c', '--check', action='store_true', default=False,
94                          help=('force payload integrity check (e.g. before '
95                                'applying)'))
96  check_args.add_argument('-r', '--report', metavar='FILE',
97                          help="dump payload report (`-' for stdout)")
98  check_args.add_argument('-t', '--type', dest='assert_type',
99                          help='assert the payload type',
100                          choices=[_TYPE_FULL, _TYPE_DELTA])
101  check_args.add_argument('-z', '--block-size', metavar='NUM', default=0,
102                          type=int,
103                          help='assert a non-default (4096) payload block size')
104  check_args.add_argument('-u', '--allow-unhashed', action='store_true',
105                          default=False, help='allow unhashed operations')
106  check_args.add_argument('-d', '--disabled_tests', default=(), metavar='',
107                          help=('space separated list of tests to disable. '
108                                'allowed options include: ' +
109                                ', '.join(update_payload.CHECKS_TO_DISABLE)),
110                          choices=update_payload.CHECKS_TO_DISABLE)
111  check_args.add_argument('-k', '--key', metavar='FILE',
112                          help=('override standard key used for signature '
113                                'validation'))
114  check_args.add_argument('-m', '--meta-sig', metavar='FILE',
115                          help='verify metadata against its signature')
116  check_args.add_argument('-s', '--metadata-size', metavar='NUM', default=0,
117                          help='the metadata size to verify with the one in'
118                          ' payload')
119  check_args.add_argument('--part_sizes', metavar='NUM', nargs='+', type=int,
120                          help='override partition size auto-inference')
121
122  apply_args = parser.add_argument_group('Applying payload')
123  # TODO(ahassani): Extent extract-bsdiff to puffdiff too.
124  apply_args.add_argument('-x', '--extract-bsdiff', action='store_true',
125                          default=False,
126                          help=('use temp input/output files with BSDIFF '
127                                'operations (not in-place)'))
128  apply_args.add_argument('--bspatch-path', metavar='FILE',
129                          help='use the specified bspatch binary')
130  apply_args.add_argument('--puffpatch-path', metavar='FILE',
131                          help='use the specified puffpatch binary')
132
133  apply_args.add_argument('--src_part_paths', metavar='FILE', nargs='+',
134                          help='source partitition files')
135  apply_args.add_argument('--dst_part_paths', metavar='FILE', nargs='+',
136                          help='destination partition files')
137  apply_args.add_argument('--out_dst_part_paths', metavar='FILE', nargs='+',
138                          help='created destination partition files')
139
140  parser.add_argument('payload', metavar='PAYLOAD', help='the payload file')
141  parser.add_argument('--part_names', metavar='NAME', nargs='+',
142                      help='names of partitions')
143
144  # Parse command-line arguments.
145  args = parser.parse_args(argv)
146
147  # There are several options that imply --check.
148  args.check = (args.check or args.report or args.assert_type or
149                args.block_size or args.allow_unhashed or
150                args.disabled_tests or args.meta_sig or args.key or
151                args.part_sizes is not None or args.metadata_size)
152
153  # Makes sure the following arguments have the same length as |part_names| if
154  # set.
155  for arg in ['part_sizes', 'src_part_paths', 'dst_part_paths',
156              'out_dst_part_paths']:
157    if getattr(args, arg) is None:
158      # Parameter is not set.
159      continue
160    if len(args.part_names) != len(getattr(args, arg, [])):
161      parser.error('partitions in --%s do not match --part_names' % arg)
162
163  def _IsSrcPartPathsProvided(args):
164    return args.src_part_paths is not None
165
166  # Makes sure parameters are coherent with payload type.
167  if ApplyPayload(args):
168    if _IsSrcPartPathsProvided(args):
169      if args.assert_type == _TYPE_FULL:
170        parser.error('%s payload does not accept source partition arguments'
171                     % _TYPE_FULL)
172      else:
173        args.assert_type = _TYPE_DELTA
174    else:
175      if args.assert_type == _TYPE_DELTA:
176        parser.error('%s payload requires source partitions arguments'
177                     % _TYPE_DELTA)
178      else:
179        args.assert_type = _TYPE_FULL
180  else:
181    # Not applying payload.
182    if args.extract_bsdiff:
183      parser.error('--extract-bsdiff can only be used when applying payloads')
184    if args.bspatch_path:
185      parser.error('--bspatch-path can only be used when applying payloads')
186    if args.puffpatch_path:
187      parser.error('--puffpatch-path can only be used when applying payloads')
188
189  # By default, look for a metadata-signature file with a name based on the name
190  # of the payload we are checking. We only do it if check was triggered.
191  if args.check and not args.meta_sig:
192    default_meta_sig = args.payload + '.metadata-signature'
193    if os.path.isfile(default_meta_sig):
194      args.meta_sig = default_meta_sig
195      print('Using default metadata signature', args.meta_sig, file=sys.stderr)
196
197  return args
198
199
200def main(argv):
201  # Parse and validate arguments.
202  args = ParseArguments(argv[1:])
203
204  with open(args.payload, 'rb') as payload_file:
205    payload = update_payload.Payload(payload_file)
206    try:
207      # Initialize payload.
208      payload.Init()
209
210      # Perform payload integrity checks.
211      if args.check:
212        report_file = None
213        do_close_report_file = False
214        metadata_sig_file = None
215        try:
216          if args.report:
217            if args.report == '-':
218              report_file = sys.stdout
219            else:
220              report_file = open(args.report, 'w')
221              do_close_report_file = True
222
223          part_sizes = (args.part_sizes and
224                        dict(zip(args.part_names, args.part_sizes)))
225          metadata_sig_file = args.meta_sig and open(args.meta_sig, 'rb')
226          payload.Check(
227              pubkey_file_name=args.key,
228              metadata_sig_file=metadata_sig_file,
229              metadata_size=int(args.metadata_size),
230              report_out_file=report_file,
231              assert_type=args.assert_type,
232              block_size=int(args.block_size),
233              part_sizes=part_sizes,
234              allow_unhashed=args.allow_unhashed,
235              disabled_tests=args.disabled_tests)
236        finally:
237          if metadata_sig_file:
238            metadata_sig_file.close()
239          if do_close_report_file:
240            report_file.close()
241
242      # Apply payload.
243      if ApplyPayload(args):
244        dargs = {'bsdiff_in_place': not args.extract_bsdiff}
245        if args.bspatch_path:
246          dargs['bspatch_path'] = args.bspatch_path
247        if args.puffpatch_path:
248          dargs['puffpatch_path'] = args.puffpatch_path
249        if args.assert_type == _TYPE_DELTA:
250          dargs['old_parts'] = dict(zip(args.part_names, args.src_part_paths))
251
252        out_dst_parts = {}
253        file_handles = []
254        if args.out_dst_part_paths is not None:
255          for name, path in zip(args.part_names, args.out_dst_part_paths):
256            handle = open(path, 'wb+')
257            file_handles.append(handle)
258            out_dst_parts[name] = handle.name
259        else:
260          for name in args.part_names:
261            handle = tempfile.NamedTemporaryFile()
262            file_handles.append(handle)
263            out_dst_parts[name] = handle.name
264
265        payload.Apply(out_dst_parts, **dargs)
266
267        # If destination kernel and rootfs partitions are not given, then this
268        # just becomes an apply operation with no check.
269        if CheckApplyPayload(args):
270          # Prior to comparing, add the unused space past the filesystem
271          # boundary in the new target partitions to become the same size as
272          # the given partitions. This will truncate to larger size.
273          for part_name, out_dst_part, dst_part in zip(args.part_names,
274                                                       file_handles,
275                                                       args.dst_part_paths):
276            out_dst_part.truncate(os.path.getsize(dst_part))
277
278            # Compare resulting partitions with the ones from the target image.
279            if not filecmp.cmp(out_dst_part.name, dst_part):
280              raise error.PayloadError(
281                  'Resulting %s partition corrupted.' % part_name)
282
283        # Close the output files. If args.out_dst_* was not given, then these
284        # files are created as temp files and will be deleted upon close().
285        for handle in file_handles:
286          handle.close()
287    except error.PayloadError as e:
288      sys.stderr.write('Error: %s\n' % e)
289      return 1
290
291  return 0
292
293
294if __name__ == '__main__':
295  sys.exit(main(sys.argv))
296