add test to verify Cr50 update and recovery from erased nvmem

This test can be used to verify the cr50-update process, or recovery
from erased nvmem.

This test will update to the given dev image. After that, each iteration
of the test will start by rolling back to the oldest release image. It
will then use the cr50-update script to update to each image in order of
lowest to highest version. After all of the iterations are complete, the
test will reflash the Cr50 image that the dut had at the start of the
test.

This test verifies that the device updated and the update script did not
exit with any unexpected exit codes.

If erase_nvmem is being tested then nvmem will also be erased during
rollback.

There are four parameters in each control file: test_type, dev_image
release_image, and old_release_image. The image parameters are the
locations of the test images. If the test type is erase_nvmem then cr50
will erase nvmem before rolling back to the release image. The
old_release_image needs to have the oldest version if it is being used.
dev_image needs to have a version higher than release_image,
old_release_image, and whatever image is running on the DUT. If this is
not the case, the test cannot guarantee the original state of the dut is
restored.

There are two control files: control and control.erase_nvmem.
control.erase_nvmem will ignore 'old_release_image' and make sure the
test_type is set to 'erase_nvmem'.

CQ-DEPEND=CL:456624
BUG=b:35833679
BUG=b:35833781
BRANCH=none
TEST=test_that $DUT_IP firmware_Cr50Update
--args="iterations=3 old_release_image=$OLD_RELEASE
release_image=$RELEASE_IMAGE dev_image=$DEV_IMAGE"

Change-Id: I4e2f6c3a720f873e44caa6426f251fe0119cba30
Signed-off-by: Mary Ruthven <mruthven@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/456523
Reviewed-by: Wai-Hong Tam <waihong@google.com>
diff --git a/client/common_lib/cros/cr50_utils.py b/client/common_lib/cros/cr50_utils.py
index da6e902..00b1534 100644
--- a/client/common_lib/cros/cr50_utils.py
+++ b/client/common_lib/cros/cr50_utils.py
@@ -16,6 +16,7 @@
 GET_CR50_VERSION = 'cat /var/cache/cr50-version'
 GET_CR50_MESSAGES ='grep "cr50-.*\[" /var/log/messages'
 UPDATE_FAILURE = 'unexpected cr50-update exit code'
+DUMMY_VER = '-1.-1.-1'
 # This dictionary is used to search the usb_updater output for the version
 # strings. There are two usb_updater commands that will return versions:
 # 'fwver' and 'binver'.
@@ -33,6 +34,35 @@
 }
 
 
+def AssertVersionsAreEqual(name_a, ver_a, name_b, ver_b):
+    """Raise an error ver_a isn't the same as ver_b
+
+    Args:
+        name_a: the name of section a
+        ver_a: the version string for section a
+        name_b: the name of section b
+        ver_b: the version string for section b
+
+    Raises:
+        AssertionError if ver_a is not equal to ver_b
+    """
+    assert ver_a == ver_b, ("Versions do not match: %s %s %s %s" %
+                            (name_a, ver_a, name_b, ver_b))
+
+
+def GetNewestVersion(ver_a, ver_b):
+    """Compare the versions. Return the newest one. If they are the same return
+    None."""
+    a = [int(x) for x in ver_a.split('.')]
+    b = [int(x) for x in ver_b.split('.')]
+
+    if a > b:
+        return ver_a
+    if b > a:
+        return ver_b
+    return None
+
+
 def GetVersion(versions, name):
     """Return the version string from the dictionary.
 
@@ -40,17 +70,25 @@
     substring name. Make sure all of the versions match and return the version
     string. Raise an error if the versions don't match.
 
-    @param version: dictionary with the partition names as keys and the
-                    partition version strings as values.
-    @param name: the string used to find the relevant items in versions.
+    Args:
+        version: dictionary with the partition names as keys and the
+                 partition version strings as values.
+        name: the string used to find the relevant items in versions.
+    Returns:
+        the version from versions or "-1.-1.-1" if an invalid RO was detected.
     """
     ver = None
+    key = None
     for k, v in versions.iteritems():
         if name in k:
-            if ver and ver != v:
-                raise error.TestFail("Versions don't match %s %s" % (ver, v))
+            if v == DUMMY_VER:
+                logging.info("Detected invalid %s %s", name, v)
+                return v
+            elif ver:
+                AssertVersionsAreEqual(key, ver, k, v)
             else:
                 ver = v
+                key = k
     return ver
 
 
@@ -91,12 +129,8 @@
     return GetVersionFromUpdater(client, ["--binver", image])
 
 
-def CompareVersions(name_a, ver_a, rw_a, ver_b):
-    """Compare ver_a to ver_b. Raise an error if they aren't the same"""
-    if ver_a != ver_b:
-        raise error.TestFail("Versions do not match: %s RO %s RW %s %s RO %s " \
-                             "RW %s" % (name_a, ver_a[0], ver_a[1], name_b,
-                                        ver_b[0], ver_b[1]))
+def GetVersionString(ver):
+    return 'RO %s RW %s' % (ver[0], ver[1])
 
 
 def GetRunningVersion(client):
@@ -115,7 +149,8 @@
     running_ver = GetFwVersion(client)
     saved_ver = GetSavedVersion(client)
 
-    CompareVersions("Running", running_ver, "Saved", saved_ver)
+    AssertVersionsAreEqual("Running", GetVersionString(running_ver),
+                           "Saved", GetVersionString(saved_ver))
     return running_ver
 
 
@@ -145,7 +180,7 @@
     return messages.rsplit('\n', 1)[-1]
 
 
-def VerifyUpdate(client, ver=None, last_message=''):
+def VerifyUpdate(client, ver='', last_message=''):
     """Verify that the saved update state is correct and there were no
     unexpected cr50-update exit codes since the last update.
 
@@ -158,8 +193,10 @@
     logging.debug("last cr50 message %s", last_message)
 
     new_ver = GetRunningVersion(client)
-    if ver:
-        CompareVersions("Old", ver, "Updated", new_ver)
+    if ver != '':
+        if DUMMY_VER != ver[0]:
+            AssertVersionsAreEqual("Old RO", ver[0], "Updated RO", new_ver[0])
+        AssertVersionsAreEqual("Old RW", ver[1], "Updated RW", new_ver[1])
     return new_ver, last_message
 
 
@@ -167,3 +204,20 @@
     """Removes the cr50 status files in /var/cache and reboots the AP"""
     client.run("rm %s" % CR50_STATE)
     client.reboot()
+
+
+def InstallImage(client, src, dest=CR50_FILE):
+    """Copy the image at src to dest on the dut
+    Args:
+        src: the image location of the server
+        dest: the desired location on the dut
+    Returns:
+        The filename where the image was copied to on the dut, a tuple
+        containing the RO and RW version of the file
+    """
+    # Send the file to the DUT
+    client.send_file(src, dest)
+
+    ver = GetBinVersion(client, dest)
+    client.run("sync")
+    return dest, ver
diff --git a/server/cros/servo/chrome_ec.py b/server/cros/servo/chrome_ec.py
index a7c812b..18f9fe6 100644
--- a/server/cros/servo/chrome_ec.py
+++ b/server/cros/servo/chrome_ec.py
@@ -4,6 +4,7 @@
 
 import ast, logging, re, time
 
+from autotest_lib.client.bin import utils
 from autotest_lib.client.common_lib import error
 from autotest_lib.client.cros import ec
 
@@ -285,6 +286,10 @@
     This class is to abstract these interfaces.
     """
     IDLE_COUNT = 'count: (\d+)'
+    VERSION_FORMAT = '\d+\.\d+\.\d+'
+    VERSION_ERROR = 'Error'
+    INACTIVE = '\nRW_(A|B): +(%s|%s)(/DBG|)?' % (VERSION_FORMAT, VERSION_ERROR)
+    ACTIVE = '\nRW_(A|B): +\* +(%s)(/DBG|)?' % (VERSION_FORMAT)
 
     def __init__(self, servo):
         super(ChromeCr50, self).__init__(servo, "cr50_console")
@@ -303,16 +308,111 @@
             raise error.TestFail("Could not clear deep sleep count")
 
 
+    def has_command(self, cmd):
+        """Returns 1 if cr50 has the command 0 if it doesn't"""
+        try:
+            self.send_command_get_output('help', [cmd])
+        except:
+            logging.info("Image does not include '%s' command", cmd)
+            return 0
+        return 1
+
+
+    def erase_nvmem(self):
+        """Use flasherase to erase both nvmem sections"""
+        if not self.has_command('flasherase'):
+            raise error.TestError("need image with 'flasherase'")
+
+        self.send_command('flasherase 0x7d000 0x3000')
+        self.send_command('flasherase 0x3d000 0x3000')
+
+
+    def reboot(self):
+        """Reboot Cr50 and wait for CCD to be enabled"""
+        self.send_command('reboot')
+        self.wait_for_ccd_disable()
+        self.ccd_enable()
+
+
+    def rollback(self):
+        """Set the reset counter high enough to force a rollback then reboot"""
+        if not self.has_command('rw') or not self.has_command('eraseflashinfo'):
+            raise error.TestError("need image with 'rw' and 'eraseflashinfo'")
+
+        if self.get_inactive_version_info()[1] == self.VERSION_ERROR:
+            raise error.TestError("Invalid image in inactive RW")
+
+        # Increase the reset count to above the rollback threshold
+        self.send_command('rw 0x40000128 1')
+        self.send_command('rw 0x4000012c 15')
+
+        self.send_command('eraseflashinfo')
+
+        self.reboot()
+
+
+    def get_version_info(self, regexp):
+        """Get information from the version command"""
+        return self.send_command_get_output('ver', [regexp])[0][1::]
+
+
+    def get_inactive_version_info(self):
+        """Get the active partition, version, and hash"""
+        return self.get_version_info(self.INACTIVE)
+
+
+    def get_active_version_info(self):
+        """Get the active partition, version, and hash"""
+        return self.get_version_info(self.ACTIVE)
+
+
+    def get_ccd_state(self):
+        """Get the CCD state from servo
+
+        Returns:
+            'off' or 'on' based on whether the cr50 console is working.
+        """
+        return self._servo.get('ccd_state')
+
+
+    def wait_for_ccd_state(self, state, timeout):
+        """Wait up to timeout seconds for CCD to be 'on' or 'off'
+        Args:
+            state: a string either 'on' or 'off'.
+            timeout: time in seconds to wait
+
+        Raises
+            TestFail if ccd never reaches the specified state
+        """
+        logging.info("Wait until ccd is '%s'", state)
+        value = utils.wait_for_value(self.get_ccd_state, state,
+                                     timeout_sec=timeout)
+        if value != state:
+            raise error.TestFail("timed out before detecting ccd '%s'" % state)
+        logging.info("ccd is '%s'", state)
+
+
+    def wait_for_ccd_disable(self, timeout=60):
+        """Wait for the cr50 console to stop working"""
+        self.wait_for_ccd_state('off', timeout)
+
+
+    def wait_for_ccd_enable(self, timeout=60):
+        """Wait for the cr50 console to start working"""
+        self.wait_for_ccd_state('on', timeout)
+
+
     def ccd_disable(self):
         """Change the values of the CC lines to disable CCD"""
+        logging.info("disable ccd")
         self._servo.set_nocheck('servo_v4_ccd_mode', 'disconnect')
-        # TODO: Add a better way to wait until usb is disconnected
-        time.sleep(3)
+        self.wait_for_ccd_disable()
 
 
     def ccd_enable(self):
         """Reenable CCD and reset servo interfaces"""
+        logging.info("reenable ccd")
         self._servo.set_nocheck('servo_v4_ccd_mode', 'ccd')
         self._servo.set('sbu_mux_enable', 'on')
         self._servo.set_nocheck('power_state', 'ccd_reset')
-        time.sleep(2)
+        self.wait_for_ccd_enable()
diff --git a/server/site_tests/firmware_Cr50Update/control b/server/site_tests/firmware_Cr50Update/control
new file mode 100644
index 0000000..3c6fbf0
--- /dev/null
+++ b/server/site_tests/firmware_Cr50Update/control
@@ -0,0 +1,62 @@
+# Copyright 2017 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.
+
+AUTHOR = "mruthven"
+NAME = "firmware_Cr50Update"
+PURPOSE = "Verify Cr50 update"
+TIME = "SHORT"
+TEST_TYPE = "server"
+
+DOC = """
+This test verifies Cr50 update works or recovery from erased nvmem.
+
+To test nvmem recovery set test to "nvmem_recovery" or use
+firmware_Cr50Update.nvmem_recovery
+
+The test will rollback to the oldest Cr50 image and then verify each update to
+the next newest image runs successfully. If testing nvmem recovery, nvmem will
+be erased during the rollback from dev_image to release_image and
+old_release_image will be ignored for nvmem_recovery tests.
+
+old_release_image should have a lower version than release_image and
+release_image should have a version lower than the dev_image. The dev_image
+needs to have a higher version than all of the given images and whatever is
+running on Cr50 to guarantee that the state can be restored.
+
+After the test is complete the original Cr50 image will be reflashed onto the
+device.
+
+@param iterations: the number of iterations to run
+@param dev_image: the location of the dev image. Must be built with CR50_DEV=1
+@param release_image: the location of the release image
+@param old_release_image: the location of the old release image. This is
+                          optional. If it is included, the test will verify that
+                          the old release can update to the new release, then
+                          the new release can update to the dev image.
+@param test: string representing the test type. use "nvmem_recovery" if nvmem
+             should be erased before updating to the release image. This can be
+             used to verify that Cr50 can recovery from erased nvmem.
+"""
+
+from autotest_lib.client.common_lib import error
+from autotest_lib.server import utils
+
+args_dict = utils.args_to_dict(args)
+servo_args = hosts.CrosHost.get_servo_arguments(args_dict)
+
+iterations = int(args_dict.get("iterations", 1))
+old_release_image = args_dict.get("old_release_image", "")
+release_image = args_dict.get("release_image", "")
+dev_image = args_dict.get("dev_image", "")
+test = args_dict.get("test", "")
+
+def run(machine):
+    host = hosts.create_host(machine, servo_args=servo_args)
+
+    job.run_test("firmware_Cr50Update", host=host, cmdline_args=args,
+                 release_image=release_image, dev_image=dev_image,
+                 old_release_image=old_release_image, test=test,
+                 iterations=iterations)
+
+parallel_simple(run, machines)
diff --git a/server/site_tests/firmware_Cr50Update/control.erase_nvmem b/server/site_tests/firmware_Cr50Update/control.erase_nvmem
new file mode 100644
index 0000000..c0d58dd
--- /dev/null
+++ b/server/site_tests/firmware_Cr50Update/control.erase_nvmem
@@ -0,0 +1,49 @@
+# Copyright 2017 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.
+
+AUTHOR = "mruthven"
+NAME = "firmware_Cr50Update.erase_nvmem"
+PURPOSE = "Verify Cr50 update"
+TIME = "SHORT"
+TEST_TYPE = "server"
+
+DOC = """
+This test verifies Cr50 can recover from erased nvmem.
+
+The test will update to the dev image and then rollback to the release image.
+During the rollback Cr50 will erase nvmem. The test verifies the device boots
+into normal mode and not recovery after nvmem has been erased.
+
+The release_image needs to have a version lower than the dev_image. The
+dev_image needs to have a higher version than the release image and whatever
+image is running on Cr50 at the start of the test to guarantee that the
+original state can be restored.
+
+After the test is complete the original Cr50 image will be reflashed onto the
+device.
+
+@param iterations: the number of iterations to run
+@param dev_image: the location of the dev image. Must be built with CR50_DEV=1
+@param release_image: the location of the release image
+"""
+
+from autotest_lib.client.common_lib import error
+from autotest_lib.server import utils
+
+args_dict = utils.args_to_dict(args)
+servo_args = hosts.CrosHost.get_servo_arguments(args_dict)
+
+iterations = int(args_dict.get("iterations", 1))
+release_image = args_dict.get("release_image", "")
+dev_image = args_dict.get("dev_image", "")
+test = "erase_nvmem"
+
+def run(machine):
+    host = hosts.create_host(machine, servo_args=servo_args)
+
+    job.run_test("firmware_Cr50Update", host=host, cmdline_args=args,
+                 release_image=release_image, dev_image=dev_image,
+                 test=test, iterations=iterations)
+
+parallel_simple(run, machines)
diff --git a/server/site_tests/firmware_Cr50Update/firmware_Cr50Update.py b/server/site_tests/firmware_Cr50Update/firmware_Cr50Update.py
new file mode 100644
index 0000000..1eec894a
--- /dev/null
+++ b/server/site_tests/firmware_Cr50Update/firmware_Cr50Update.py
@@ -0,0 +1,389 @@
+# Copyright 2017 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.
+
+import logging
+import os
+import time
+
+from autotest_lib.client.common_lib import error
+from autotest_lib.client.common_lib.cros import cr50_utils
+from autotest_lib.server import autotest, test
+from autotest_lib.server.cros import debugd_dev_tools
+from autotest_lib.server.cros.faft.firmware_test import FirmwareTest
+
+
+class firmware_Cr50Update(FirmwareTest):
+    """
+    Verify a dut can update to the given image.
+
+    Copy the new image onto the device and clear the update state to force
+    cr50-update to run. The test will fail if Cr50 does not update or if the
+    update script encounters any errors.
+
+    @param image: the location of the update image
+    @param image_type: string representing the image type. If it is "dev" then
+                       don't check the RO versions when comparing versions.
+    """
+    version = 1
+    UPDATE_TIMEOUT = 20
+    ERASE_NVMEM = "erase_nvmem"
+
+    DEV_NAME = "dev_image"
+    OLD_RELEASE_NAME = "old_release_image"
+    RELEASE_NAME = "release_image"
+    ORIGINAL_NAME = "original_image"
+    RESTORE_ORIGINAL_TRIES = 3
+    SUCCESS = 0
+    UPSTART_SUCCESS = 1
+
+
+    def initialize(self, host, cmdline_args, release_image, dev_image,
+                   old_release_image="", test=""):
+        """Initialize servo and process the given images"""
+        self.processed_images = False
+
+        super(firmware_Cr50Update, self).initialize(host, cmdline_args)
+        if not hasattr(self, "cr50"):
+            raise error.TestNAError('Test can only be run on devices with '
+                                    'access to the Cr50 console')
+        # Make sure ccd is disabled so it won't interfere with the update
+        self.cr50.ccd_disable()
+
+        self.rootfs_tool = debugd_dev_tools.RootfsVerificationTool()
+        self.rootfs_tool.initialize(host)
+        if not self.rootfs_tool.is_enabled():
+            logging.debug('Removing rootfs verification.')
+            # 'enable' actually disables rootfs verification
+            self.rootfs_tool.enable()
+
+        self.host = host
+        self.test = test.lower()
+
+        # A dict used to store relevant information for each image
+        self.images = {}
+
+        running_rw = cr50_utils.GetRunningVersion(self.host)[1]
+        # Get the original image from the cr50 firmware directory on the dut
+        self.save_original_image(cr50_utils.CR50_FILE)
+
+        # If Cr50 is not running the image from the cr50 firmware directory,
+        # then raise an error, otherwise the test will not be able to restore
+        # the original state during cleanup.
+        if running_rw != self.original_rw:
+            raise error.TestError("Can't determine original Cr50 version. "
+                                  "Running %s, but saved %s." %
+                                  (running_rw, self.original_rw))
+
+        # Process the given images in order of oldest to newest. Get the version
+        # info and add them to the update order
+        self.update_order = []
+        if self.test != self.ERASE_NVMEM and old_release_image:
+            self.add_image_to_update_order(self.OLD_RELEASE_NAME,
+                                           old_release_image)
+        self.add_image_to_update_order(self.RELEASE_NAME, release_image)
+        self.add_image_to_update_order(self.DEV_NAME, dev_image)
+        self.verify_update_order()
+        self.processed_images = True
+        logging.info("Update %s", self.update_order)
+
+        # Update to the dev image
+        self.run_update(self.DEV_NAME)
+
+
+    def restore_original_image(self):
+        """Update to the image that was running at the start of the test.
+
+        Returns SUCCESS if the update was successful or the update error if it
+        failed.
+        """
+        rv = self.SUCCESS
+        _, running_rw, is_dev = self.cr50.get_active_version_info()
+        new_rw = cr50_utils.GetNewestVersion(running_rw, self.original_rw)
+
+        # If Cr50 is running the original image, then no update is needed.
+        if new_rw is None:
+            return rv
+
+        try:
+            # If a rollback is needed, update to the dev image so it can
+            # rollback to the original image.
+            if new_rw != self.original_rw and not is_dev:
+                logging.info("Updating to dev image to enable rollback")
+                self.run_update(self.DEV_NAME, use_usb_update=True)
+
+            logging.info("Updating to the original image %s",
+                         self.original_rw)
+            self.run_update(self.ORIGINAL_NAME, use_usb_update=True)
+        except Exception, e:
+            logging.info("cleanup update from %s to %s failed", running_rw,
+                          self.original_rw)
+            logging.debug(e)
+            rv = e
+        self.cr50.ccd_enable()
+        return rv
+
+
+    def cleanup(self):
+        """Update Cr50 to the image it was running at the start of the test"""
+        logging.warning('rootfs verification is disabled')
+
+        # Make sure keepalive is disabled
+        self.cr50.ccd_enable()
+        self.cr50.send_command("ccd keepalive disable")
+
+        # Restore the original Cr50 image
+        if self.processed_images:
+            for i in xrange(self.RESTORE_ORIGINAL_TRIES):
+                rv = self.restore_original_image()
+                if rv == self.SUCCESS:
+                    logging.info("Successfully restored the original image")
+                    break
+            if rv != self.SUCCESS:
+                logging.info("Could not restore the original image")
+                raise rv
+
+        super(firmware_Cr50Update, self).cleanup()
+
+
+    def run_usb_update(self, dest, is_newer):
+        """Run usb_update with the given image.
+
+        If the new image version is newer than the one Cr50 is running, then
+        the upstart option will be used. This will reboot the AP after
+        usb_updater runs, so Cr50 will finish the update and jump to the new
+        image.
+
+        If the new image version is older than the one Cr50 is running, then the
+        best usb_updater can do is flash the image into the inactive partition.
+        To finish th update, the 'rw' command will need to be used to force a
+        rollback.
+
+        @param dest: the image location.
+        @param is_newer: True if the rw version of the update image is newer
+                         than the one Cr50 is running.
+        """
+
+        result = self.host.run("status trunksd")
+        if 'running' in result.stdout:
+            self.host.run("stop trunksd")
+
+        # Enable CCD, so we can detect the Cr50 reboot.
+        self.cr50.ccd_enable()
+        if is_newer:
+            # Using -u usb_updater will post a reboot request but not reboot
+            # immediately.
+            result = self.host.run("usb_updater -s -u %s" % dest,
+                                   ignore_status=True)
+            # After a posted reboot, the usb_update exit code should equal 1.
+            if result.exit_status != self.UPSTART_SUCCESS:
+                logging.debug(result)
+                raise error.TestError("Got unexpected usb_update exit code")
+            # Reset the AP to finish the Cr50 update.
+            self.cr50.send_command("sysrst pulse")
+        else:
+            logging.info("Flashing image into inactive partition")
+            # The image at 'dest' is older than the one Cr50 is running, so
+            # upstart cannot be used. Without -u Cr50 will flash the image into
+            # the inactive partition and reboot immediately.
+            result = self.host.run("usb_updater -s %s" % dest,
+                                   ignore_timeout=True,
+                                   timeout=self.UPDATE_TIMEOUT)
+            logging.info(result)
+
+        # After usb_updater finishes running, Cr50 will reboot. Wait until Cr50
+        # reboots before continuing. Cr50 reboot can be detected by detecting
+        # when CCD stops working.
+        self.cr50.wait_for_ccd_disable()
+
+
+    def finish_rollback(self, image_rw, erase_nvmem):
+        """Rollback to the image in the inactive partition.
+
+        Use the cr50 'rw' command to set the reset counter high enough to
+        trigger a rollback, erase nvmem if requested, and then reboot cr50 to
+        finish the rollback.
+
+        @param image_rw: the rw version of the update image.
+        @param erase_nvmem: True if nvmem needs to be erased during the
+                            rollback.
+        """
+        self.cr50.ccd_enable()
+        inactive_info = self.cr50.get_inactive_version_info()
+        if inactive_info[1] != image_rw:
+            raise error.TestError("Image is not in inactive partition")
+
+        logging.info("Attempting Rollback")
+
+        # Enable CCD keepalive before turning off the DUT. It will be removed
+        # after the Cr50 reboot
+        self.cr50.send_command("ccd keepalive enable")
+
+        # Shutdown the device so it won't react to any of the following commands
+        self.ec.reboot("ap-off")
+
+        # CCD may disapppear after resetting the EC. If it does, re-enable it.
+        # TODO: remove this when CCD is no longer disabled after ec reset.
+        try:
+            self.cr50.wait_for_ccd_disable(timeout=15)
+        except error.TestFail, e:
+            pass
+        self.cr50.ccd_enable()
+
+        if erase_nvmem:
+            logging.info("Erasing nvmem")
+            self.cr50.erase_nvmem()
+
+        self.cr50.rollback()
+
+        # Verify the inactive partition is the active one after the rollback.
+        if inactive_info != self.cr50.get_active_version_info():
+            raise error.TestError("Rollback failed")
+
+        # Verify the system boots normally after erasing nvmem
+        self.check_state((self.checkers.crossystem_checker,
+                          {'mainfw_type': 'normal'}))
+
+
+    def run_update(self, image_name, erase_nvmem=False, use_usb_update=False):
+        """Copy the image to the DUT and upate to it.
+
+        Normal updates will use the cr50-update script to update. If a rollback
+        is True, use usb_update to flash the image and then use the 'rw'
+        commands to force a rollback. On rollback updates erase_nvmem can be
+        used to request that nvmem be erased during the rollback.
+
+        @param image_name: the key in the images dict that can be used to
+                           retrieve the image info.
+        @param erase_nvmem: True if nvmem needs to be erased during the
+                            rollback.
+        @param use_usb_update: True if usb_updater should be used directly
+                               instead of running the update script.
+        """
+        self.cr50.ccd_disable()
+        # Get the current update information
+        image_ver, image_ver_str, image_path = self.images[image_name]
+
+        dest, ver = cr50_utils.InstallImage(self.host, image_path)
+        assert ver == image_ver, "Install failed"
+        image_rw = image_ver[1]
+
+        running_ver = cr50_utils.GetRunningVersion(self.host)
+        running_ver_str = cr50_utils.GetVersionString(running_ver)
+
+        # If the given image is older than the running one, then we will need
+        # to do a rollback to complete the update.
+        rollback = (cr50_utils.GetNewestVersion(running_ver[1], image_rw) !=
+                    image_rw)
+        logging.info("Attempting %s from %s to %s",
+                     "rollback" if rollback else "update", running_ver_str,
+                     image_ver_str)
+
+        # If a rollback is needed, flash the image into the inactive partition,
+        # on or use usb_update to update to the new image if it is requested.
+        if use_usb_update or rollback:
+            self.run_usb_update(dest, not rollback)
+        # Use cr50 console commands to rollback to the old image.
+        if rollback:
+            self.finish_rollback(image_rw, erase_nvmem)
+        # Running the usb update or rollback will enable ccd. Disable it again.
+        self.cr50.ccd_disable()
+
+        # Get the last cr50 update related message from /var/log/messages
+        last_message = cr50_utils.CheckForFailures(self.host, '')
+
+        # Clear the update state and reboot, so cr50-update will run again.
+        cr50_utils.ClearUpdateStateAndReboot(self.host)
+
+        # Verify the version has been updated and that there have been no
+        # unexpected usb_updater exit codes.
+        cr50_utils.VerifyUpdate(self.host, image_ver, last_message)
+
+        logging.info("Successfully updated from %s to %s %s", running_ver_str,
+                     image_name, image_ver_str)
+
+
+    def add_image_to_update_order(self, image_name, image_path):
+        """Process the image. Add it to the update_order list and images dict.
+
+        Copy the image to the DUT and get version information.
+
+        Store the image information in the images dictionary and add it to the
+        update_order list.
+
+        Args:
+            image_name: string that is what the image should be called. Used as
+                        the key in the images dict.
+            image_path: the path for the image.
+
+        Raises:
+            TestError if the image could not be found.
+        """
+        tmp_file = '/tmp/%s.bin' % image_name
+
+        if not os.path.isfile(image_path):
+            raise error.TestError("Failed to locate %s" % image_name)
+
+        _, ver = cr50_utils.InstallImage(self.host, image_path, tmp_file)
+        ver_str = cr50_utils.GetVersionString(ver)
+
+        self.update_order.append([image_name, ver[1]])
+        self.images[image_name] = (ver, ver_str, image_path)
+        logging.info("%s stored at %s with version %s", image_name, image_path,
+                     ver_str)
+
+
+    def verify_update_order(self):
+        """Verify each image in the update order has a higher version than the
+        last.
+
+        The test uses the cr50 update script to update to the next image in the
+        update order. If the versions are not in ascending order then the update
+        won't work. Cr50 cannot update to an older version using the standard
+        update process.
+
+        Raises:
+            TestError if the versions are not in ascending order.
+        """
+        for i, update_info in enumerate(self.update_order[1::]):
+            last_name, last_rw = self.update_order[i]
+            name, rw = update_info
+            if cr50_utils.GetNewestVersion(last_rw, rw) != rw:
+                raise error.TestError("%s is version %s. %s needs to have a "
+                                      "higher version, but it has %s" %
+                                      (last_name, last_rw, name, rw))
+
+
+    def save_original_image(self, dut_path):
+        """Save the image currently running on the DUT.
+
+        Copy the image from the DUT to the local tmp directory and get version
+        information. Store the information in the images dict.
+
+        @param dut_path: the location of the cr50 prod image on the DUT.
+        """
+        name = self.ORIGINAL_NAME
+        local_dest = '/tmp/%s.bin' % name
+
+        self.host.get_file(dut_path, local_dest)
+
+        ver = cr50_utils.GetBinVersion(self.host, dut_path)
+        ver_str = cr50_utils.GetVersionString(ver)
+
+        self.images[name] = (ver, ver_str, local_dest)
+        logging.info("%s stored at %s with version %s", name, local_dest,
+                     ver_str)
+
+        self.original_rw = ver[1]
+
+
+    def after_run_once(self):
+        """Add log printing what iteration we just completed"""
+        logging.info("Update iteration %s ran successfully", self.iteration)
+
+
+    def run_once(self, host, cmdline_args, release_image, dev_image,
+                 old_release_image="", test=""):
+        for i, update_info in enumerate(self.update_order):
+            erase_nvmem = self.test == self.ERASE_NVMEM
+            self.run_update(update_info[0], erase_nvmem=erase_nvmem)