Snap for 4498106 from dcd8f845c5131c3f5f6be2991bc33e225fe04cae to pi-release

Change-Id: Iddea6dfb8b695e871d5a1af950d34977a8fd9b6b
diff --git a/at-factory-tool/atft.py b/at-factory-tool/atft.py
index 1556227..2ee4346 100644
--- a/at-factory-tool/atft.py
+++ b/at-factory-tool/atft.py
@@ -295,6 +295,9 @@
     # The field to sort target devices
     self.sort_by = self.atft_manager.SORT_BY_LOCATION
 
+    # List of serial numbers for the devices in auto provisioning mode.
+    self.auto_dev_serials = []
+
     # Store the last refreshed target list, we use this list to prevent
     # refreshing the same list.
     self.last_target_list = []
@@ -806,7 +809,7 @@
     self.cmd_output = wx.TextCtrl(
         self.panel,
         wx.ID_ANY,
-        size=(800, 320),
+        size=(800, 190),
         style=wx.TE_MULTILINE | wx.TE_READONLY | wx.HSCROLL)
     self.vbox.Add(self.cmd_output, 0, wx.ALL | wx.EXPAND, 5)
 
@@ -814,7 +817,7 @@
     self.toolbar.Realize()
     self.statusbar = self.CreateStatusBar()
     self.statusbar.SetStatusText('Ready')
-    self.SetSize((800, 800))
+    self.SetSize((800, 720))
     self.SetTitle(self.TITLE)
     self.Center()
     self.Show(True)
@@ -1000,9 +1003,9 @@
     else:
       # Leave auto provisioning mode.
       for device in self.atft_manager.target_devs:
-        # Change all waiting devices' status to idle.
+        # Change all waiting devices' status to it's original state.
         if device.provision_status == ProvisionStatus.WAITING:
-          device.provision_status = ProvisionStatus.IDLE
+          self.atft_manager.CheckProvisionStatus(device)
       message = 'Automatic key provisioning end'
       self.PrintToCommandWindow(message)
       self.log.Info('Autoprov', message)
@@ -1244,11 +1247,14 @@
     """
     # All idle devices -> waiting.
     for target_dev in self.atft_manager.target_devs:
-      if target_dev.provision_status == ProvisionStatus.IDLE:
+      if (target_dev.serial_number not in self.auto_dev_serials and
+          target_dev.provision_status != ProvisionStatus.PROVISION_SUCCESS and
+          not ProvisionStatus.isFailed(target_dev.provision_status)
+          ):
+        self.auto_dev_serials.append(target_dev.serial_number)
         target_dev.provision_status = ProvisionStatus.WAITING
+        self._CreateThread(self._HandleStateTransition, target_dev)
 
-    for target_dev in self.atft_manager.target_devs:
-      self._CreateThread(self._HandleStateTransition, target_dev)
 
   def _HandleKeysLeft(self):
     """Display how many keys left in the ATFA device.
@@ -1596,8 +1602,7 @@
       if not target:
         continue
       # Start state could be IDLE or FUSEVBOOT_FAILED
-      if (TEST_MODE or target.provision_status == ProvisionStatus.IDLE or
-          target.provision_status == ProvisionStatus.FUSEVBOOT_FAILED):
+      if (TEST_MODE or not target.provision_state.bootloader_locked):
         target.provision_status = ProvisionStatus.WAITING
         pending_targets.append(target)
       else:
@@ -1617,6 +1622,7 @@
       target: The target device DeviceInfo object.
     """
     operation = 'Fuse bootloader verified boot key'
+    serial = target.serial_number
     self._SendOperationStartEvent(operation, target)
     self.PauseRefresh()
 
@@ -1666,6 +1672,13 @@
     # released.
     reboot_lock.acquire()
 
+    target = self.atft_manager.GetTargetDevice(serial)
+    if target and not target.provision_state.bootloader_locked:
+      target.provision_status = ProvisionStatus.FUSEVBOOT_FAILED
+      e = FastbootFailure('Status not updated.')
+      self._HandleException('E', e, operation)
+      return
+
   def _RebootSuccessCallback(self, msg, lock):
     """The callback if reboot succeed.
 
@@ -1702,10 +1715,10 @@
       # Start state could be FUSEVBOOT_SUCCESS or REBOOT_SUCCESS
       # or FUSEATTR_FAILED
       # Note: Reboot to check vboot is optional, user can skip that manually.
-      if (TEST_MODE or
-          target.provision_status == ProvisionStatus.FUSEVBOOT_SUCCESS or
-          target.provision_status == ProvisionStatus.REBOOT_SUCCESS or
-          target.provision_status == ProvisionStatus.FUSEATTR_FAILED):
+      if (TEST_MODE or (
+            target.provision_state.bootloader_locked and
+            not target.provision_state.avb_perm_attr_set
+          )):
         pending_targets.append(target)
       else:
         self._SendAlertEvent(self.ALERT_FUSE_PERM_ATTR_FUSED)
@@ -1748,9 +1761,11 @@
       if not target:
         continue
       # Start state could be FUSEATTR_SUCCESS or LOCKAVB_FAIELD
-      if (TEST_MODE or
-          target.provision_status == ProvisionStatus.FUSEATTR_SUCCESS or
-          target.provision_status == ProvisionStatus.LOCKAVB_FAILED):
+      if (TEST_MODE or(
+            target.provision_state.bootloader_locked and
+            target.provision_state.avb_perm_attr_set and
+            not target.provision_state.avb_locked
+          )):
         target.provision_status = ProvisionStatus.WAITING
         pending_targets.append(target)
       else:
@@ -1871,8 +1886,12 @@
         continue
       pending_targets.append(target_dev)
       status = target_dev.provision_status
-      if (TEST_MODE or status == ProvisionStatus.LOCKAVB_SUCCESS or
-          status == ProvisionStatus.PROVISION_FAILED):
+      if (TEST_MODE or (
+          target_dev.provision_state.bootloader_locked and
+          target_dev.provision_state.avb_perm_attr_set and
+          target_dev.provision_state.avb_locked and
+          not target_dev.provision_state.provisioned
+        )):
         target_dev.provision_status = ProvisionStatus.WAITING
       else:
         self._SendAlertEvent(self.ALERT_PROV_PROVED)
@@ -1918,32 +1937,33 @@
       target: The target device object.
     """
     self.auto_prov_lock.acquire()
-    while target.provision_status != ProvisionStatus.PROVISION_SUCCESS:
-      if target.provision_status == ProvisionStatus.WAITING:
+    serial = target.serial_number
+    while not ProvisionStatus.isFailed(target.provision_status):
+      target = self.atft_manager.GetTargetDevice(serial)
+      if not target:
+        # The target disappear somehow.
+        break
+      if not self.auto_prov:
+        # Auto provision mode exited.
+        break
+      if not target.provision_state.bootloader_locked:
         self._FuseVbootKeyTarget(target)
-        if (target.provision_status == ProvisionStatus.FUSEVBOOT_FAILED or
-            target.provision_status == ProvisionStatus.REBOOT_FAILED):
-          break
-      elif (target.provision_status == ProvisionStatus.FUSEVBOOT_SUCCESS or
-            target.provision_status == ProvisionStatus.REBOOT_SUCCESS):
+        continue
+      elif not target.provision_state.avb_perm_attr_set:
         self._FusePermAttrTarget(target)
-        if target.provision_status == ProvisionStatus.FUSEATTR_FAILED:
-          break
-      elif target.provision_status == ProvisionStatus.FUSEATTR_SUCCESS:
+        continue
+      elif not target.provision_state.avb_locked:
         self._LockAvbTarget(target)
-        if target.provision_status == ProvisionStatus.LOCKAVB_FAILED:
-          break
-      elif target.provision_status == ProvisionStatus.LOCKAVB_SUCCESS:
+        continue
+      elif not target.provision_state.provisioned:
         self._ProvisionTarget(target)
         if self.atft_manager.GetATFAKeysLeft() == 0:
           # No keys left. If it's auto provisioning mode, exit.
           self._SendAlertEvent(self.ALERT_NO_KEYS_LEFT_LEAVE_PROV)
           self.toolbar.ToggleTool(self.ID_TOOL_PROVISION, False)
           self.OnToggleAutoProv(None)
-        if target.provision_status == ProvisionStatus.PROVISION_FAILED:
-          break
-      else:
-        break
+      break
+    self.auto_dev_serials.remove(serial)
     self.auto_prov_lock.release()
 
   def _ProcessKey(self):
diff --git a/at-factory-tool/atft_unittest.py b/at-factory-tool/atft_unittest.py
index d60a6b2..5da0b1a 100644
--- a/at-factory-tool/atft_unittest.py
+++ b/at-factory-tool/atft_unittest.py
@@ -18,6 +18,7 @@
 
 import atft
 from atftman import ProvisionStatus
+from atftman import ProvisionState
 import fastboot_exceptions
 from mock import call
 from mock import MagicMock
@@ -58,6 +59,7 @@
     self.serial_number = serial_number
     self.location = location
     self.provision_status = provision_status
+    self.provision_state = ProvisionState()
     self.time_set = False
 
   def __eq__(self, other):
@@ -314,11 +316,15 @@
                                ProvisionStatus.WAITING)
     mock_atft.atft_manager.target_devs.append(test_dev1)
     mock_atft.atft_manager.target_devs.append(test_dev2)
+    mock_atft.atft_manager.CheckProvisionStatus.side_effect = (
+        lambda target=test_dev2, state=ProvisionStatus.LOCKAVB_SUCCESS:
+        self.MockStateChange(target, state))
     mock_atft.OnToggleAutoProv(None)
     self.assertEqual(False, mock_atft.auto_prov)
     self.assertEqual(test_dev1.provision_status,
                      ProvisionStatus.PROVISION_SUCCESS)
-    self.assertEqual(test_dev2.provision_status, ProvisionStatus.IDLE)
+    mock_atft.atft_manager.CheckProvisionStatus.assert_called_once()
+    self.assertEqual(test_dev2.provision_status, ProvisionStatus.LOCKAVB_SUCCESS)
 
   # Test atft.OnChangeKeyThreshold
   def testOnChangeKeyThreshold(self):
@@ -358,7 +364,7 @@
     mock_atft._HandleStateTransition = MagicMock()
     mock_atft._HandleAutoProv()
     self.assertEqual(test_dev2.provision_status, ProvisionStatus.WAITING)
-    self.assertEqual(2, mock_atft._CreateThread.call_count)
+    self.assertEqual(1, mock_atft._CreateThread.call_count)
 
   # Test atft._HandleKeysLeft
   def MockGetKeysLeft(self, keys_left_array):
@@ -406,6 +412,14 @@
   # Test atft._HandleStateTransition
   def MockStateChange(self, target, state):
     target.provision_status = state
+    if state == ProvisionStatus.REBOOT_SUCCESS:
+      target.provision_state.bootloader_locked = True
+    if state == ProvisionStatus.FUSEATTR_SUCCESS:
+      target.provision_state.avb_perm_attr_set = True
+    if state == ProvisionStatus.LOCKAVB_SUCCESS:
+      target.provision_state.avb_locked = True
+    if state == ProvisionStatus.PROVISION_SUCCESS:
+      target.provision_state.provisioned = True
 
   def testHandleStateTransition(self):
     mock_atft = MockAtft()
@@ -427,6 +441,11 @@
     mock_atft._ProvisionTarget.side_effect = (
         lambda target=mock_atft, state=ProvisionStatus.PROVISION_SUCCESS:
         self.MockStateChange(target, state))
+    mock_atft.auto_dev_serials = [self.TEST_SERIAL1]
+    mock_atft.auto_prov = True
+    mock_atft.atft_manager = MagicMock()
+    mock_atft.atft_manager.GetTargetDevice = MagicMock()
+    mock_atft.atft_manager.GetTargetDevice.return_value = test_dev1
     mock_atft._HandleStateTransition(test_dev1)
     self.assertEqual(ProvisionStatus.PROVISION_SUCCESS,
                      test_dev1.provision_status)
@@ -451,6 +470,11 @@
     mock_atft._ProvisionTarget.side_effect = (
         lambda target=mock_atft, state=ProvisionStatus.PROVISION_SUCCESS:
             self.MockStateChange(target, state))
+    mock_atft.auto_dev_serials = [self.TEST_SERIAL1]
+    mock_atft.auto_prov = True
+    mock_atft.atft_manager = MagicMock()
+    mock_atft.atft_manager.GetTargetDevice = MagicMock()
+    mock_atft.atft_manager.GetTargetDevice.return_value = test_dev1
     mock_atft._HandleStateTransition(test_dev1)
     self.assertEqual(ProvisionStatus.FUSEVBOOT_FAILED,
                      test_dev1.provision_status)
@@ -475,6 +499,11 @@
     mock_atft._ProvisionTarget.side_effect = (
         lambda target=mock_atft, state=ProvisionStatus.PROVISION_SUCCESS:
         self.MockStateChange(target, state))
+    mock_atft.auto_dev_serials = [self.TEST_SERIAL1]
+    mock_atft.auto_prov = True
+    mock_atft.atft_manager = MagicMock()
+    mock_atft.atft_manager.GetTargetDevice = MagicMock()
+    mock_atft.atft_manager.GetTargetDevice.return_value = test_dev1
     mock_atft._HandleStateTransition(test_dev1)
     self.assertEqual(ProvisionStatus.REBOOT_FAILED, test_dev1.provision_status)
 
@@ -498,6 +527,11 @@
     mock_atft._ProvisionTarget.side_effect = (
         lambda target=mock_atft, state=ProvisionStatus.PROVISION_SUCCESS:
         self.MockStateChange(target, state))
+    mock_atft.auto_dev_serials = [self.TEST_SERIAL1]
+    mock_atft.auto_prov = True
+    mock_atft.atft_manager = MagicMock()
+    mock_atft.atft_manager.GetTargetDevice = MagicMock()
+    mock_atft.atft_manager.GetTargetDevice.return_value = test_dev1
     mock_atft._HandleStateTransition(test_dev1)
     self.assertEqual(ProvisionStatus.FUSEATTR_FAILED,
                      test_dev1.provision_status)
@@ -522,6 +556,11 @@
     mock_atft._ProvisionTarget.side_effect = (
         lambda target=mock_atft, state=ProvisionStatus.PROVISION_SUCCESS:
         self.MockStateChange(target, state))
+    mock_atft.auto_dev_serials = [self.TEST_SERIAL1]
+    mock_atft.auto_prov = True
+    mock_atft.atft_manager = MagicMock()
+    mock_atft.atft_manager.GetTargetDevice = MagicMock()
+    mock_atft.atft_manager.GetTargetDevice.return_value = test_dev1
     mock_atft._HandleStateTransition(test_dev1)
     self.assertEqual(ProvisionStatus.LOCKAVB_FAILED, test_dev1.provision_status)
 
@@ -545,10 +584,55 @@
     mock_atft._ProvisionTarget.side_effect = (
         lambda target=mock_atft, state=ProvisionStatus.PROVISION_FAILED:
         self.MockStateChange(target, state))
+    mock_atft.auto_dev_serials = [self.TEST_SERIAL1]
+    mock_atft.auto_prov = True
+    mock_atft.atft_manager = MagicMock()
+    mock_atft.atft_manager.GetTargetDevice = MagicMock()
+    mock_atft.atft_manager.GetTargetDevice.return_value = test_dev1
     mock_atft._HandleStateTransition(test_dev1)
     self.assertEqual(ProvisionStatus.PROVISION_FAILED,
                      test_dev1.provision_status)
 
+  def testHandleStateTransitionSkipStep(self):
+    mock_atft = MockAtft()
+    test_dev1 = TestDeviceInfo(self.TEST_SERIAL1, self.TEST_LOCATION1,
+                               ProvisionStatus.WAITING)
+    mock_atft._FuseVbootKeyTarget = MagicMock()
+    mock_atft._FuseVbootKeyTarget.side_effect = (
+        lambda target=mock_atft, state=ProvisionStatus.REBOOT_SUCCESS:
+        self.MockStateChange(target, state))
+    mock_atft._FusePermAttrTarget = MagicMock()
+    mock_atft._FusePermAttrTarget.side_effect = (
+        lambda target=mock_atft, state=ProvisionStatus.FUSEATTR_SUCCESS:
+        self.MockStateChange(target, state))
+    mock_atft._LockAvbTarget = MagicMock()
+    mock_atft._LockAvbTarget.side_effect = (
+        lambda target=mock_atft, state=ProvisionStatus.LOCKAVB_SUCCESS:
+        self.MockStateChange(target, state))
+    mock_atft._ProvisionTarget = MagicMock()
+    mock_atft._ProvisionTarget.side_effect = (
+        lambda target=mock_atft, state=ProvisionStatus.PROVISION_SUCCESS:
+        self.MockStateChange(target, state))
+
+    # The device has bootloader_locked and avb_locked set. Should fuse perm attr
+    # and provision key.
+    test_dev1.provision_state.bootloader_locked = True
+    test_dev1.provision_state.avb_locked = True
+    mock_atft.auto_dev_serials = [self.TEST_SERIAL1]
+    mock_atft.auto_prov = True
+    mock_atft.atft_manager = MagicMock()
+    mock_atft.atft_manager.GetTargetDevice = MagicMock()
+    mock_atft.atft_manager.GetTargetDevice.return_value = test_dev1
+    mock_atft._HandleStateTransition(test_dev1)
+    self.assertEqual(ProvisionStatus.PROVISION_SUCCESS,
+                     test_dev1.provision_status)
+    mock_atft._FusePermAttrTarget.assert_called_once()
+    mock_atft._ProvisionTarget.assert_called_once()
+    self.assertEqual(True, test_dev1.provision_state.bootloader_locked)
+    self.assertEqual(True, test_dev1.provision_state.avb_perm_attr_set)
+    self.assertEqual(True, test_dev1.provision_state.avb_locked)
+    self.assertEqual(True, test_dev1.provision_state.provisioned)
+
   # Test atft._CheckATFAStatus
   def testCheckATFAStatus(self):
     mock_atft = MockAtft()
@@ -565,6 +649,7 @@
 
   def MockReboot(self, target, timeout, success, fail):
     success()
+    target.provision_state.bootloader_locked = True
 
   @patch('wx.QueueEvent')
   def testFuseVbootKey(self, mock_queue_event):
@@ -584,6 +669,7 @@
                                ProvisionStatus.FUSEVBOOT_FAILED)
     test_dev3 = TestDeviceInfo(self.TEST_SERIAL3, self.TEST_LOCATION2,
                                ProvisionStatus.FUSEVBOOT_SUCCESS)
+    test_dev3.provision_state.bootloader_locked = True
     self.device_map[self.TEST_SERIAL1] = test_dev1
     self.device_map[self.TEST_SERIAL2] = test_dev2
     self.device_map[self.TEST_SERIAL3] = test_dev3
@@ -613,6 +699,7 @@
                                ProvisionStatus.FUSEVBOOT_FAILED)
     test_dev3 = TestDeviceInfo(self.TEST_SERIAL3, self.TEST_LOCATION2,
                                ProvisionStatus.FUSEVBOOT_SUCCESS)
+    test_dev3.provision_state.bootloader_locked = True
     self.device_map[self.TEST_SERIAL1] = test_dev1
     self.device_map[self.TEST_SERIAL2] = test_dev2
     self.device_map[self.TEST_SERIAL3] = test_dev3
@@ -632,6 +719,7 @@
                                ProvisionStatus.FUSEVBOOT_FAILED)
     test_dev3 = TestDeviceInfo(self.TEST_SERIAL3, self.TEST_LOCATION2,
                                ProvisionStatus.FUSEVBOOT_SUCCESS)
+    test_dev3.provision_state.bootloader_locked = True
     self.device_map[self.TEST_SERIAL1] = test_dev1
     self.device_map[self.TEST_SERIAL2] = test_dev2
     self.device_map[self.TEST_SERIAL3] = test_dev3
@@ -649,6 +737,7 @@
                                ProvisionStatus.FUSEVBOOT_FAILED)
     test_dev3 = TestDeviceInfo(self.TEST_SERIAL3, self.TEST_LOCATION2,
                                ProvisionStatus.FUSEVBOOT_SUCCESS)
+    test_dev3.provision_state.bootloader_locked = True
     self.device_map[self.TEST_SERIAL1] = test_dev1
     self.device_map[self.TEST_SERIAL2] = test_dev2
     self.device_map[self.TEST_SERIAL3] = test_dev3
@@ -670,14 +759,23 @@
     mock_atft._SendMessageEvent = MagicMock()
     mock_atft.atft_manager.GetTargetDevice.side_effect = (
         self.MockGetTargetDevice)
-    test_dev1 = TestDeviceInfo(self.TEST_SERIAL1, self.TEST_LOCATION1,
-                               ProvisionStatus.FUSEVBOOT_SUCCESS)
-    test_dev2 = TestDeviceInfo(self.TEST_SERIAL2, self.TEST_LOCATION2,
-                               ProvisionStatus.REBOOT_SUCCESS)
-    test_dev3 = TestDeviceInfo(self.TEST_SERIAL3, self.TEST_LOCATION2,
-                               ProvisionStatus.FUSEATTR_FAILED)
-    test_dev4 = TestDeviceInfo(self.TEST_SERIAL4, self.TEST_LOCATION2,
-                               ProvisionStatus.FUSEATTR_SUCCESS)
+    test_dev1 = TestDeviceInfo(
+        self.TEST_SERIAL1, self.TEST_LOCATION1,
+        ProvisionStatus.FUSEVBOOT_SUCCESS)
+    test_dev1.provision_state.bootloader_locked = True
+    test_dev2 = TestDeviceInfo(
+        self.TEST_SERIAL2, self.TEST_LOCATION2,
+        ProvisionStatus.REBOOT_SUCCESS)
+    test_dev2.provision_state.bootloader_locked = True
+    test_dev3 = TestDeviceInfo(
+        self.TEST_SERIAL3, self.TEST_LOCATION2,
+        ProvisionStatus.FUSEATTR_FAILED)
+    test_dev3.provision_state.bootloader_locked = True
+    test_dev4 = TestDeviceInfo(
+        self.TEST_SERIAL4, self.TEST_LOCATION2,
+        ProvisionStatus.FUSEATTR_SUCCESS)
+    test_dev4.provision_state.bootloader_locked = True
+    test_dev4.provision_state.avb_perm_attr_set = True
     self.device_map[self.TEST_SERIAL1] = test_dev1
     self.device_map[self.TEST_SERIAL2] = test_dev2
     self.device_map[self.TEST_SERIAL3] = test_dev3
@@ -703,12 +801,17 @@
         self.MockGetTargetDevice)
     test_dev1 = TestDeviceInfo(self.TEST_SERIAL1, self.TEST_LOCATION1,
                                ProvisionStatus.FUSEVBOOT_SUCCESS)
+    test_dev1.provision_state.bootloader_locked = True
     test_dev2 = TestDeviceInfo(self.TEST_SERIAL2, self.TEST_LOCATION2,
                                ProvisionStatus.REBOOT_SUCCESS)
+    test_dev2.provision_state.bootloader_locked = True
     test_dev3 = TestDeviceInfo(self.TEST_SERIAL3, self.TEST_LOCATION2,
                                ProvisionStatus.FUSEATTR_FAILED)
+    test_dev3.provision_state.bootloader_locked = True
     test_dev4 = TestDeviceInfo(self.TEST_SERIAL4, self.TEST_LOCATION2,
                                ProvisionStatus.FUSEATTR_SUCCESS)
+    test_dev4.provision_state.bootloader_locked = True
+    test_dev4.provision_state.avb_perm_attr_set = True
     self.device_map[self.TEST_SERIAL1] = test_dev1
     self.device_map[self.TEST_SERIAL2] = test_dev2
     self.device_map[self.TEST_SERIAL3] = test_dev3
@@ -725,10 +828,13 @@
     mock_atft._HandleException.reset_mock()
     test_dev1 = TestDeviceInfo(self.TEST_SERIAL1, self.TEST_LOCATION1,
                                ProvisionStatus.FUSEVBOOT_SUCCESS)
+    test_dev1.provision_state.bootloader_locked = True
     test_dev2 = TestDeviceInfo(self.TEST_SERIAL2, self.TEST_LOCATION2,
                                ProvisionStatus.REBOOT_SUCCESS)
+    test_dev2.provision_state.bootloader_locked = True
     test_dev3 = TestDeviceInfo(self.TEST_SERIAL3, self.TEST_LOCATION2,
                                ProvisionStatus.FUSEATTR_FAILED)
+    test_dev3.provision_state.bootloader_locked = True
     test_dev4 = TestDeviceInfo(self.TEST_SERIAL4, self.TEST_LOCATION2,
                                ProvisionStatus.FUSEATTR_SUCCESS)
     self.device_map[self.TEST_SERIAL1] = test_dev1
@@ -754,10 +860,16 @@
         self.MockGetTargetDevice)
     test_dev1 = TestDeviceInfo(self.TEST_SERIAL1, self.TEST_LOCATION1,
                                ProvisionStatus.FUSEATTR_SUCCESS)
+    test_dev1.provision_state.bootloader_locked = True
+    test_dev1.provision_state.avb_perm_attr_set = True
     test_dev2 = TestDeviceInfo(self.TEST_SERIAL2, self.TEST_LOCATION2,
                                ProvisionStatus.LOCKAVB_FAILED)
+    test_dev2.provision_state.bootloader_locked = True
+    test_dev2.provision_state.avb_perm_attr_set = True
     test_dev3 = TestDeviceInfo(self.TEST_SERIAL3, self.TEST_LOCATION2,
                                ProvisionStatus.FUSEATTR_FAILED)
+    test_dev3.provision_state.bootloader_locked = True
+    test_dev3.provision_state.avb_perm_attr_set = False
     test_dev4 = TestDeviceInfo(self.TEST_SERIAL4, self.TEST_LOCATION2,
                                ProvisionStatus.IDLE)
     self.device_map[self.TEST_SERIAL1] = test_dev1
@@ -785,10 +897,16 @@
         self.MockGetTargetDevice)
     test_dev1 = TestDeviceInfo(self.TEST_SERIAL1, self.TEST_LOCATION1,
                                ProvisionStatus.FUSEATTR_SUCCESS)
+    test_dev1.provision_state.bootloader_locked = True
+    test_dev1.provision_state.avb_perm_attr_set = True
     test_dev2 = TestDeviceInfo(self.TEST_SERIAL2, self.TEST_LOCATION2,
                                ProvisionStatus.LOCKAVB_FAILED)
+    test_dev2.provision_state.bootloader_locked = True
+    test_dev2.provision_state.avb_perm_attr_set = True
     test_dev3 = TestDeviceInfo(self.TEST_SERIAL3, self.TEST_LOCATION2,
                                ProvisionStatus.FUSEATTR_FAILED)
+    test_dev3.provision_state.bootloader_locked = True
+    test_dev3.provision_state.avb_perm_attr_set = False
     test_dev4 = TestDeviceInfo(self.TEST_SERIAL4, self.TEST_LOCATION2,
                                ProvisionStatus.IDLE)
     self.device_map[self.TEST_SERIAL1] = test_dev1
@@ -962,10 +1080,17 @@
         self.MockGetTargetDevice)
     test_dev1 = TestDeviceInfo(self.TEST_SERIAL1, self.TEST_LOCATION1,
                                ProvisionStatus.PROVISION_FAILED)
+    test_dev1.provision_state.bootloader_locked = True
+    test_dev1.provision_state.avb_perm_attr_set = True
+    test_dev1.provision_state.avb_locked = True
     test_dev2 = TestDeviceInfo(self.TEST_SERIAL2, self.TEST_LOCATION2,
                                ProvisionStatus.LOCKAVB_SUCCESS)
+    test_dev2.provision_state.bootloader_locked = True
+    test_dev2.provision_state.avb_perm_attr_set = True
+    test_dev2.provision_state.avb_locked = True
     test_dev3 = TestDeviceInfo(self.TEST_SERIAL3, self.TEST_LOCATION2,
                                ProvisionStatus.FUSEATTR_FAILED)
+    test_dev3.provision_state.bootloader_locked = True
     self.device_map[self.TEST_SERIAL1] = test_dev1
     self.device_map[self.TEST_SERIAL2] = test_dev2
     self.device_map[self.TEST_SERIAL3] = test_dev3
@@ -988,10 +1113,17 @@
         self.MockGetTargetDevice)
     test_dev1 = TestDeviceInfo(self.TEST_SERIAL1, self.TEST_LOCATION1,
                                ProvisionStatus.PROVISION_FAILED)
+    test_dev1.provision_state.bootloader_locked = True
+    test_dev1.provision_state.avb_perm_attr_set = True
+    test_dev1.provision_state.avb_locked = True
     test_dev2 = TestDeviceInfo(self.TEST_SERIAL2, self.TEST_LOCATION2,
                                ProvisionStatus.LOCKAVB_SUCCESS)
+    test_dev2.provision_state.bootloader_locked = True
+    test_dev2.provision_state.avb_perm_attr_set = True
+    test_dev2.provision_state.avb_locked = True
     test_dev3 = TestDeviceInfo(self.TEST_SERIAL3, self.TEST_LOCATION2,
                                ProvisionStatus.FUSEATTR_FAILED)
+    test_dev3.provision_state.bootloader_locked = True
     self.device_map[self.TEST_SERIAL1] = test_dev1
     self.device_map[self.TEST_SERIAL2] = test_dev2
     self.device_map[self.TEST_SERIAL3] = test_dev3
@@ -1002,10 +1134,17 @@
     self.assertEqual(2, mock_atft._HandleException.call_count)
     test_dev1 = TestDeviceInfo(self.TEST_SERIAL1, self.TEST_LOCATION1,
                                ProvisionStatus.PROVISION_FAILED)
+    test_dev1.provision_state.bootloader_locked = True
+    test_dev1.provision_state.avb_perm_attr_set = True
+    test_dev1.provision_state.avb_locked = True
     test_dev2 = TestDeviceInfo(self.TEST_SERIAL2, self.TEST_LOCATION2,
                                ProvisionStatus.LOCKAVB_SUCCESS)
+    test_dev2.provision_state.bootloader_locked = True
+    test_dev2.provision_state.avb_perm_attr_set = True
+    test_dev2.provision_state.avb_locked = True
     test_dev3 = TestDeviceInfo(self.TEST_SERIAL3, self.TEST_LOCATION2,
                                ProvisionStatus.FUSEATTR_FAILED)
+    test_dev3.provision_state.bootloader_locked = True
     self.device_map[self.TEST_SERIAL1] = test_dev1
     self.device_map[self.TEST_SERIAL2] = test_dev2
     self.device_map[self.TEST_SERIAL3] = test_dev3
diff --git a/at-factory-tool/atftman.py b/at-factory-tool/atftman.py
index c51814c..1a1fadc 100644
--- a/at-factory-tool/atftman.py
+++ b/at-factory-tool/atftman.py
@@ -93,6 +93,18 @@
   def ToString(provision_status, language_index):
     return ProvisionStatus.STRING_MAP[provision_status][language_index]
 
+  @staticmethod
+  def isSuccess(provision_status):
+    return provision_status % 10 == ProvisionStatus._SUCCESS
+
+  @staticmethod
+  def isProcessing(provision_status):
+    return provision_status % 10 == ProvisionStatus._PROCESSING
+
+  @staticmethod
+  def isFailed(provision_status):
+    return provision_status % 10 == ProvisionStatus._FAILED
+
 
 class ProvisionState(object):
   """The provision state of the target device."""
@@ -150,6 +162,9 @@
   def Oem(self, oem_command, err_to_out=False):
     return self._fastboot_device_controller.Oem(oem_command, err_to_out)
 
+  def Flash(self, partition, file_path):
+    return self._fastboot_device_controller.Flash(partition, file_path)
+
   def Upload(self, file_path):
     return self._fastboot_device_controller.Upload(file_path)
 
@@ -442,15 +457,21 @@
       serial_location_map: The serial location map.
       check_status: Whether to check provision status for the target device.
     """
-    controller = self._fastboot_device_controller(serial)
-    location = None
-    if serial in serial_location_map:
-      location = serial_location_map[serial]
+    try:
+      controller = self._fastboot_device_controller(serial)
+      location = None
+      if serial in serial_location_map:
+        location = serial_location_map[serial]
 
-    new_target_dev = DeviceInfo(controller, serial, location)
-    if check_status:
-      self.CheckProvisionStatus(new_target_dev)
-    self.target_devs.append(new_target_dev)
+      new_target_dev = DeviceInfo(controller, serial, location)
+      if check_status:
+        self.CheckProvisionStatus(new_target_dev)
+      self.target_devs.append(new_target_dev)
+    except FastbootFailure as e:
+      e.msg = ('Error while creating new device: ' + str(new_target_dev) +
+               '\n'+ e.msg)
+      self.stable_serials.remove(serial)
+      raise e
 
   def _AddNewAtfa(self, atfa_serial):
     """Create a new ATFA device object.
@@ -532,12 +553,8 @@
         if callback_lock and callback_lock.acquire(False):
           success_serials.append(serial)
 
-    serial_location_map = self._serial_mapper.get_serial_map()
     for serial in success_serials:
       self._reboot_callbacks[serial].success()
-      self._CreateNewTargetDevice(serial, serial_location_map, False)
-      self.GetTargetDevice(serial).provision_status = (
-          ProvisionStatus.REBOOT_SUCCESS)
 
   def _ParseStateString(self, state_string):
     """Parse the string returned by 'at-vboot-state' to a key-value map.
@@ -607,6 +624,7 @@
         target_dev.provision_status = ProvisionStatus.FUSEVBOOT_SUCCESS
       target_dev.provision_state.bootloader_locked = True
 
+
   def TransferContent(self, src, dst):
     """Transfer content from a device to another device.
 
@@ -686,7 +704,7 @@
       self.CheckProvisionStatus(target)
       if not target.provision_state.provisioned:
         raise FastbootFailure('Status not updated.')
-    except FastbootFailure as e:
+    except (FastbootFailure, DeviceNotFoundException) as e:
       target.provision_status = ProvisionStatus.PROVISION_FAILED
       raise e
 
@@ -716,6 +734,10 @@
       self.CheckProvisionStatus(target)
       if not target.provision_state.bootloader_locked:
         raise FastbootFailure('Status not updated.')
+
+      # # Another possible flow:
+      # target.Flash('sec', temp_file_name)
+      # os.remove(temp_file_name)
     except FastbootFailure as e:
       target.provision_status = ProvisionStatus.FUSEVBOOT_FAILED
       raise e
@@ -793,15 +815,15 @@
 
       reboot_callback = RebootCallback(
           timeout,
-          self.RebootCallbackWrapper(success_callback, serial),
-          self.RebootCallbackWrapper(timeout_callback, serial))
+          self.RebootCallbackWrapper(success_callback, serial, True),
+          self.RebootCallbackWrapper(timeout_callback, serial, False))
       self._reboot_callbacks[serial] = reboot_callback
 
     except FastbootFailure as e:
       target.provision_status = ProvisionStatus.REBOOT_FAILED
       raise e
 
-  def RebootCallbackWrapper(self, callback, serial):
+  def RebootCallbackWrapper(self, callback, serial, success):
     """This wrapper function wraps the original callback function.
 
     Some clean up operations are added. We need to remove the handler if
@@ -812,17 +834,30 @@
     Args:
       callback: The original callback function.
       serial: The serial number for the device.
+      success: Whether this is the success callback.
     Returns:
       An extended callback function.
     """
-    def RebootCallbackFunc(callback=callback, serial=serial):
-      rebooting_dev = self.GetTargetDevice(serial)
-      if rebooting_dev:
-        self.target_devs.remove(rebooting_dev)
-        del rebooting_dev
-      callback()
-      self._reboot_callbacks[serial].Release()
-      del self._reboot_callbacks[serial]
+    def RebootCallbackFunc(callback=callback, serial=serial, success=success):
+      try:
+        rebooting_dev = self.GetTargetDevice(serial)
+        if rebooting_dev:
+          self.target_devs.remove(rebooting_dev)
+          del rebooting_dev
+        if success:
+          serial_location_map = self._serial_mapper.get_serial_map()
+          self._CreateNewTargetDevice(serial, serial_location_map, True)
+          self.GetTargetDevice(serial).provision_status = (
+              ProvisionStatus.REBOOT_SUCCESS)
+        callback()
+        self._reboot_callbacks[serial].Release()
+        del self._reboot_callbacks[serial]
+      except FastbootFailure as e:
+        # Release the lock so that it can be obtained again.
+        self._reboot_callbacks[serial].lock.release()
+        raise e
+
+
 
     return RebootCallbackFunc
 
diff --git a/at-factory-tool/fastbootsh.py b/at-factory-tool/fastbootsh.py
index 593e9f1..ceba490 100644
--- a/at-factory-tool/fastbootsh.py
+++ b/at-factory-tool/fastbootsh.py
@@ -109,6 +109,27 @@
     finally:
       self._lock.release()
 
+  def Flash(self, partition, file_path):
+    """Flash a file to a partition.
+
+    Args:
+      file_path: The partition file to be flashed.
+      partition: The partition to be flashed.
+    Returns:
+      The output for the fastboot command required.
+    Raises:
+      FastbootFailure: If failure happens during the command.
+    """
+    try:
+      self._lock.acquire()
+      out = self.fastboot_command(
+          '-s', self.serial_number, 'flash', partition, file_path)
+      return out
+    except sh.ErrorReturnCode as e:
+      raise fastboot_exceptions.FastbootFailure(e.stderr)
+    finally:
+      self._lock.release()
+
   def Upload(self, file_path):
     """Pulls a file from the fastboot device to the local file system.
 
diff --git a/at-factory-tool/fastbootsubp.py b/at-factory-tool/fastbootsubp.py
index bd8ca12..7447c0b 100644
--- a/at-factory-tool/fastbootsubp.py
+++ b/at-factory-tool/fastbootsubp.py
@@ -118,6 +118,30 @@
     finally:
       self._lock.release()
 
+  def Flash(self, partition, file_path):
+    """Flash a file to a partition.
+
+    Args:
+      file_path: The partition file to be flashed.
+      partition: The partition to be flashed.
+    Returns:
+      The output for the fastboot command required.
+    Raises:
+      FastbootFailure: If failure happens during the command.
+    """
+    try:
+      self._lock.acquire()
+      return subprocess.check_output(
+          [
+              FastbootDevice.fastboot_command, '-s', self.serial_number,
+              'flash', partition, file_path
+          ],
+          creationflags=CREATE_NO_WINDOW)
+    except subprocess.CalledProcessError as e:
+      raise fastboot_exceptions.FastbootFailure(e.output)
+    finally:
+      self._lock.release()
+
   def Upload(self, file_path):
     """Pulls a file from the fastboot device to the local file system.