Update test harness to take in optional public and private keys to sign payloads.

bin/cros_au_test_harness.py

 #!/usr/bin/python
+import shutil
 
 # Copyright (c) 2011 The Chromium OS Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """This module runs a suite of Auto Update tests.
 
   The tests can be run on either a virtual machine or actual device depending
   on parameters given.  Specific tests can be run by invoking --test_prefix.
   Verbose is useful for many of the tests if you want to see individual commands
   being run during the update process.
 """
 
 import optparse
 import os
 import re
 import subprocess
 import sys
+import tempfile
 import threading
 import time
 import unittest
 import urllib
 
 sys.path.append(os.path.join(os.path.dirname(__file__), '../lib'))
 from cros_build_lib import Die
 from cros_build_lib import GetIPAddress
 from cros_build_lib import Info
 from cros_build_lib import ReinterpretPathForChroot
@@ skipping 16 matching lines @@
 class AUTest(object):
   """Abstract interface that defines an Auto Update test."""
   verbose = False
 
   def setUp(self):
     unittest.TestCase.setUp(self)
     # Set these up as they are used often.
     self.crosutils = os.path.join(os.path.dirname(__file__), '..')
     self.crosutilsbin = os.path.join(os.path.dirname(__file__))
     self.download_folder = os.path.join(self.crosutils, 'latest_download')
+    self.vm_image_path = None
     if not os.path.exists(self.download_folder):
       os.makedirs(self.download_folder)
 
   # -------- Helper functions ---------
 
+  def _PrepareRealBase(self, image_path):
+    self.PerformUpdate(image_path)
+
+  def _PrepareVMBase(self, image_path):
+    # VM Constants.
+    FULL_VDISK_SIZE = 6072
+    FULL_STATEFULFS_SIZE = 3074
+    # Needed for VM delta updates.  We need to use the qemu image rather
+    # than the base image on a first update.  By tracking the first_update
+    # we can set src_image to the qemu form of the base image when
+    # performing generating the delta payload.
+    self._first_update = True
+    self.vm_image_path = '%s/chromiumos_qemu_image.bin' % os.path.dirname(
+        image_path)
+    if not os.path.exists(self.vm_image_path):
+      Info('Creating %s' % self.vm_image_path)
+      RunCommand(['%s/image_to_vm.sh' % self.crosutils,
+                  '--full',
+                  '--from=%s' % ReinterpretPathForChroot(
+                      os.path.dirname(image_path)),
+                  '--vdisk_size=%s' % FULL_VDISK_SIZE,
+                  '--statefulfs_size=%s' % FULL_STATEFULFS_SIZE,
+                  '--board=%s' % self.board,
+                  '--test_image'], enter_chroot=True)
+
+    Info('Using %s as base' % self.vm_image_path)
+    self.assertTrue(os.path.exists(self.vm_image_path))
+
+  def AppendUpdateFlags(self, cmd, image_path, src_image_path, proxy_port,
+                        private_key_path):
+    """Appends common args to an update cmd defined by an array.
+
+    Modifies cmd in places by appending appropriate items given args.
+    """
+    if proxy_port: cmd.append('--proxy_port=%s' % proxy_port)
+
+    # Get pregenerated update if we have one.
+    update_id = _GenerateUpdateId(target=image_path, src=src_image_path,
+                                  key=private_key_path)
+    cache_path = dev_server_cache[update_id]
+    if cache_path:
+      update_url = DevServerWrapper.GetDevServerURL(proxy_port, cache_path)
+      cmd.append('--update_url=%s' % update_url)
+    else:
+      cmd.append('--image=%s' % image_path)
+      if src_image_path: cmd.append('--src_image=%s' % src_image_path)
+
+  def RunUpdateCmd(self, cmd):
+    """Runs the given update cmd given verbose options.
+
+    Raises an UpdateException if the update fails.
+    """
+    if self.verbose:
+      try:
+        RunCommand(cmd)
+      except Exception, e:
+        raise UpdateException(1, e.message)
+    else:
+      (code, stdout, stderr) = RunCommandCaptureOutput(cmd)
+      if code != 0:
+        raise UpdateException(code, stdout)
+
   def GetStatefulChangeFlag(self, stateful_change):
     """Returns the flag to pass to image_to_vm for the stateful change."""
     stateful_change_flag = ''
     if stateful_change:
       stateful_change_flag = '--stateful_update_flag=%s' % stateful_change
 
     return stateful_change_flag
 
   def _ParseGenerateTestReportOutput(self, output):
     """Returns the percentage of tests that passed based on output."""
@@ skipping 23 matching lines @@
     Info('Output from VerifyImage():')
     print >> sys.stderr, output
     sys.stderr.flush()
     percent_passed = self._ParseGenerateTestReportOutput(output)
     Info('Percent passed: %d vs. Percent required: %d' % (
         percent_passed, percent_required_to_pass))
     unittest.assertTrue(percent_passed >= percent_required_to_pass)
     return percent_passed
 
   def PerformUpdate(self, image_path, src_image_path='', stateful_change='old',
-                    proxy_port=None):
+                    proxy_port=None, private_key_path=None):
     """Performs an update using  _UpdateImage and reports any error.
 
     Subclasses should not override this method but override _UpdateImage
     instead.
 
     Args:
       image_path:  Path to the image to update with.  This image must be a test
         image.
       src_image_path:  Optional.  If set, perform a delta update using the
         image specified by the path as the source image.
       stateful_change: How to modify the stateful partition.  Values are:
           'old':  Don't modify stateful partition.  Just update normally.
           'clean':  Uses clobber-state to wipe the stateful partition with the
             exception of code needed for ssh.
       proxy_port:  Port to have the client connect to. For use with
         CrosTestProxy.
     Raises an UpdateException if _UpdateImage returns an error.
     """
     try:
-      if not self.use_delta_updates:
-        src_image_path = ''
+      if not self.use_delta_updates: src_image_path = ''
+      if private_key_path:
+        key_to_use = private_key_path
+      else:
+        key_to_use = self.private_key
 
-      self._UpdateImage(image_path, src_image_path, stateful_change, proxy_port)
+      self._UpdateImage(image_path, src_image_path, stateful_change, proxy_port,
+                        key_to_use)
     except UpdateException as err:
       # If the update fails, print it out
       Warning(err.stdout)
       raise
 
   def AttemptUpdateWithPayloadExpectedFailure(self, payload, expected_msg):
     """Attempt a payload update, expect it to fail with expected log"""
     try:
       self._UpdateUsingPayload(payload)
     except UpdateException as err:
@@ skipping 33 matching lines @@
   def ProcessOptions(cls, parser, options):
     """Processes options.
 
     Static method that should be called from main.  Subclasses should also
     call their parent method if they override it.
     """
     cls.verbose = options.verbose
     cls.base_image_path = options.base_image
     cls.target_image_path = options.target_image
     cls.use_delta_updates = options.delta
+    cls.board = options.board
+    cls.private_key = options.private_key
     if options.quick_test:
       cls.verify_suite = 'build_RootFilesystemSize'
     else:
       cls.verify_suite = 'suite_Smoke'
 
     # Sanity checks.
     if not cls.base_image_path:
       parser.error('Need path to base image for vm.')
     elif not os.path.exists(cls.base_image_path):
       Die('%s does not exist' % cls.base_image_path)
 
     if not cls.target_image_path:
       parser.error('Need path to target image to update with.')
     elif not os.path.exists(cls.target_image_path):
       Die('%s does not exist' % cls.target_image_path)
 
   def PrepareBase(self, image_path):
     """Prepares target with base_image_path."""
     pass
 
   def _UpdateImage(self, image_path, src_image_path='', stateful_change='old',
-                   proxy_port=None):
+                   proxy_port=None, private_key_path=None):
     """Implementation of an actual update.
 
     See PerformUpdate for description of args.  Subclasses must override this
     method with the correct update procedure for the class.
     """
     pass
 
   def _UpdateUsingPayload(self, update_path, stateful_change='old',
                          proxy_port=None):
     """Updates target with the pre-generated update stored in update_path.
@@ skipping 190 matching lines @@
   def ProcessOptions(cls, parser, options):
     """Processes non-vm-specific options."""
     AUTest.ProcessOptions(parser, options)
     cls.remote = options.remote
 
     if not cls.remote:
       parser.error('We require a remote address for real tests.')
 
   def PrepareBase(self, image_path):
     """Auto-update to base image to prepare for test."""
-    self.PerformUpdate(image_path)
+    _PrepareRealBase(image_path)
 
   def _UpdateImage(self, image_path, src_image_path='', stateful_change='old',
-                   proxy_port=None):
+                   proxy_port=None, private_key_path=None):
     """Updates a remote image using image_to_live.sh."""
     stateful_change_flag = self.GetStatefulChangeFlag(stateful_change)
     cmd = ['%s/image_to_live.sh' % self.crosutils,
-           '--image=%s' % image_path,
            '--remote=%s' % self.remote,
            stateful_change_flag,
            '--verify',
-           '--src_image=%s' % src_image_path
           ]
-
-    if proxy_port:
-      cmd.append('--proxy_port=%s' % proxy_port)
-
-    if self.verbose:
-      try:
-        RunCommand(cmd)
-      except Exception, e:
-        raise UpdateException(1, e.message)
-    else:
-      (code, stdout, stderr) = RunCommandCaptureOutput(cmd)
-      if code != 0:
-        raise UpdateException(code, stdout)
+    self.AppendUpdateFlags(cmd, image_path, src_image_path, proxy_port,
+                           private_key_path)
+    self.RunUpdateCmd(cmd)
 
   def _UpdateUsingPayload(self, update_path, stateful_change='old',
                          proxy_port=None):
     """Updates a remote image using image_to_live.sh."""
     stateful_change_flag = self.GetStatefulChangeFlag(stateful_change)
     cmd = ['%s/image_to_live.sh' % self.crosutils,
            '--payload=%s' % update_path,
            '--remote=%s' % self.remote,
            stateful_change_flag,
            '--verify',
           ]
-
-    if proxy_port:
-      cmd.append('--proxy_port=%s' % proxy_port)
-
-    if self.verbose:
-      try:
-        RunCommand(cmd)
-      except Exception, e:
-        raise UpdateException(1, e.message)
-    else:
-      (code, stdout, stderr) = RunCommandCaptureOutput(cmd)
-      if code != 0:
-        raise UpdateException(code, stdout)
+    if proxy_port: cmd.append('--proxy_port=%s' % proxy_port)
+    self.RunUpdateCmd(cmd)
 
   def VerifyImage(self, percent_required_to_pass):
     """Verifies an image using run_remote_tests.sh with verification suite."""
     output = RunCommand([
         '%s/run_remote_tests.sh' % self.crosutils,
         '--remote=%s' % self.remote,
         self.verify_suite,
        ], error_ok=True, enter_chroot=False, redirect_stdout=True)
     return self.AssertEnoughTestsPassed(self, output, percent_required_to_pass)
 
 
 class VirtualAUTest(unittest.TestCase, AUTest):
   """Test harness for updating virtual machines."""
 
-  # VM Constants.
-  _FULL_VDISK_SIZE = 6072
-  _FULL_STATEFULFS_SIZE = 3074
-
   # Class variables used to acquire individual VM variables per test.
   _vm_lock = threading.Lock()
   _next_port = 9222
 
   def _KillExistingVM(self, pid_file):
     if os.path.exists(pid_file):
       Warning('Existing %s found.  Deleting and killing process' %
               pid_file)
       RunCommand(['./cros_stop_vm', '--kvm_pid=%s' % pid_file],
                  cwd=self.crosutilsbin)
 
     assert not os.path.exists(pid_file)
 
   def _AcquireUniquePortAndPidFile(self):
     """Acquires unique ssh port and pid file for VM."""
     with VirtualAUTest._vm_lock:
       self._ssh_port = VirtualAUTest._next_port
       self._kvm_pid_file = '/tmp/kvm.%d' % self._ssh_port
       VirtualAUTest._next_port += 1
 
   def setUp(self):
     """Unit test overriden method.  Is called before every test."""
     AUTest.setUp(self)
-    self.vm_image_path = None
     self._AcquireUniquePortAndPidFile()
     self._KillExistingVM(self._kvm_pid_file)
 
   def tearDown(self):
     self._KillExistingVM(self._kvm_pid_file)
 
   @classmethod
   def ProcessOptions(cls, parser, options):
     """Processes vm-specific options."""
     AUTest.ProcessOptions(parser, options)
-    cls.board = options.board
 
     # Communicate flags to tests.
     cls.graphics_flag = ''
     if options.no_graphics: cls.graphics_flag = '--no_graphics'
-
-    if not cls.board:
-      parser.error('Need board to convert base image to vm.')
+    if not cls.board: parser.error('Need board to convert base image to vm.')
 
   def PrepareBase(self, image_path):
     """Creates an update-able VM based on base image."""
-    # Needed for VM delta updates.  We need to use the qemu image rather
-    # than the base image on a first update.  By tracking the first_update
-    # we can set src_image to the qemu form of the base image when
-    # performing generating the delta payload.
-    self._first_update = True
-    self.vm_image_path = '%s/chromiumos_qemu_image.bin' % os.path.dirname(
-        image_path)
-    if not os.path.exists(self.vm_image_path):
-      Info('Creating %s' % vm_image_path)
-      RunCommand(['%s/image_to_vm.sh' % self.crosutils,
-                  '--full',
-                  '--from=%s' % ReinterpretPathForChroot(
-                      os.path.dirname(image_path)),
-                  '--vdisk_size=%s' % self._FULL_VDISK_SIZE,
-                  '--statefulfs_size=%s' % self._FULL_STATEFULFS_SIZE,
-                  '--board=%s' % self.board,
-                  '--test_image'], enter_chroot=True)
-
-    Info('Using %s as base' % self.vm_image_path)
-    self.assertTrue(os.path.exists(self.vm_image_path))
+    self._PrepareVMBase(image_path)
 
   def _UpdateImage(self, image_path, src_image_path='', stateful_change='old',
-                   proxy_port=''):
+                   proxy_port='', private_key_path=None):
     """Updates VM image with image_path."""
     stateful_change_flag = self.GetStatefulChangeFlag(stateful_change)
     if src_image_path and self._first_update:
       src_image_path = self.vm_image_path
       self._first_update = False
 
-    # Check image payload cache first.
-    update_id = _GenerateUpdateId(target=image_path, src=src_image_path)
-    cache_path = dev_server_cache[update_id]
-    if cache_path:
-      Info('Using cache %s' % cache_path)
-      update_url = DevServerWrapper.GetDevServerURL(proxy_port, cache_path)
-      cmd = ['%s/cros_run_vm_update' % self.crosutilsbin,
-             '--vm_image_path=%s' % self.vm_image_path,
-             '--snapshot',
-             self.graphics_flag,
-             '--persist',
-             '--kvm_pid=%s' % self._kvm_pid_file,
-             '--ssh_port=%s' % self._ssh_port,
-             stateful_change_flag,
-             '--update_url=%s' % update_url,
-             ]
-    else:
-      cmd = ['%s/cros_run_vm_update' % self.crosutilsbin,
-             '--update_image_path=%s' % image_path,
-             '--vm_image_path=%s' % self.vm_image_path,
-             '--snapshot',
-             self.graphics_flag,
-             '--persist',
-             '--kvm_pid=%s' % self._kvm_pid_file,
-             '--ssh_port=%s' % self._ssh_port,
-             stateful_change_flag,
-             '--src_image=%s' % src_image_path,
-             '--proxy_port=%s' % proxy_port
-             ]
+    cmd = ['%s/cros_run_vm_update' % self.crosutilsbin,
+           '--vm_image_path=%s' % self.vm_image_path,
+           '--snapshot',
+           self.graphics_flag,
+           '--persist',
+           '--kvm_pid=%s' % self._kvm_pid_file,
+           '--ssh_port=%s' % self._ssh_port,
+           stateful_change_flag,
+          ]
 
-    if self.verbose:
-      try:
-        RunCommand(cmd)
-      except Exception, e:
-        raise UpdateException(1, e.message)
-    else:
-      (code, stdout, stderr) = RunCommandCaptureOutput(cmd)
-      if code != 0:
-        raise UpdateException(code, stdout)
+    self.AppendUpdateFlags(cmd, image_path, src_image_path, proxy_port,
+                           private_key_path)
+    self.RunUpdateCmd(cmd)
 
   def _UpdateUsingPayload(self, update_path, stateful_change='old',
                          proxy_port=None):
     """Updates a vm image using cros_run_vm_update."""
     stateful_change_flag = self.GetStatefulChangeFlag(stateful_change)
     cmd = ['%s/cros_run_vm_update' % self.crosutilsbin,
            '--payload=%s' % update_path,
            '--vm_image_path=%s' % self.vm_image_path,
            '--snapshot',
            self.graphics_flag,
            '--persist',
            '--kvm_pid=%s' % self._kvm_pid_file,
            '--ssh_port=%s' % self._ssh_port,
            stateful_change_flag,
            ]
-
-    if proxy_port:
-      cmd.append('--proxy_port=%s' % proxy_port)
-
-    if self.verbose:
-      try:
-        RunCommand(cmd)
-      except Exception, e:
-        raise UpdateException(1, e.message)
-    else:
-      (code, stdout, stderr) = RunCommandCaptureOutput(cmd)
-      if code != 0:
-        raise UpdateException(code, stdout)
+    if proxy_port: cmd.append('--proxy_port=%s' % proxy_port)
+    self.RunUpdateCmd(cmd)
 
   def VerifyImage(self, percent_required_to_pass):
     """Runs vm smoke suite to verify image."""
     # image_to_live already verifies lsb-release matching.  This is just
     # for additional steps.
 
     commandWithArgs = ['%s/cros_run_vm_test' % self.crosutilsbin,
                        '--image_path=%s' % self.vm_image_path,
                        '--snapshot',
                        '--persist',
                        '--kvm_pid=%s' % self._kvm_pid_file,
                        '--ssh_port=%s' % self._ssh_port,
                        self.verify_suite,
                        ]
 
     if self.graphics_flag:
       commandWithArgs.append(self.graphics_flag)
 
     output = RunCommand(commandWithArgs, error_ok=True, enter_chroot=False,
                         redirect_stdout=True)
     return self.AssertEnoughTestsPassed(self, output, percent_required_to_pass)
 
 
-class GenerateVirtualAUDeltasTest(VirtualAUTest):
-  """Class the overrides VirtualAUTest and stores deltas we will generate."""
+class PregenerateAUDeltas(unittest.TestCase, AUTest):
+  """Magical class that emulates an AUTest to store deltas we will generate.
+  
+  This class emulates an AUTest such that when it runs as a TestCase it runs
+  through the exact up
+  """
   delta_list = {}
 
   def setUp(self):
     AUTest.setUp(self)
 
   def tearDown(self):
     pass
 
+  @classmethod
+  def ProcessOptions(cls, parser, options):
+    AUTest.ProcessOptions(parser, options)
+    cls.au_type = options.type
+
+  def PrepareBase(self, image_path):
+    if self.au_type == 'vm':
+      self._PrepareVMBase(image_path)
+    else:
+      self._PrepareRealBase(image_path)
+
   def _UpdateImage(self, image_path, src_image_path='', stateful_change='old',
-                   proxy_port=None):
-    if src_image_path and self._first_update:
+                   proxy_port=None, private_key_path=None):
+    if self.au_type == 'vm' and src_image_path and self._first_update:
       src_image_path = self.vm_image_path
       self._first_update = False
 
+    # Generate a value that combines delta with private key path.
+    val = '%s+%s' % (src_image_path, private_key_path)
     if not self.delta_list.has_key(image_path):
-      self.delta_list[image_path] = set([src_image_path])
+      self.delta_list[image_path] = set([val])
     else:
-      self.delta_list[image_path].add(src_image_path)
+      self.delta_list[image_path].add(val)
 
   def AttemptUpdateWithPayloadExpectedFailure(self, payload, expected_msg):
     pass
 
   def VerifyImage(self, percent_required_to_pass):
     pass
 
 
 class ParallelJob(threading.Thread):
   """Small wrapper for threading.  Thread that releases a semaphores on exit."""
@@ skipping 74 matching lines @@
   def GetDevServerURL(cls, port, sub_dir):
     """Returns the dev server url for a given port and sub directory."""
     ip_addr = GetIPAddress()
     if not port: port = 8080
     url = 'http://%(ip)s:%(port)s/%(dir)s' % {'ip': ip_addr,
                                               'port': str(port),
                                               'dir': sub_dir}
     return url
 
 
-def _GenerateUpdateId(target, src):
+def _GenerateUpdateId(target, src, key):
   """Returns a simple representation id of target and src paths."""
-  if src:
-    return '%s->%s' % (target, src)
-  else:
-    return target
+  update_id = target
+  if src: update_id = '->'.join([update_id, src])
+  if key: update_id = '+'.join([update_id, key])
+  return update_id
 
 
-def _RunParallelJobs(number_of_sumultaneous_jobs, jobs, jobs_args, print_status):
+def _RunParallelJobs(number_of_sumultaneous_jobs, jobs, jobs_args,
+                     print_status):
 
   """Runs set number of specified jobs in parallel.
 
   Args:
     number_of_simultaneous_jobs:  Max number of threads to be run in parallel.
     jobs:  Array of methods to run.
     jobs_args:  Array of args associated with method calls.
     print_status: True if you'd like this to print out .'s as it runs jobs.
   Returns:
     Returns an array of results corresponding to each thread.
@@ skipping 49 matching lines @@
   This method effectively pre-generates the dev server cache for all tests.
 
   Args:
     parser: parser from main.
     options: options from parsed parser.
   Returns:
     Dictionary of Update Identifiers->Relative cache locations.
   Raises:
     UpdateException if we fail to generate an update.
   """
-  def _GenerateVMUpdate(target, src):
+  def _GenerateVMUpdate(target, src, private_key_path):
     """Generates an update using the devserver."""
-    target = ReinterpretPathForChroot(target)
-    if src:
-      src = ReinterpretPathForChroot(src)
+    command = ['./enter_chroot.sh',
+               '--nogit_config',
+               '--',
+               'sudo',
+               './start_devserver',
+               '--pregenerate_update',
+               '--exit',
+              ]
+    # Add actual args to command.
+    command.append('--image=%s' % ReinterpretPathForChroot(target))
+    if src: command.append('--src_image=%s' % ReinterpretPathForChroot(src))
+    if options.type == 'vm': command.append('--for_vm')
+    if private_key_path:
+      command.append('--private_key=%s' %
+                     ReinterpretPathForChroot(private_key_path))
 
-    return RunCommandCaptureOutput(['./enter_chroot.sh',
-                                    '--nogit_config',
-                                    '--',
-                                    'sudo',
-                                    './start_devserver',
-                                    '--pregenerate_update',
-                                    '--exit',
-                                    '--image=%s' % target,
-                                    '--src_image=%s' % src,
-                                    '--for_vm',
-                                   ], combine_stdout_stderr=True,
-                                   print_cmd=False)
+    return RunCommandCaptureOutput(command, combine_stdout_stderr=True,
+                                   print_cmd=True)
 
   # Get the list of deltas by mocking out update method in test class.
-  test_suite = _PrepareTestSuite(parser, options, GenerateVirtualAUDeltasTest)
+  test_suite = _PrepareTestSuite(parser, options, PregenerateAUDeltas)
   test_result = unittest.TextTestRunner(verbosity=0).run(test_suite)
+  if not test_result.wasSuccessful():
+    raise UpdateException(1, 'Error finding updates to generate.')
 
   Info('The following delta updates are required.')
   update_ids = []
   jobs = []
   args = []
-  for target, srcs in GenerateVirtualAUDeltasTest.delta_list.items():
-    for src in srcs:
-      update_id = _GenerateUpdateId(target=target, src=src)
+  for target, srcs in PregenerateAUDeltas.delta_list.items():
+    for src_key in srcs:
+      (src, key) = src_key.split('+')
+      # TODO(sosa): Add private key as part of caching name once devserver can
+      # handle it its own cache.
+      update_id = _GenerateUpdateId(target=target, src=src, key=key)
       print >> sys.stderr, 'AU: %s' % update_id
       update_ids.append(update_id)
       jobs.append(_GenerateVMUpdate)
-      args.append((target, src))
+      args.append((target, src, key))
 
   raw_results = _RunParallelJobs(options.jobs, jobs, args, print_status=True)
   results = []
 
   # Looking for this line in the output.
   key_line_re = re.compile('^PREGENERATED_UPDATE=([\w/.]+)')
   for result in raw_results:
     (return_code, output, _) = result
     if return_code != 0:
       Warning(output)
@@ skipping 31 matching lines @@
     test_case = unittest.TestLoader().loadTestsFromName(test_name)
     threads.append(unittest.TextTestRunner().run)
     args.append(test_case)
 
   results = _RunParallelJobs(options.jobs, threads, args, print_status=False)
   for test_result in results:
     if not test_result.wasSuccessful():
       Die('Test harness was not successful')
 
 
+def _InsertPrivateKeyIntoImage(image_path, private_key_path):
+    from_dir = os.path.dirname(image_path)
+    image = os.path.basename(image_path)
+    crosutils_dir = os.path.abspath(__file__).rsplit('/', 2)[0]
+    target_key_path = 'usr/share/update_engine/update-payload-key.pub.pem'
+
+    # Temporary directories for this function.
+    rootfs_dir = tempfile.mkdtemp(suffix='rootfs', prefix='tmp')
+    stateful_dir = tempfile.mkdtemp(suffix='stateful', prefix='tmp')
+
+    Info('Copying %s into %s' % (private_key_path, image_path))
+    try:
+      RunCommand(['%s/mount_gpt_image.sh' % crosutils_dir,
+                  '--from=%s' % from_dir,
+                  '--image=%s' % image,
+                  '--rootfs_mountpt=%s' % rootfs_dir,
+                  '--stateful_mountpt=%s' % stateful_dir,
+                 ], print_cmd=False, redirect_stdout=True,
+                 redirect_stderr=True)
+      path = os.path.join(rootfs_dir, target_key_path)
+      dir_path = os.path.dirname(path)
+      RunCommand(['sudo', 'mkdir', '--parents', dir_path], print_cmd=False)
+      RunCommand(['sudo', 'cp', '--force', '-p', private_key_path, path],
+                 print_cmd=False)
+    finally:
+      # Unmount best effort regardless.
+      RunCommand(['%s/mount_gpt_image.sh' % crosutils_dir,
+                  '--unmount',
+                  '--rootfs_mountpt=%s' % rootfs_dir,
+                  '--stateful_mountpt=%s' % stateful_dir,
+                 ], print_cmd=False, redirect_stdout=True, redirect_stderr=True)
+      # Clean up our directories.
+      os.rmdir(rootfs_dir)
+      os.rmdir(stateful_dir)
+
+
 def main():
   parser = optparse.OptionParser()
   parser.add_option('-b', '--base_image',
                     help='path to the base image.')
   parser.add_option('-r', '--board',
                     help='board for the images.')
   parser.add_option('--no_delta', action='store_false', default=True,
                     dest='delta',
                     help='Disable using delta updates.')
   parser.add_option('--no_graphics', action='store_true',
                     help='Disable graphics for the vm test.')
   parser.add_option('-j', '--jobs', default=8, type=int,
                      help='Number of simultaneous jobs')
+  parser.add_option('--public_key', default=None,
+                     help='Public key to use on images and updates.')
+  parser.add_option('--private_key', default=None,
+                     help='Private key to use on images and updates.')
   parser.add_option('-q', '--quick_test', default=False, action='store_true',
                     help='Use a basic test to verify image.')
   parser.add_option('-m', '--remote',
                     help='Remote address for real test.')
   parser.add_option('-t', '--target_image',
                     help='path to the target image.')
   parser.add_option('--test_prefix', default='test',
                     help='Only runs tests with specific prefix i.e. '
                          'testFullUpdateWipeStateful.')
   parser.add_option('-p', '--type', default='vm',
                     help='type of test to run: [vm, real]. Default: vm.')
   parser.add_option('--verbose', default=True, action='store_true',
                     help='Print out rather than capture output as much as '
                          'possible.')
   (options, leftover_args) = parser.parse_args()
 
   if leftover_args:
     parser.error('Found extra options we do not support: %s' % leftover_args)
 
+  assert options.target_image and os.path.exists(options.target_image), \
+    'Target image path does not exist'
+  if not options.base_image:
+    Warning('Base image not specified.  Using target image as base image.')
+    options.base_image = options.target_image
+
   # Figure out the test_class.
   if options.type == 'vm':  test_class = VirtualAUTest
   elif options.type == 'real': test_class = RealAUTest
   else: parser.error('Could not parse harness type %s.' % options.type)
 
+  if options.private_key or options.public_key:
+    assert os.path.exists(options.private_key), 'Could not find private key.'
+    assert os.path.exists(options.public_key), 'Could not find public key.'
+    _InsertPrivateKeyIntoImage(options.target_image, options.public_key)
+    _InsertPrivateKeyIntoImage(options.base_image, options.public_key)
+
   # TODO(sosa): Caching doesn't really make sense on non-vm images (yet).
   global dev_server_cache
   if options.type == 'vm' and options.jobs > 1:
     dev_server_cache = _PregenerateUpdates(parser, options)
     my_server = DevServerWrapper()
     my_server.start()
     try:
       _RunTestsInParallel(parser, options, test_class)
     finally:
       my_server.Stop()
 
   else:
     dev_server_cache = None
     test_suite = _PrepareTestSuite(parser, options, test_class)
     test_result = unittest.TextTestRunner(verbosity=2).run(test_suite)
     if not test_result.wasSuccessful():
       Die('Test harness was not successful.')
 
 
 if __name__ == '__main__':
   main()