update_payload: use argparse

optparse is deprecated (or going to be). Change paycheck.py and
blockdiff.py to use argparse instead. Both of these files are being used
manually and it would be a good time to fix these before major changes
in update_payload.

paycheck.sh -h:

usage: paycheck.py [-h] [-c] [-D] [-r FILE] [-t {full,delta}] [-z NUM] [-u]
                   [-d] [-k FILE] [-m FILE] [-p NUM] [-P NUM] [-x]
                   [--bspatch-path FILE] [--puffpatch-path FILE]
                   [--dst_kern FILE] [--dst_root FILE] [--src_kern FILE]
                   [--src_root FILE] [-b BLOCK] [-B BLOCK] [-s NUM]
                   PAYLOAD

Applies a Chrome OS update PAYLOAD to src_kern and src_root emitting dst_kern and dst_root, respectively. src_kern and src_root are only needed for delta payloads. When no partitions are provided, verifies the payload integrity.

positional arguments:
  PAYLOAD               the payload file

optional arguments:
  -h, --help            show this help message and exit

Checking payload integrity:
  -c, --check           force payload integrity check (e.g. before applying)
  -D, --describe        Print a friendly description of the payload.
  -r FILE, --report FILE
                        dump payload report (`-' for stdout)
  -t {full,delta}, --type {full,delta}
                        assert the payload type
  -z NUM, --block-size NUM
                        assert a non-default (4096) payload block size
  -u, --allow-unhashed  allow unhashed operations
  -d , --disabled_tests
                        space separated list of tests to disable. allowed
                        options include: dst-pseudo-extents, move-same-src-
                        dst-block, payload-sig
  -k FILE, --key FILE   override standard key used for signature validation
  -m FILE, --meta-sig FILE
                        verify metadata against its signature
  -p NUM, --root-part-size NUM
                        override rootfs partition size auto-inference
  -P NUM, --kern-part-size NUM
                        override kernel partition size auto-inference

Applying payload:
  -x, --extract-bsdiff  use temp input/output files with BSDIFF operations
                        (not in-place)
  --bspatch-path FILE   use the specified bspatch binary
  --puffpatch-path FILE
                        use the specified puffpatch binary
  --dst_kern FILE       destination kernel partition file
  --dst_root FILE       destination root partition file
  --src_kern FILE       source kernel partition file
  --src_root FILE       source root partition file

Block tracing:
  -b BLOCK, --root-block BLOCK
                        trace the origin for a rootfs block
  -B BLOCK, --kern-block BLOCK
                        trace the origin for a kernel block
  -s NUM, --skip NUM    skip first NUM occurrences of traced block

Note: a payload may verify correctly but fail to apply, and vice versa; this is by design and can be thought of as static vs dynamic correctness. A payload that both verifies and applies correctly should be safe for use by the Chrome OS Update Engine. Use --check to verify a payload prior to applying it.

BUG=chromium:796338
TEST=unitests
TEST=test_paycheck.sh
TEST=blockdiff.py

Change-Id: I794b5f61e6ba6f92939947c97c432f9fea0b6b3c
Reviewed-on: https://chromium-review.googlesource.com/834876
Commit-Ready: Amin Hassani <ahassani@chromium.org>
Tested-by: Amin Hassani <ahassani@chromium.org>
Reviewed-by: Ben Chan <benchan@chromium.org>
Reviewed-by: Sen Jiang <senj@chromium.org>

3 files changed, 161 insertions, 159 deletions

diff --git a/scripts/blockdiff.py b/scripts/blockdiff.py
index 1dc60a68..1f580c94 100755
--- a/scripts/blockdiff.py
+++ b/scripts/blockdiff.py
@@ -8,7 +8,8 @@
 
 from __future__ import print_function
 
-import optparse
+# pylint: disable=import-error
+import argparse
 import sys
 
 
@@ -71,28 +72,25 @@ def BlockDiff(block_size, file1, file2, name1, name2, max_length=-1):
 
 def main(argv):
   # Parse command-line arguments.
-  parser = optparse.OptionParser(
-      usage='Usage: %prog FILE1 FILE2',
-      description='Compare FILE1 and FILE2 by blocks.')
+  parser = argparse.ArgumentParser(
+      description='Compare FILE1 and FILE2 by blocks.',
+      formatter_class=argparse.ArgumentDefaultsHelpFormatter)
 
-  parser.add_option('-b', '--block-size', metavar='NUM', type=int, default=4096,
-                    help='the block size to use (default: %default)')
-  parser.add_option('-m', '--max-length', metavar='NUM', type=int, default=-1,
-                    help='maximum number of bytes to compared')
+  parser.add_argument('-b', '--block-size', metavar='NUM', type=int,
+                      default=4096, help='the block size to use')
+  parser.add_argument('-m', '--max-length', metavar='NUM', type=int, default=-1,
+                      help='maximum number of bytes to compare')
+  parser.add_argument('file1', metavar='FILE1')
+  parser.add_argument('file2', metavar='FILE2')
 
-  opts, args = parser.parse_args(argv[1:])
-
-  try:
-    name1, name2 = args
-  except ValueError:
-    parser.error('unexpected number of arguments')
+  args = parser.parse_args(argv[1:])
 
   # Perform the block diff.
   try:
-    with open(name1) as file1:
-      with open(name2) as file2:
-        diff_list = BlockDiff(opts.block_size, file1, file2, name1, name2,
-                              opts.max_length)
+    with open(args.file1) as file1:
+      with open(args.file2) as file2:
+        diff_list = BlockDiff(args.block_size, file1, file2,
+                              args.file1, args.file2, args.max_length)
   except BlockDiffError as e:
     print('Error: ' % e, file=sys.stderr)
     return 2
diff --git a/scripts/paycheck.py b/scripts/paycheck.py
index 8df1bf07..7f0b9a32 100755
--- a/scripts/paycheck.py
+++ b/scripts/paycheck.py
@@ -8,7 +8,8 @@
 
 from __future__ import print_function
 
-import optparse
+# pylint: disable=import-error
+import argparse
 import os
 import sys
 
@@ -29,17 +30,12 @@ def ParseArguments(argv):
     argv: command-line arguments to parse (excluding the program name)
 
   Returns:
-    A tuple (opts, payload, extra_args), where `opts' are the options
-    returned by the parser, `payload' is the name of the payload file
-    (mandatory argument) and `extra_args' are any additional command-line
-    arguments.
+    Returns the arguments returned by the argument parser.
   """
-  parser = optparse.OptionParser(
-      usage=('Usage: %prog [OPTION...] PAYLOAD [DST_KERN DST_ROOT '
-             '[SRC_KERN SRC_ROOT]]'),
-      description=('Applies a Chrome OS update PAYLOAD to SRC_KERN and '
-                   'SRC_ROOT emitting DST_KERN and DST_ROOT, respectively. '
-                   'SRC_KERN and SRC_ROOT are only needed for delta payloads. '
+  parser = argparse.ArgumentParser(
+      description=('Applies a Chrome OS update PAYLOAD to src_kern and '
+                   'src_root emitting dst_kern and dst_root, respectively. '
+                   'src_kern and src_root are only needed for delta payloads. '
                    'When no partitions are provided, verifies the payload '
                    'integrity.'),
       epilog=('Note: a payload may verify correctly but fail to apply, and '
@@ -47,160 +43,166 @@ def ParseArguments(argv):
               'vs dynamic correctness. A payload that both verifies and '
               'applies correctly should be safe for use by the Chrome OS '
               'Update Engine. Use --check to verify a payload prior to '
-              'applying it.'))
-
-  check_opts = optparse.OptionGroup(parser, 'Checking payload integrity')
-  check_opts.add_option('-c', '--check', action='store_true', default=False,
-                        help=('force payload integrity check (e.g. before '
-                              'applying)'))
-  check_opts.add_option('-D', '--describe', action='store_true', default=False,
-                        help='Print a friendly description of the payload.')
-  check_opts.add_option('-r', '--report', metavar='FILE',
-                        help="dump payload report (`-' for stdout)")
-  check_opts.add_option('-t', '--type', metavar='TYPE', dest='assert_type',
-                        help=("assert that payload is either `%s' or `%s'" %
-                              (_TYPE_FULL, _TYPE_DELTA)))
-  check_opts.add_option('-z', '--block-size', metavar='NUM', default=0,
-                        type='int',
-                        help='assert a non-default (4096) payload block size')
-  check_opts.add_option('-u', '--allow-unhashed', action='store_true',
-                        default=False, help='allow unhashed operations')
-  check_opts.add_option('-d', '--disabled_tests', metavar='TESTLIST',
-                        default=(),
-                        help=('comma-separated list of tests to disable; '
-                              'available values: ' +
-                              ', '.join(update_payload.CHECKS_TO_DISABLE)))
-  check_opts.add_option('-k', '--key', metavar='FILE',
-                        help=('Override standard key used for signature '
-                              'validation'))
-  check_opts.add_option('-m', '--meta-sig', metavar='FILE',
-                        help='verify metadata against its signature')
-  check_opts.add_option('-p', '--root-part-size', metavar='NUM',
-                        default=0, type='int',
-                        help=('override rootfs partition size auto-inference'))
-  check_opts.add_option('-P', '--kern-part-size', metavar='NUM',
-                        default=0, type='int',
-                        help=('override kernel partition size auto-inference'))
-  parser.add_option_group(check_opts)
-
-  trace_opts = optparse.OptionGroup(parser, 'Applying payload')
-  trace_opts.add_option('-x', '--extract-bsdiff', action='store_true',
-                        default=False,
-                        help=('use temp input/output files with BSDIFF '
-                              'operations (not in-place)'))
-  trace_opts.add_option('--bspatch-path', metavar='FILE',
-                        help=('use the specified bspatch binary'))
-  trace_opts.add_option('--puffpatch-path', metavar='FILE',
-                        help=('use the specified puffpatch binary'))
-  parser.add_option_group(trace_opts)
-
-  trace_opts = optparse.OptionGroup(parser, 'Block tracing')
-  trace_opts.add_option('-b', '--root-block', metavar='BLOCK', type='int',
-                        help='trace the origin for a rootfs block')
-  trace_opts.add_option('-B', '--kern-block', metavar='BLOCK', type='int',
-                        help='trace the origin for a kernel block')
-  trace_opts.add_option('-s', '--skip', metavar='NUM', default='0', type='int',
-                        help='skip first NUM occurrences of traced block')
-  parser.add_option_group(trace_opts)
+              'applying it.'),
+      formatter_class=argparse.RawDescriptionHelpFormatter
+  )
+
+  check_args = parser.add_argument_group('Checking payload integrity')
+  check_args.add_argument('-c', '--check', action='store_true', default=False,
+                          help=('force payload integrity check (e.g. before '
+                                'applying)'))
+  check_args.add_argument('-D', '--describe', action='store_true',
+                          default=False,
+                          help='Print a friendly description of the payload.')
+  check_args.add_argument('-r', '--report', metavar='FILE',
+                          help="dump payload report (`-' for stdout)")
+  check_args.add_argument('-t', '--type', dest='assert_type',
+                          help='assert the payload type',
+                          choices=[_TYPE_FULL, _TYPE_DELTA])
+  check_args.add_argument('-z', '--block-size', metavar='NUM', default=0,
+                          type=int,
+                          help='assert a non-default (4096) payload block size')
+  check_args.add_argument('-u', '--allow-unhashed', action='store_true',
+                          default=False, help='allow unhashed operations')
+  check_args.add_argument('-d', '--disabled_tests', default=(), metavar='',
+                          help=('space separated list of tests to disable. '
+                                'allowed options include: ' +
+                                ', '.join(update_payload.CHECKS_TO_DISABLE)),
+                          choices=update_payload.CHECKS_TO_DISABLE)
+  check_args.add_argument('-k', '--key', metavar='FILE',
+                          help=('override standard key used for signature '
+                                'validation'))
+  check_args.add_argument('-m', '--meta-sig', metavar='FILE',
+                          help='verify metadata against its signature')
+  check_args.add_argument('-p', '--root-part-size', metavar='NUM',
+                          default=0, type=int,
+                          help='override rootfs partition size auto-inference')
+  check_args.add_argument('-P', '--kern-part-size', metavar='NUM',
+                          default=0, type=int,
+                          help='override kernel partition size auto-inference')
+
+  apply_args = parser.add_argument_group('Applying payload')
+  # TODO(ahassani): Extent extract-bsdiff to puffdiff too.
+  apply_args.add_argument('-x', '--extract-bsdiff', action='store_true',
+                          default=False,
+                          help=('use temp input/output files with BSDIFF '
+                                'operations (not in-place)'))
+  apply_args.add_argument('--bspatch-path', metavar='FILE',
+                          help='use the specified bspatch binary')
+  apply_args.add_argument('--puffpatch-path', metavar='FILE',
+                          help='use the specified puffpatch binary')
+  apply_args.add_argument('--dst_kern', metavar='FILE',
+                          help='destination kernel partition file')
+  apply_args.add_argument('--dst_root', metavar='FILE',
+                          help='destination root partition file')
+  apply_args.add_argument('--src_kern', metavar='FILE',
+                          help='source kernel partition file')
+  apply_args.add_argument('--src_root', metavar='FILE',
+                          help='source root partition file')
+
+  trace_args = parser.add_argument_group('Block tracing')
+  trace_args.add_argument('-b', '--root-block', metavar='BLOCK', type=int,
+                          help='trace the origin for a rootfs block')
+  trace_args.add_argument('-B', '--kern-block', metavar='BLOCK', type=int,
+                          help='trace the origin for a kernel block')
+  trace_args.add_argument('-s', '--skip', metavar='NUM', default='0', type=int,
+                          help='skip first NUM occurrences of traced block')
+
+  parser.add_argument('payload', metavar='PAYLOAD', help='the payload file')
 
   # Parse command-line arguments.
-  opts, args = parser.parse_args(argv)
-
-  # Validate a value given to --type, if any.
-  if opts.assert_type not in (None, _TYPE_FULL, _TYPE_DELTA):
-    parser.error('invalid argument to --type: %s' % opts.assert_type)
-
-  # Convert and validate --disabled_tests value list, if provided.
-  if opts.disabled_tests:
-    opts.disabled_tests = opts.disabled_tests.split(',')
-    for test in opts.disabled_tests:
-      if test not in update_payload.CHECKS_TO_DISABLE:
-        parser.error('invalid argument to --disabled_tests: %s' % test)
+  args = parser.parse_args(argv)
 
   # Ensure consistent use of block tracing options.
-  do_block_trace = not (opts.root_block is None and opts.kern_block is None)
-  if opts.skip and not do_block_trace:
+  do_block_trace = not (args.root_block is None and args.kern_block is None)
+  if args.skip and not do_block_trace:
     parser.error('--skip must be used with either --root-block or --kern-block')
 
   # There are several options that imply --check.
-  opts.check = (opts.check or opts.report or opts.assert_type or
-                opts.block_size or opts.allow_unhashed or
-                opts.disabled_tests or opts.meta_sig or opts.key or
-                opts.root_part_size or opts.kern_part_size)
-
-  # Check number of arguments, enforce payload type accordingly.
-  if len(args) == 3:
-    if opts.assert_type == _TYPE_DELTA:
-      parser.error('%s payload requires source partition arguments' %
-                   _TYPE_DELTA)
-    opts.assert_type = _TYPE_FULL
-  elif len(args) == 5:
-    if opts.assert_type == _TYPE_FULL:
-      parser.error('%s payload does not accept source partition arguments' %
-                   _TYPE_FULL)
-    opts.assert_type = _TYPE_DELTA
-  elif len(args) == 1:
+  args.check = (args.check or args.report or args.assert_type or
+                args.block_size or args.allow_unhashed or
+                args.disabled_tests or args.meta_sig or args.key or
+                args.root_part_size or args.kern_part_size)
+
+  # Check the arguments, enforce payload type accordingly.
+  if (args.src_kern is None) != (args.src_root is None):
+    parser.error('--src_kern and --src_root should be given together')
+  if (args.dst_kern is None) != (args.dst_root is None):
+    parser.error('--dst_kern and --dst_root should be given together')
+
+  if args.dst_kern and args.dst_root:
+    if args.src_kern and args.src_root:
+      if args.assert_type == _TYPE_FULL:
+        parser.error('%s payload does not accept source partition arguments'
+                     % _TYPE_FULL)
+      else:
+        args.assert_type = _TYPE_DELTA
+    else:
+      if args.assert_type == _TYPE_DELTA:
+        parser.error('%s payload requires source partitions arguments'
+                     % _TYPE_DELTA)
+      else:
+        args.assert_type = _TYPE_FULL
+  else:
     # Not applying payload; if block tracing not requested either, do an
     # integrity check.
     if not do_block_trace:
-      opts.check = True
-    if opts.extract_bsdiff:
+      args.check = True
+    if args.extract_bsdiff:
       parser.error('--extract-bsdiff can only be used when applying payloads')
-    if opts.bspatch_path:
+    if args.bspatch_path:
       parser.error('--bspatch-path can only be used when applying payloads')
-    if opts.puffpatch_path:
+    if args.puffpatch_path:
       parser.error('--puffpatch-path can only be used when applying payloads')
-  else:
-    parser.error('unexpected number of arguments')
 
   # By default, look for a metadata-signature file with a name based on the name
   # of the payload we are checking. We only do it if check was triggered.
-  if opts.check and not opts.meta_sig:
-    default_meta_sig = args[0] + '.metadata-signature'
+  if args.check and not args.meta_sig:
+    default_meta_sig = args.payload + '.metadata-signature'
     if os.path.isfile(default_meta_sig):
-      opts.meta_sig = default_meta_sig
-      print('Using default metadata signature', opts.meta_sig, file=sys.stderr)
+      args.meta_sig = default_meta_sig
+      print('Using default metadata signature', args.meta_sig, file=sys.stderr)
 
-  return opts, args[0], args[1:]
+  return args
 
 
 def main(argv):
   # Parse and validate arguments.
-  options, payload_file_name, extra_args = ParseArguments(argv[1:])
+  args = ParseArguments(argv[1:])
 
-  with open(payload_file_name) as payload_file:
+  with open(args.payload) as payload_file:
     payload = update_payload.Payload(payload_file)
     try:
       # Initialize payload.
       payload.Init()
 
-      if options.describe:
+      if args.describe:
         payload.Describe()
 
       # Perform payload integrity checks.
-      if options.check:
+      if args.check:
         report_file = None
         do_close_report_file = False
         metadata_sig_file = None
         try:
-          if options.report:
-            if options.report == '-':
+          if args.report:
+            if args.report == '-':
               report_file = sys.stdout
             else:
-              report_file = open(options.report, 'w')
+              report_file = open(args.report, 'w')
               do_close_report_file = True
 
-          metadata_sig_file = options.meta_sig and open(options.meta_sig)
+          metadata_sig_file = args.meta_sig and open(args.meta_sig)
           payload.Check(
-              pubkey_file_name=options.key,
+              pubkey_file_name=args.key,
               metadata_sig_file=metadata_sig_file,
               report_out_file=report_file,
-              assert_type=options.assert_type,
-              block_size=int(options.block_size),
-              rootfs_part_size=options.root_part_size,
-              kernel_part_size=options.kern_part_size,
-              allow_unhashed=options.allow_unhashed,
-              disabled_tests=options.disabled_tests)
+              assert_type=args.assert_type,
+              block_size=int(args.block_size),
+              rootfs_part_size=args.root_part_size,
+              kernel_part_size=args.kern_part_size,
+              allow_unhashed=args.allow_unhashed,
+              disabled_tests=args.disabled_tests)
         finally:
           if metadata_sig_file:
             metadata_sig_file.close()
@@ -208,23 +210,23 @@ def main(argv):
             report_file.close()
 
       # Trace blocks.
-      if options.root_block is not None:
-        payload.TraceBlock(options.root_block, options.skip, sys.stdout, False)
-      if options.kern_block is not None:
-        payload.TraceBlock(options.kern_block, options.skip, sys.stdout, True)
+      if args.root_block is not None:
+        payload.TraceBlock(args.root_block, args.skip, sys.stdout, False)
+      if args.kern_block is not None:
+        payload.TraceBlock(args.kern_block, args.skip, sys.stdout, True)
 
       # Apply payload.
-      if extra_args:
-        dargs = {'bsdiff_in_place': not options.extract_bsdiff}
-        if options.bspatch_path:
-          dargs['bspatch_path'] = options.bspatch_path
-        if options.puffpatch_path:
-          dargs['puffpatch_path'] = options.puffpatch_path
-        if options.assert_type == _TYPE_DELTA:
-          dargs['old_kernel_part'] = extra_args[2]
-          dargs['old_rootfs_part'] = extra_args[3]
-
-        payload.Apply(extra_args[0], extra_args[1], **dargs)
+      if args.dst_root or args.dst_kern:
+        dargs = {'bsdiff_in_place': not args.extract_bsdiff}
+        if args.bspatch_path:
+          dargs['bspatch_path'] = args.bspatch_path
+        if args.puffpatch_path:
+          dargs['puffpatch_path'] = args.puffpatch_path
+        if args.assert_type == _TYPE_DELTA:
+          dargs['old_kernel_part'] = args.src_kern
+          dargs['old_rootfs_part'] = args.src_root
+
+        payload.Apply(args.dst_kern, args.dst_root, **dargs)
 
     except update_payload.PayloadError, e:
       sys.stderr.write('Error: %s\n' % e)
diff --git a/scripts/test_paycheck.sh b/scripts/test_paycheck.sh
index c395db4d..92f7de24 100755
--- a/scripts/test_paycheck.sh
+++ b/scripts/test_paycheck.sh
@@ -97,7 +97,8 @@ apply_full_payload() {
   dst_kern_part="$2/$3"
   dst_root_part="$2/$4"
 
-  time ${paycheck} ${payload_file} ${dst_kern_part} ${dst_root_part}
+  time ${paycheck} ${payload_file} \
+    --dst_kern ${dst_kern_part} --dst_root ${dst_root_part}
 }
 
 apply_delta_payload() {
@@ -107,8 +108,9 @@ apply_delta_payload() {
   src_kern_part="$2/$5"
   src_root_part="$2/$6"
 
-  time ${paycheck} ${payload_file} ${dst_kern_part} ${dst_root_part} \
-    ${src_kern_part} ${src_root_part}
+  time ${paycheck} ${payload_file} \
+    --dst_kern ${dst_kern_part} --dst_root ${dst_root_part} \
+    --src_kern ${src_kern_part} --dst_kern ${src_root_part}
 }
 
 main() {