src/third_party/chromite/lib/paygen/paygen_payload_lib.py

   1 # Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 """Hold the functions that do the real work generating payloads."""
   6
   7 from __future__ import print_function
   8
   9 import base64
  10 import datetime
  11 import filecmp
  12 import json
  13 import logging
  14 import os
  15 import shutil
  16 import tempfile
  17
  18 import fixup_path
  19 fixup_path.FixupPath()
  20
  21 from chromite.lib import cros_build_lib
  22 from chromite.lib import osutils
  23 from chromite.lib.paygen import dryrun_lib
  24 from chromite.lib.paygen import filelib
  25 from chromite.lib.paygen import gspaths
  26 from chromite.lib.paygen import signer_payloads_client
  27 from chromite.lib.paygen import urilib
  28 from chromite.lib.paygen import utils
  29
  30
  31 DESCRIPTION_FILE_VERSION = 1
  32
  33 class Error(Exception):
  34   """Base class for payload generation errors."""
  35
  36
  37 class UnexpectedSignerResultsError(Error):
  38   """This is raised when signer results don't match our expectations."""
  39
  40
  41 class PayloadVerificationError(Error):
  42   """Raised when the generated payload fails to verify."""
  43
  44
  45 class _PaygenPayload(object):
  46   """Class to manage the process of generating and signing a payload."""
  47
  48   # GeneratorUri uses these to ensure we don't use generators that are too
  49   # old to be supported.
  50   MINIMUM_GENERATOR_VERSION = '6303.0.0'
  51   MINIMUM_GENERATOR_URI = (
  52       'gs://chromeos-releases/canary-channel/x86-mario/%s/au-generator.zip' %
  53       MINIMUM_GENERATOR_VERSION)
  54
  55   # What keys do we sign payloads with, and what size are they?
  56   PAYLOAD_SIGNATURE_KEYSETS = ('update_signer',)
  57   PAYLOAD_SIGNATURE_SIZES_BYTES = (2048 / 8,)  # aka 2048 bits in bytes.
  58
  59   TEST_IMAGE_NAME = 'chromiumos_test_image.bin'
  60   RECOVERY_IMAGE_NAME = 'chromiumos_recovery_image.bin'
  61
  62   # Default names used by cros_generate_update_payload for extracting old/new
  63   # kernel/rootfs partitions.
  64   _DEFAULT_OLD_KERN_PART = 'old_kern.dat'
  65   _DEFAULT_OLD_ROOT_PART = 'old_root.dat'
  66   _DEFAULT_NEW_KERN_PART = 'new_kern.dat'
  67   _DEFAULT_NEW_ROOT_PART = 'new_root.dat'
  68
  69   # TODO(garnold)(chromium:243559) stop using these constants once we start
  70   # embedding partition sizes in payloads.
  71   _DEFAULT_ROOTFS_PART_SIZE = 2 * 1024 * 1024 * 1024
  72   _DEFAULT_KERNEL_PART_SIZE = 16 * 1024 * 1024
  73
  74   def __init__(self, payload, cache, work_dir, sign, verify,
  75                au_generator_uri_override, dry_run=False):
  76     """Init for _PaygenPayload.
  77
  78     Args:
  79       payload: An instance of gspaths.Payload describing the payload to
  80                generate.
  81       cache: An instance of DownloadCache for retrieving files.
  82       work_dir: A working directory for output files. Can NOT be shared.
  83       sign: Boolean saying if the payload should be signed (normally, you do).
  84       verify: whether the payload should be verified after being generated
  85       au_generator_uri_override: URI to override standard au_generator.zip
  86           rules.
  87       dry_run: do not do any actual work
  88     """
  89     self.payload = payload
  90     self.cache = cache
  91     self.work_dir = work_dir
  92     self._verify = verify
  93     self._au_generator_uri_override = au_generator_uri_override
  94     self._drm = dryrun_lib.DryRunMgr(dry_run)
  95
  96     self.generator_dir = os.path.join(work_dir, 'au-generator')
  97     self.src_image_file = os.path.join(work_dir, 'src_image.bin')
  98     self.tgt_image_file = os.path.join(work_dir, 'tgt_image.bin')
  99
 100     self.payload_file = os.path.join(work_dir, 'delta.bin')
 101     self.delta_log_file = os.path.join(work_dir, 'delta.log')
 102     self.description_file = os.path.join(work_dir, 'delta.json')
 103
 104     self.signer = None
 105
 106     # If we are a bootstrap environment, this import will fail, so don't
 107     # perform it until we need it.
 108     from dev.host.lib import update_payload
 109
 110     self._update_payload = update_payload
 111
 112     if sign:
 113       self.signed_payload_file = self.payload_file + '.signed'
 114       self.metadata_signature_file = self._MetadataUri(self.signed_payload_file)
 115
 116       self.signer = signer_payloads_client.SignerPayloadsClientGoogleStorage(
 117           payload.tgt_image.channel,
 118           payload.tgt_image.board,
 119           payload.tgt_image.version)
 120
 121   def _MetadataUri(self, uri):
 122     """Given a payload uri, find the uri for the metadata signature."""
 123     return uri + '.metadata-signature'
 124
 125   def _DeltaLogsUri(self, uri):
 126     """Given a payload uri, find the uri for the delta generator logs."""
 127     return uri + '.log'
 128
 129   def _JsonUri(self, uri):
 130     """Given a payload uri, find the uri for the json payload description."""
 131     return uri + '.json'
 132
 133   def _GeneratorUri(self):
 134     """Find the URI for the au-generator.zip to use to generate this payload.
 135
 136     The intent is to always find a generator compatible with the version
 137     that will process the update generated. Notice that Full updates must
 138     be compatible with all versions, no matter how old.
 139
 140     Returns:
 141       URI of an au-generator.zip in string form.
 142     """
 143     if self._au_generator_uri_override:
 144       return self._au_generator_uri_override
 145
 146     if (self.payload.src_image and
 147         gspaths.VersionGreater(self.payload.src_image.version,
 148                                self.MINIMUM_GENERATOR_VERSION)):
 149       # If we are a delta, and newer than the minimum delta age,
 150       # Use the generator from the src.
 151       return gspaths.ChromeosReleases.GeneratorUri(
 152           self.payload.src_image.channel,
 153           self.payload.src_image.board,
 154           self.payload.src_image.version)
 155     else:
 156       # If we are a full update, or a delta from older than minimum, use
 157       # the minimum generator version.
 158       return self.MINIMUM_GENERATOR_URI
 159
 160   def _PrepareGenerator(self):
 161     """Download, and extract au-generate.zip into self.generator_dir."""
 162     generator_uri = self._GeneratorUri()
 163
 164     logging.info('Preparing au-generate.zip from %s.', generator_uri)
 165
 166     # Extract zipped delta generator files to the expected directory.
 167     tmp_zip = self.cache.GetFileInTempFile(generator_uri)
 168     utils.RunCommand(['unzip', '-o', '-d', self.generator_dir, tmp_zip.name],
 169                      redirect_stdout=True, redirect_stderr=True)
 170     tmp_zip.close()
 171
 172   def _RunGeneratorCmd(self, cmd):
 173     """Wrapper for RunCommand for programs in self.generator_dir.
 174
 175     Adjusts the program name for the current self.au_generator directory, and
 176     sets up the special requirements needed for these 'out of chroot'
 177     programs. Will automatically log the command output if execution resulted
 178     in a nonzero exit code. Note that the command's stdout and stderr are
 179     combined into a single string. This also sets the TMPDIR variable
 180     accordingly in the spawned process' environment.
 181
 182     Args:
 183       cmd: Program and argument list in a list. ['delta_generator', '--help']
 184
 185     Returns:
 186       The output of the executed command.
 187
 188     Raises:
 189       cros_build_lib.RunCommandError if the command exited with a nonzero code.
 190
 191     """
 192     # Adjust the command name to match the directory it's in.
 193     cmd[0] = os.path.join(self.generator_dir, cmd[0])
 194
 195     # Modify the PATH and TMPDIR when running the script.
 196     extra_env = {
 197         'PATH': utils.PathPrepend(self.generator_dir),
 198         'TMPDIR': self.work_dir}
 199
 200     # Run the command.
 201     result = cros_build_lib.RunCommand(
 202         cmd,
 203         cwd=self.generator_dir,
 204         redirect_stdout=True,
 205         combine_stdout_stderr=True,
 206         error_code_ok=True,
 207         extra_env=extra_env)
 208
 209     # Dump error output and raise an exception if things went awry.
 210     if result.returncode:
 211       logging.error('Nonzero exit code (%d), dumping command output:\n%s',
 212                     result.returncode, result.output)
 213       raise cros_build_lib.RunCommandError(
 214           'Command failed: %s (cwd=%s)' % (' '.join(cmd), self.generator_dir),
 215           result)
 216
 217     return result.output
 218
 219   @staticmethod
 220   def _BuildArg(flag, dict_obj, key, default=None):
 221     """Returns a command-line argument iff its value is present in a dictionary.
 222
 223     Args:
 224       flag: the flag name to use with the argument value, e.g. --foo; if None
 225             or an empty string, no flag will be used
 226       dict_obj: a dictionary mapping possible keys to values
 227       key: the key of interest; e.g. 'foo'
 228       default: a default value to use if key is not in dict_obj (optional)
 229
 230     Returns:
 231       If dict_obj[key] contains a non-False value or default is non-False,
 232       returns a list containing the flag and value arguments (e.g. ['--foo',
 233       'bar']), unless flag is empty/None, in which case returns a list
 234       containing only the value argument (e.g.  ['bar']). Otherwise, returns an
 235       empty list.
 236
 237     """
 238     arg_list = []
 239     val = dict_obj.get(key) or default
 240     if val:
 241       arg_list = [str(val)]
 242       if flag:
 243         arg_list.insert(0, flag)
 244
 245     return arg_list
 246
 247   def _PrepareImage(self, image, image_file):
 248     """Download an prepare an image for delta generation.
 249
 250     Preparation includes downloading, extracting and converting the image into
 251     an on-disk format, as necessary.
 252
 253     Args:
 254       image: an object representing the image we're processing
 255       image_file: file into which the prepared image should be copied.
 256     """
 257
 258     logging.info('Preparing image from %s as %s', image.uri, image_file)
 259
 260     # Figure out what we're downloading and how to handle it.
 261     image_handling_by_type = {
 262         'signed': (None, True),
 263         'test': (self.TEST_IMAGE_NAME, False),
 264         'recovery': (self.RECOVERY_IMAGE_NAME, True),
 265     }
 266     extract_file, _ = image_handling_by_type[image.get('image_type', 'signed')]
 267
 268     # Are we donwloading an archive that contains the image?
 269     if extract_file:
 270       # Archive will be downloaded to a temporary location.
 271       with tempfile.NamedTemporaryFile(
 272           prefix='image-archive-', suffix='.tar.xz', dir=self.work_dir,
 273           delete=False) as temp_file:
 274         download_file = temp_file.name
 275     else:
 276       download_file = image_file
 277
 278     # Download the image file or archive.
 279     self.cache.GetFileCopy(image.uri, download_file)
 280
 281     # If we downloaded an archive, extract the image file from it.
 282     if extract_file:
 283       cmd = ['tar', '-xJf', download_file, extract_file]
 284       cros_build_lib.RunCommand(cmd, cwd=self.work_dir)
 285
 286       # Rename it into the desired image name.
 287       shutil.move(os.path.join(self.work_dir, extract_file), image_file)
 288
 289       # It's safe to delete the archive at this point.
 290       os.remove(download_file)
 291
 292   def _GenerateUnsignedPayload(self):
 293     """Generate the unsigned delta into self.payload_file."""
 294     # Note that the command run here requires sudo access.
 295
 296     logging.info('Generating unsigned payload as %s', self.payload_file)
 297
 298     tgt_image = self.payload.tgt_image
 299     cmd = ['cros_generate_update_payload',
 300            '--outside_chroot',
 301            '--output', self.payload_file,
 302            '--image', self.tgt_image_file,
 303            '--channel', tgt_image.channel,
 304            '--board', tgt_image.board,
 305            '--version', tgt_image.version,
 306           ]
 307     cmd += self._BuildArg('--key', tgt_image, 'key', default='test')
 308     cmd += self._BuildArg('--build_channel', tgt_image, 'image_channel',
 309                           default=tgt_image.channel)
 310     cmd += self._BuildArg('--build_version', tgt_image, 'image_version',
 311                           default=tgt_image.version)
 312
 313     if self.payload.src_image:
 314       src_image = self.payload.src_image
 315       cmd += ['--src_image', self.src_image_file,
 316               '--src_channel', src_image.channel,
 317               '--src_board', src_image.board,
 318               '--src_version', src_image.version,
 319              ]
 320       cmd += self._BuildArg('--src_key', src_image, 'key', default='test')
 321       cmd += self._BuildArg('--src_build_channel', src_image, 'image_channel',
 322                             default=src_image.channel)
 323       cmd += self._BuildArg('--src_build_version', src_image, 'image_version',
 324                             default=src_image.version)
 325
 326     delta_log = self._RunGeneratorCmd(cmd)
 327     self._StoreDeltaLog(delta_log)
 328
 329   def _GenPayloadHash(self):
 330     """Generate a hash of payload and metadata.
 331
 332     Works from an unsigned update payload.
 333
 334     Returns:
 335       payload_hash as a string.
 336     """
 337     logging.info('Calculating payload hashes on %s.', self.payload_file)
 338
 339     # How big will the signatures be.
 340     signature_sizes = [str(size) for size in self.PAYLOAD_SIGNATURE_SIZES_BYTES]
 341
 342     with tempfile.NamedTemporaryFile('rb') as payload_hash_file:
 343       cmd = ['delta_generator',
 344              '-in_file', self.payload_file,
 345              '-out_hash_file', payload_hash_file.name,
 346              '-signature_size', ':'.join(signature_sizes)]
 347
 348       self._RunGeneratorCmd(cmd)
 349       return payload_hash_file.read()
 350
 351   def _GenMetadataHash(self):
 352     """Generate a hash of payload and metadata.
 353
 354     Works from an unsigned update payload.
 355
 356     Returns:
 357       metadata_hash as a string.
 358     """
 359     logging.info('Calculating payload hashes on %s.', self.payload_file)
 360
 361     # How big will the signatures be.
 362     signature_sizes = [str(size) for size in self.PAYLOAD_SIGNATURE_SIZES_BYTES]
 363
 364     with tempfile.NamedTemporaryFile('rb') as metadata_hash_file:
 365       cmd = ['delta_generator',
 366              '-in_file', self.payload_file,
 367              '-out_metadata_hash_file', metadata_hash_file.name,
 368              '-signature_size', ':'.join(signature_sizes)]
 369
 370       self._RunGeneratorCmd(cmd)
 371       return metadata_hash_file.read()
 372
 373   def _GenerateSignerResultsError(self, format_str, *args):
 374     """Helper for reporting errors with signer results."""
 375     msg = format_str % args
 376     logging.error(msg)
 377     raise UnexpectedSignerResultsError(msg)
 378
 379   def _SignHashes(self, hashes):
 380     """Get the signer to sign the hashes with the update payload key via GS.
 381
 382     May sign each hash with more than one key, based on how many keysets are
 383     required.
 384
 385     Args:
 386       hashes: List of hashes to be signed.
 387
 388     Returns:
 389       List of lists which contain each signed hash.
 390       [[hash_1_sig_1, hash_1_sig_2], [hash_2_sig_1, hash_2_sig_2]]
 391     """
 392     logging.info('Signing payload hashes with %s.',
 393                  ', '.join(self.PAYLOAD_SIGNATURE_KEYSETS))
 394
 395     # Results look like:
 396     #  [[hash_1_sig_1, hash_1_sig_2], [hash_2_sig_1, hash_2_sig_2]]
 397     hashes_sigs = self.signer.GetHashSignatures(
 398         hashes,
 399         keysets=self.PAYLOAD_SIGNATURE_KEYSETS)
 400
 401     if hashes_sigs is None:
 402       self._GenerateSignerResultsError('Signing of hashes failed')
 403     if len(hashes_sigs) != len(hashes):
 404       self._GenerateSignerResultsError(
 405           'Count of hashes signed (%d) != Count of hashes (%d).',
 406           len(hashes_sigs),
 407           len(hashes))
 408
 409     # Make sure that the results we get back the expected number of signatures.
 410     for hash_sigs in hashes_sigs:
 411       # Make sure each hash has the right number of signatures.
 412       if len(hash_sigs) != len(self.PAYLOAD_SIGNATURE_SIZES_BYTES):
 413         self._GenerateSignerResultsError(
 414             'Signature count (%d) != Expected signature count (%d)',
 415             len(hash_sigs),
 416             len(self.PAYLOAD_SIGNATURE_SIZES_BYTES))
 417
 418       # Make sure each hash signature is the expected size.
 419       for sig, sig_size in zip(hash_sigs, self.PAYLOAD_SIGNATURE_SIZES_BYTES):
 420         if len(sig) != sig_size:
 421           self._GenerateSignerResultsError(
 422               'Signature size (%d) != expected size(%d)',
 423               len(sig),
 424               sig_size)
 425
 426     return hashes_sigs
 427
 428   def _InsertPayloadSignatures(self, signatures):
 429     """Put payload signatures into the payload they sign.
 430
 431     Args:
 432       signatures: List of signatures for the payload.
 433     """
 434     logging.info('Inserting payload signatures into %s.',
 435                  self.signed_payload_file)
 436
 437     signature_files = [utils.CreateTempFileWithContents(s) for s in signatures]
 438     signature_file_names = [f.name for f in signature_files]
 439
 440     cmd = ['delta_generator',
 441            '-in_file', self.payload_file,
 442            '-signature_file', ':'.join(signature_file_names),
 443            '-out_file', self.signed_payload_file]
 444
 445     self._RunGeneratorCmd(cmd)
 446
 447     for f in signature_files:
 448       f.close()
 449
 450   def _StoreMetadataSignatures(self, signatures):
 451     """Store metadata signatures related to the payload.
 452
 453     Our current format for saving metadata signatures only supports a single
 454     signature at this time.
 455
 456     Args:
 457       signatures: A list of metadata signatures in binary string format.
 458     """
 459     if len(signatures) != 1:
 460       self._GenerateSignerResultsError(
 461           'Received %d metadata signatures, only a single signature supported.',
 462           len(signatures))
 463
 464     logging.info('Saving metadata signatures in %s.',
 465                  self.metadata_signature_file)
 466
 467     encoded_signature = base64.b64encode(signatures[0])
 468
 469     with open(self.metadata_signature_file, 'w+') as f:
 470       f.write(encoded_signature)
 471
 472   def _StorePayloadJson(self, metadata_signatures):
 473     """Generate the payload description json file.
 474
 475     The payload description contains a dictionary with the following
 476     fields populated.
 477
 478     {
 479       "version": 1,
 480       "sha1_hex": <payload sha1 hash as a hex encoded string>,
 481       "sha256_hex": <payload sha256 hash as a hex encoded string>,
 482       "metadata_signature": <metadata signature as base64 encoded string or nil>
 483     }
 484
 485     Args:
 486       metadata_signatures: A list of signatures in binary string format.
 487     """
 488     # Locate everything we put in the json.
 489     sha1_hex, sha256_hex = filelib.ShaSums(self.payload_file)
 490
 491     metadata_signature = None
 492     if metadata_signatures:
 493       if len(metadata_signatures) != 1:
 494         self._GenerateSignerResultsError(
 495             'Received %d metadata signatures, only one supported.',
 496             len(metadata_signatures))
 497       metadata_signature = base64.b64encode(metadata_signatures[0])
 498
 499     # Bundle it up in a map matching the Json format.
 500     # Increment DESCRIPTION_FILE_VERSION, if changing this map.
 501     payload_map = {
 502       'version': DESCRIPTION_FILE_VERSION,
 503       'sha1_hex': sha1_hex,
 504       'sha256_hex': sha256_hex,
 505       'metadata_signature': metadata_signature,
 506     }
 507
 508     # Convert to Json.
 509     payload_json = json.dumps(payload_map, sort_keys=True)
 510
 511     # Write out the results.
 512     osutils.WriteFile(self.description_file, payload_json)
 513
 514   def _StoreDeltaLog(self, delta_log):
 515     """Store delta log related to the payload.
 516
 517     Write out the delta log to a known file name. Mostly in it's own function
 518     to simplify unittest mocks.
 519
 520     Args:
 521       delta_log: The delta logs as a single string.
 522     """
 523     with open(self.delta_log_file, 'w+') as f:
 524       f.write(delta_log)
 525
 526   def _SignPayload(self):
 527     """Wrap all the steps for signing an existing payload.
 528
 529     Returns:
 530       List of payload signatures, List of metadata signatures.
 531       """
 532     # Create hashes to sign.
 533     payload_hash = self._GenPayloadHash()
 534     metadata_hash = self._GenMetadataHash()
 535
 536     # Sign them.
 537     payload_signatures, metadata_signatures = self._SignHashes(
 538         [payload_hash, metadata_hash])
 539
 540     # Insert payload signature(s).
 541     self._InsertPayloadSignatures(payload_signatures)
 542
 543     # Store Metadata signature(s).
 544     self._StoreMetadataSignatures(metadata_signatures)
 545
 546     return (payload_signatures, metadata_signatures)
 547
 548   def _Create(self):
 549     """Create a given payload, if it doesn't already exist."""
 550
 551     logging.info('Generating %s payload %s',
 552                  'delta' if self.payload.src_image else 'full', self.payload)
 553
 554     # Fetch and extract the delta generator.
 555     self._PrepareGenerator()
 556
 557     # Fetch and prepare the tgt image.
 558     self._PrepareImage(self.payload.tgt_image, self.tgt_image_file)
 559
 560     # Fetch and prepare the src image.
 561     if self.payload.src_image:
 562       self._PrepareImage(self.payload.src_image, self.src_image_file)
 563
 564     # Generate the unsigned payload.
 565     self._GenerateUnsignedPayload()
 566
 567     # Sign the payload, if needed.
 568     metadata_signatures = None
 569     if self.signer:
 570       _, metadata_signatures = self._SignPayload()
 571
 572     # Store hash and signatures json.
 573     self._StorePayloadJson(metadata_signatures)
 574
 575
 576   def _CheckPayloadIntegrity(self, payload, is_delta, metadata_sig_file_name):
 577     """Checks the integrity of a generated payload.
 578
 579     Args:
 580       payload: an pre-initialized update_payload.Payload object.
 581       is_delta: whether or not this is a delta payload (Boolean).
 582       metadata_sig_file_name: metadata signature file.
 583
 584     Raises:
 585       PayloadVerificationError: when an error is encountered.
 586     """
 587     logging.info('Checking payload integrity')
 588     with utils.CheckedOpen(metadata_sig_file_name) as metadata_sig_file:
 589       try:
 590         # TODO(garnold)(chromium:243559) partition sizes should be embedded in
 591         # the payload; ditch the default values once it's done.
 592         # TODO(garnold)(chromium:261417) this disables the check for unmoved
 593         # blocks in MOVE sequences, which is an inefficiency but not
 594         # necessarily a problem.  It should be re-enabled once the delta
 595         # generator can optimize away such cases.
 596         payload.Check(metadata_sig_file=metadata_sig_file,
 597                       assert_type=('delta' if is_delta else 'full'),
 598                       rootfs_part_size=self._DEFAULT_ROOTFS_PART_SIZE,
 599                       kernel_part_size=self._DEFAULT_KERNEL_PART_SIZE,
 600                       disabled_tests=['move-same-src-dst-block'])
 601       except self._update_payload.PayloadError as e:
 602         raise PayloadVerificationError(
 603             'Payload integrity check failed: %s' % e)
 604
 605   def _ApplyPayload(self, payload, is_delta):
 606     """Applies a generated payload and verifies the result.
 607
 608     Args:
 609       payload: an pre-initialized update_payload.Payload object.
 610       is_delta: whether or not this is a delta payload (Boolean).
 611
 612     Raises:
 613       PayloadVerificationError: when an error occurs.
 614     """
 615     # Extract the source/target kernel/rootfs partitions.
 616     # TODO(garnold)(chromium:243561) this is a redundant operation as the
 617     # partitions are already extracted (in some form) for the purpose of
 618     # payload generation. We should only do this once.
 619     cmd = ['cros_generate_update_payload',
 620            '--outside_chroot',
 621            '--extract',
 622            '--image', self.tgt_image_file]
 623     part_files = {}
 624     part_files['new_kernel_part'] = self._DEFAULT_NEW_KERN_PART
 625     part_files['new_rootfs_part'] = self._DEFAULT_NEW_ROOT_PART
 626     if is_delta:
 627       cmd += ['--src_image', self.src_image_file]
 628       part_files['old_kernel_part'] = self._DEFAULT_OLD_KERN_PART
 629       part_files['old_rootfs_part'] = self._DEFAULT_OLD_ROOT_PART
 630
 631     self._RunGeneratorCmd(cmd)
 632
 633     for part_name, part_file in part_files.items():
 634       part_file = os.path.join(self.generator_dir, part_file)
 635       if not os.path.isfile(part_file):
 636         raise PayloadVerificationError('Failed to extract partition (%s)' %
 637                                        part_file)
 638       part_files[part_name] = part_file
 639
 640     # Apply the payload and verify the result; make sure to pass in the
 641     # explicit path to the bspatch binary in the au-generator directory (the
 642     # one we need to be using), and not to depend on PATH resolution etc. Also
 643     # note that we instruct the call to generate files with a .test suffix,
 644     # which we can later compare to the actual target partition (as it was
 645     # extracted from the target image above).
 646     logging.info('Applying %s payload and verifying result',
 647                  'delta' if is_delta else 'full')
 648     ref_new_kern_part = part_files['new_kernel_part']
 649     part_files['new_kernel_part'] += '.test'
 650     ref_new_root_part = part_files['new_rootfs_part']
 651     part_files['new_rootfs_part'] += '.test'
 652     bspatch_path = os.path.join(self.generator_dir, 'bspatch')
 653     try:
 654       payload.Apply(bspatch_path=bspatch_path, **part_files)
 655     except self._update_payload.PayloadError as e:
 656       raise PayloadVerificationError('Payload failed to apply: %s' % e)
 657
 658     # Prior to comparing, remove unused space past the filesystem boundary
 659     # in the extracted target partitions.
 660     filelib.TruncateToSize(ref_new_kern_part,
 661                            os.path.getsize(part_files['new_kernel_part']))
 662     filelib.TruncateToSize(ref_new_root_part,
 663                            os.path.getsize(part_files['new_rootfs_part']))
 664
 665     # Compare resulting partitions with the ones from the target image.
 666     if not filecmp.cmp(ref_new_kern_part, part_files['new_kernel_part']):
 667       raise PayloadVerificationError('Resulting kernel partition corrupted')
 668     if not filecmp.cmp(ref_new_root_part, part_files['new_rootfs_part']):
 669       raise PayloadVerificationError('Resulting rootfs partition corrupted')
 670
 671   def _VerifyPayload(self):
 672     """Checks the integrity of the generated payload.
 673
 674     Raises:
 675       PayloadVerificationError when the payload fails to verify.
 676     """
 677     if self.signer:
 678       payload_file_name = self.signed_payload_file
 679       metadata_sig_file_name = self.metadata_signature_file
 680     else:
 681       payload_file_name = self.payload_file
 682       metadata_sig_file_name = None
 683
 684     with open(payload_file_name) as payload_file:
 685       payload = self._update_payload.Payload(payload_file)
 686       is_delta = bool(self.payload.src_image)
 687       try:
 688         payload.Init()
 689
 690         # First, verify the payload's integrity.
 691         self._CheckPayloadIntegrity(payload, is_delta, metadata_sig_file_name)
 692
 693         # Second, try to apply the payload and check the result.
 694         self._ApplyPayload(payload, is_delta)
 695
 696       except self._update_payload.PayloadError as e:
 697         raise PayloadVerificationError('Payload failed to verify: %s' % e)
 698
 699   def _UploadResults(self):
 700     """Copy the payload generation results to the specified destination."""
 701
 702     logging.info('Uploading payload to %s.', self.payload.uri)
 703
 704     # Deliver the payload to the final location.
 705     if self.signer:
 706       urilib.Copy(self.signed_payload_file, self.payload.uri)
 707       urilib.Copy(self.metadata_signature_file,
 708                   self._MetadataUri(self.payload.uri))
 709     else:
 710       urilib.Copy(self.payload_file, self.payload.uri)
 711
 712     # Upload payload related artifacts.
 713     urilib.Copy(self.delta_log_file, self._DeltaLogsUri(self.payload.uri))
 714     urilib.Copy(self.description_file, self._JsonUri(self.payload.uri))
 715
 716   def Run(self):
 717     """Create, verify and upload the results."""
 718     self._drm(self._Create)
 719     if self._verify:
 720       self._drm(self._VerifyPayload)
 721     self._drm(self._UploadResults)
 722
 723
 724 def DefaultPayloadUri(payload, random_str=None):
 725   """Compute the default output URI for a payload.
 726
 727   For a glob that matches all potential URIs for this
 728   payload, pass in a random_str of '*'.
 729
 730   Args:
 731     payload: gspaths.Payload instance.
 732     random_str: A hook to force a specific random_str. None means generate it.
 733
 734   Returns:
 735     Default URI for the payload.
 736   """
 737   src_version = None
 738   if payload.src_image:
 739     src_version = payload.src_image['version']
 740
 741   if payload.tgt_image.get('image_type', 'signed') == 'signed':
 742     # Signed payload.
 743     return gspaths.ChromeosReleases.PayloadUri(
 744         channel=payload.tgt_image.channel,
 745         board=payload.tgt_image.board,
 746         version=payload.tgt_image.version,
 747         random_str=random_str,
 748         key=payload.tgt_image.key,
 749         image_channel=payload.tgt_image.image_channel,
 750         image_version=payload.tgt_image.image_version,
 751         src_version=src_version,
 752         bucket=payload.tgt_image.bucket)
 753   else:
 754     # Unsigned test payload.
 755     return gspaths.ChromeosReleases.PayloadUri(
 756         channel=payload.tgt_image.channel,
 757         board=payload.tgt_image.board,
 758         version=payload.tgt_image.version,
 759         random_str=random_str,
 760         src_version=src_version,
 761         bucket=payload.tgt_image.bucket)
 762
 763
 764 def SetPayloadUri(payload, uri):
 765   """Sets (overrides) the URI in a payload object.
 766
 767   Args:
 768     payload: gspaths.Payload instance.
 769     uri: A URI (string) to the payload file.
 770   """
 771   payload.uri = uri
 772
 773
 774 def FillInPayloadUri(payload, random_str=None):
 775   """Fill in default output URI for a payload if missing.
 776
 777   Args:
 778     payload: gspaths.Payload instance.
 779     random_str: A hook to force a specific random_str. None means generate it.
 780   """
 781   if not payload.uri:
 782     SetPayloadUri(payload, DefaultPayloadUri(payload, random_str))
 783
 784
 785 def _FilterNonPayloadUris(payload_uris):
 786   """Filters out non-payloads from a list of GS URIs.
 787
 788   This essentially filters out known auxiliary artifacts whose names resemble /
 789   derive from a respective payload name, such as files with .log and
 790   .metadata-signature extensions.
 791
 792   Args:
 793     payload_uris: a list of GS URIs (potentially) corresopnding to payloads
 794
 795   Returns:
 796     A filtered list of URIs.
 797   """
 798   return [uri for uri in payload_uris
 799           if not (uri.endswith('.log') or uri.endswith('.metadata-signature'))]
 800
 801
 802 def FindExistingPayloads(payload):
 803   """Look to see if any matching payloads already exist.
 804
 805   Since payload names contain a random component, there can be multiple
 806   names for a given payload. This function lists all existing payloads
 807   that match the default URI for the given payload.
 808
 809   Args:
 810     payload: gspaths.Payload instance.
 811
 812   Returns:
 813     List of URIs for existing payloads that match the default payload pattern.
 814   """
 815   search_uri = DefaultPayloadUri(payload, random_str='*')
 816   return _FilterNonPayloadUris(urilib.ListFiles(search_uri))
 817
 818
 819 def FindCacheDir(work_dir=None):
 820   """Helper for deciding what cache directory to use.
 821
 822   Args:
 823     work_dir: Directory that contains ALL work files, cache will
 824               be created inside it, if present.
 825
 826   Returns:
 827     Returns a directory suitable for use with a DownloadCache. Will
 828     always be consistent if a consistent work_dir is passed in.
 829   """
 830   # Discover which directory to use for caching
 831   if work_dir:
 832     return os.path.join(work_dir, 'cache')
 833   else:
 834     return '/usr/local/google/payloads'
 835
 836
 837 def CreateAndUploadPayload(payload, cache, work_dir, sign=True, verify=True,
 838                            dry_run=False, au_generator_uri=None):
 839   """Helper to create a PaygenPayloadLib instance and use it.
 840
 841   Args:
 842     payload: An instance of utils.Payload describing the payload to generate.
 843     cache: An instance of DownloadCache for retrieving files.
 844     work_dir: A working directory that can hold scratch files. Will be cleaned
 845               up when done, and won't interfere with other users. None for /tmp.
 846     sign: Boolean saying if the payload should be signed (normally, you do).
 847     verify: whether the payload should be verified (default: True)
 848     dry_run: don't perform actual work
 849     au_generator_uri: URI to override standard au_generator.zip rules.
 850   """
 851   with osutils.TempDir(prefix='paygen_payload.', base_dir=work_dir) as gen_dir:
 852     logging.info('* Starting payload generation')
 853     start_time = datetime.datetime.now()
 854
 855     _PaygenPayload(payload, cache, gen_dir, sign, verify, au_generator_uri,
 856                    dry_run=dry_run).Run()
 857
 858     end_time = datetime.datetime.now()
 859     logging.info('* Finished payload generation in %s', end_time - start_time)