xref: /aosp_15_r20/external/autotest/utils/frozen_chromite/lib/stateful_updater.py (revision 9c5db1993ded3edbeafc8092d69fe5de2ee02df7) 
1# -*- coding: utf-8 -*-
2# Copyright 2020 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""Module for updating the stateful partition on the device.
7
8Use this module to update the stateful partition given a stateful payload
9(e.g. stateful.tgz) on the device. This module untars/uncompresses the payload
10on the device into var_new and dev_image_new directories. Optinonally, you can
11ask this module to reset a stateful partition by preparing it to be clobbered on
12reboot.
13"""
14
15from __future__ import print_function
16
17import os
18import tempfile
19
20from autotest_lib.utils.frozen_chromite.lib import constants
21from autotest_lib.utils.frozen_chromite.lib import cros_build_lib
22from autotest_lib.utils.frozen_chromite.lib import cros_logging as logging
23from autotest_lib.utils.frozen_chromite.lib import osutils
24
25
26class Error(Exception):
27  """Base exception class of StatefulUpdater errors."""
28
29
30class StatefulUpdater(object):
31  """The module for updating the stateful partition."""
32
33  UPDATE_TYPE_STANDARD = 'standard'
34  UPDATE_TYPE_CLOBBER = 'clobber'
35
36  _VAR_DIR = 'var_new'
37  _DEV_IMAGE_DIR = 'dev_image_new'
38  _UPDATE_TYPE_FILE = '.update_available'
39
40  def __init__(self, device, stateful_dir=constants.STATEFUL_DIR):
41    """Initializes the module.
42
43    Args:
44      device: The ChromiumOsDevice to be updated.
45      stateful_dir: The stateful directory on the Chromium OS device.
46    """
47    self._device = device
48    self._stateful_dir = stateful_dir
49    self._var_dir = os.path.join(self._stateful_dir, self._VAR_DIR)
50    self._dev_image_dir = os.path.join(self._stateful_dir, self._DEV_IMAGE_DIR)
51    self._update_type_file = os.path.join(self._stateful_dir,
52                                          self._UPDATE_TYPE_FILE)
53
54  def Update(self, payload_path_on_device, update_type=None):
55    """Updates the stateful partition given the update file.
56
57    Args:
58      payload_path_on_device: The path to the stateful update (stateful.tgz)
59        on the DUT.
60      update_type: The type of the stateful update to be marked. Accepted
61        values: 'standard' (default) and 'clobber'.
62    """
63    if not self._device.IfPathExists(payload_path_on_device):
64      raise Error('Missing the file: %s' % payload_path_on_device)
65
66    try:
67      cmd = ['tar', '--ignore-command-error', '--overwrite',
68             '--directory', self._stateful_dir, '-xzf', payload_path_on_device]
69      self._device.run(cmd)
70    except cros_build_lib.RunCommandError as e:
71      raise Error('Failed to untar the stateful update with error %s' % e)
72
73    # Make sure target directories are generated on the device.
74    if (not self._device.IfPathExists(self._var_dir) or
75        not self._device.IfPathExists(self._dev_image_dir)):
76      raise Error('Missing var or dev_image in stateful payload.')
77
78    self._MarkUpdateType(update_type if update_type is not None
79                         else self.UPDATE_TYPE_STANDARD)
80
81  def _MarkUpdateType(self, update_type):
82    """Marks the type of the update.
83
84    Args:
85      update_type: The type of the update to be marked. See Update()
86    """
87    if update_type not in (self.UPDATE_TYPE_CLOBBER, self.UPDATE_TYPE_STANDARD):
88      raise Error('Invalid update type %s' % update_type)
89
90    with tempfile.NamedTemporaryFile() as f:
91      if update_type == self.UPDATE_TYPE_STANDARD:
92        logging.notice('Performing standard stateful update...')
93      elif update_type == self.UPDATE_TYPE_CLOBBER:
94        logging.notice('Restoring stateful to factory_install '
95                       'with dev_image...')
96        osutils.WriteFile(f.name, 'clobber')
97
98      try:
99        self._device.CopyToDevice(f.name, self._update_type_file, 'scp')
100      except cros_build_lib.RunCommandError as e:
101        raise Error('Failed to copy update type file to device with error %s' %
102                    e)
103
104  def Reset(self):
105    """Resets the stateful partition."""
106    logging.info('Resetting stateful update state.')
107
108    try:
109      self._device.run(['rm', '-rf', self._update_type_file,
110                        self._var_dir, self._dev_image_dir])
111    except cros_build_lib.RunCommandError as e:
112      logging.warning('(ignoring) Failed to delete stateful update paths with'
113                      ' error: %s', e)
114