xref: /aosp_15_r20/external/toolchain-utils/image_chromeos.py (revision 760c253c1ed00ce9abd48f8546f08516e57485fe)
1*760c253cSXin Li#!/usr/bin/env python3
2*760c253cSXin Li# -*- coding: utf-8 -*-
3*760c253cSXin Li#
4*760c253cSXin Li# Copyright 2019 The ChromiumOS Authors
5*760c253cSXin Li# Use of this source code is governed by a BSD-style license that can be
6*760c253cSXin Li# found in the LICENSE file.
7*760c253cSXin Li
8*760c253cSXin Li"""Script to image a ChromeOS device.
9*760c253cSXin Li
10*760c253cSXin LiThis script images a remote ChromeOS device with a specific image."
11*760c253cSXin Li"""
12*760c253cSXin Li
13*760c253cSXin Li
14*760c253cSXin Li__author__ = "[email protected] (Ahmad Sharif)"
15*760c253cSXin Li
16*760c253cSXin Liimport argparse
17*760c253cSXin Liimport filecmp
18*760c253cSXin Liimport glob
19*760c253cSXin Liimport os
20*760c253cSXin Liimport re
21*760c253cSXin Liimport shutil
22*760c253cSXin Liimport sys
23*760c253cSXin Liimport tempfile
24*760c253cSXin Liimport time
25*760c253cSXin Li
26*760c253cSXin Lifrom cros_utils import command_executer
27*760c253cSXin Lifrom cros_utils import locks
28*760c253cSXin Lifrom cros_utils import logger
29*760c253cSXin Lifrom cros_utils import misc
30*760c253cSXin Lifrom cros_utils.file_utils import FileUtils
31*760c253cSXin Li
32*760c253cSXin Li
33*760c253cSXin Lichecksum_file = "/usr/local/osimage_checksum_file"
34*760c253cSXin Lilock_file = "/tmp/image_chromeos_lock/image_chromeos_lock"
35*760c253cSXin Li
36*760c253cSXin Li
37*760c253cSXin Lidef Usage(parser, message):
38*760c253cSXin Li    print("ERROR: %s" % message)
39*760c253cSXin Li    parser.print_help()
40*760c253cSXin Li    sys.exit(0)
41*760c253cSXin Li
42*760c253cSXin Li
43*760c253cSXin Lidef CheckForCrosFlash(chromeos_root, remote, log_level):
44*760c253cSXin Li    cmd_executer = command_executer.GetCommandExecuter(log_level=log_level)
45*760c253cSXin Li
46*760c253cSXin Li    # Check to see if remote machine has cherrypy, ctypes
47*760c253cSXin Li    command = "python -c 'import cherrypy, ctypes'"
48*760c253cSXin Li    ret = cmd_executer.CrosRunCommand(
49*760c253cSXin Li        command, chromeos_root=chromeos_root, machine=remote
50*760c253cSXin Li    )
51*760c253cSXin Li    logger.GetLogger().LogFatalIf(
52*760c253cSXin Li        ret == 255, f"Failed ssh to {remote} (for checking cherrypy)"
53*760c253cSXin Li    )
54*760c253cSXin Li    logger.GetLogger().LogFatalIf(
55*760c253cSXin Li        ret != 0,
56*760c253cSXin Li        f"Failed to find cherrypy or ctypes on '{remote}', "
57*760c253cSXin Li        "cros flash cannot work.",
58*760c253cSXin Li    )
59*760c253cSXin Li
60*760c253cSXin Li
61*760c253cSXin Lidef DisableCrosBeeps(chromeos_root, remote, log_level):
62*760c253cSXin Li    """Disable annoying chromebooks beeps after reboots."""
63*760c253cSXin Li    cmd_executer = command_executer.GetCommandExecuter(log_level=log_level)
64*760c253cSXin Li
65*760c253cSXin Li    command = "/usr/bin/futility gbb --set --flash --flags=0x1"
66*760c253cSXin Li    logger.GetLogger().LogOutput("Trying to disable beeping.")
67*760c253cSXin Li
68*760c253cSXin Li    ret, o, _ = cmd_executer.CrosRunCommandWOutput(
69*760c253cSXin Li        command, chromeos_root=chromeos_root, machine=remote
70*760c253cSXin Li    )
71*760c253cSXin Li    if ret != 0:
72*760c253cSXin Li        logger.GetLogger().LogOutput(o)
73*760c253cSXin Li        logger.GetLogger().LogOutput("Failed to disable beeps.")
74*760c253cSXin Li
75*760c253cSXin Li
76*760c253cSXin Lidef FindChromeOSImage(image_file, chromeos_root):
77*760c253cSXin Li    """Find path for ChromeOS image inside chroot.
78*760c253cSXin Li
79*760c253cSXin Li    This function could be called with image paths that are either inside
80*760c253cSXin Li    or outside the chroot.  In either case the path needs to be translated
81*760c253cSXin Li    to an real/absolute path inside the chroot.
82*760c253cSXin Li    Example input paths:
83*760c253cSXin Li    /usr/local/google/home/uname/chromeos/out/tmp/my-test-images/image
84*760c253cSXin Li    ~/chromiumos/src/build/images/board/latest/image
85*760c253cSXin Li    /tmp/peppy-release/R67-1235.0.0/image
86*760c253cSXin Li
87*760c253cSXin Li    Corresponding example output paths:
88*760c253cSXin Li    /tmp/my-test-images/image
89*760c253cSXin Li    /mnt/host/source/src/build/images/board/latest/image
90*760c253cSXin Li    /tmp/peppy-release/R67-1235.0,0/image
91*760c253cSXin Li    """
92*760c253cSXin Li
93*760c253cSXin Li    sys.path.insert(0, chromeos_root)
94*760c253cSXin Li
95*760c253cSXin Li    from chromite.lib import path_util
96*760c253cSXin Li
97*760c253cSXin Li    return path_util.ToChrootPath(image_file, source_path=chromeos_root)
98*760c253cSXin Li
99*760c253cSXin Li
100*760c253cSXin Lidef DoImage(argv):
101*760c253cSXin Li    """Image ChromeOS."""
102*760c253cSXin Li
103*760c253cSXin Li    parser = argparse.ArgumentParser()
104*760c253cSXin Li    parser.add_argument(
105*760c253cSXin Li        "-c",
106*760c253cSXin Li        "--chromeos_root",
107*760c253cSXin Li        dest="chromeos_root",
108*760c253cSXin Li        help="Target directory for ChromeOS installation.",
109*760c253cSXin Li    )
110*760c253cSXin Li    parser.add_argument("-r", "--remote", dest="remote", help="Target device.")
111*760c253cSXin Li    parser.add_argument(
112*760c253cSXin Li        "-i", "--image", dest="image", help="Image binary file."
113*760c253cSXin Li    )
114*760c253cSXin Li    parser.add_argument(
115*760c253cSXin Li        "-b", "--board", dest="board", help="Target board override."
116*760c253cSXin Li    )
117*760c253cSXin Li    parser.add_argument(
118*760c253cSXin Li        "-f",
119*760c253cSXin Li        "--force",
120*760c253cSXin Li        dest="force",
121*760c253cSXin Li        action="store_true",
122*760c253cSXin Li        default=False,
123*760c253cSXin Li        help="Force an image even if it is non-test.",
124*760c253cSXin Li    )
125*760c253cSXin Li    parser.add_argument(
126*760c253cSXin Li        "-n",
127*760c253cSXin Li        "--no_lock",
128*760c253cSXin Li        dest="no_lock",
129*760c253cSXin Li        default=False,
130*760c253cSXin Li        action="store_true",
131*760c253cSXin Li        help="Do not attempt to lock remote before imaging.  "
132*760c253cSXin Li        "This option should only be used in cases where the "
133*760c253cSXin Li        "exclusive lock has already been acquired (e.g. in "
134*760c253cSXin Li        "a script that calls this one).",
135*760c253cSXin Li    )
136*760c253cSXin Li    parser.add_argument(
137*760c253cSXin Li        "-l",
138*760c253cSXin Li        "--logging_level",
139*760c253cSXin Li        dest="log_level",
140*760c253cSXin Li        default="verbose",
141*760c253cSXin Li        help="Amount of logging to be used. Valid levels are "
142*760c253cSXin Li        "'quiet', 'average', and 'verbose'.",
143*760c253cSXin Li    )
144*760c253cSXin Li    parser.add_argument("-a", "--image_args", dest="image_args")
145*760c253cSXin Li    parser.add_argument(
146*760c253cSXin Li        "--keep_stateful",
147*760c253cSXin Li        dest="keep_stateful",
148*760c253cSXin Li        default=False,
149*760c253cSXin Li        action="store_true",
150*760c253cSXin Li        help="Do not clobber the stateful partition.",
151*760c253cSXin Li    )
152*760c253cSXin Li
153*760c253cSXin Li    options = parser.parse_args(argv[1:])
154*760c253cSXin Li
155*760c253cSXin Li    if not options.log_level in command_executer.LOG_LEVEL:
156*760c253cSXin Li        Usage(parser, "--logging_level must be 'quiet', 'average' or 'verbose'")
157*760c253cSXin Li    else:
158*760c253cSXin Li        log_level = options.log_level
159*760c253cSXin Li
160*760c253cSXin Li    # Common initializations
161*760c253cSXin Li    cmd_executer = command_executer.GetCommandExecuter(log_level=log_level)
162*760c253cSXin Li    l = logger.GetLogger()
163*760c253cSXin Li
164*760c253cSXin Li    if options.chromeos_root is None:
165*760c253cSXin Li        Usage(parser, "--chromeos_root must be set")
166*760c253cSXin Li
167*760c253cSXin Li    if options.remote is None:
168*760c253cSXin Li        Usage(parser, "--remote must be set")
169*760c253cSXin Li
170*760c253cSXin Li    options.chromeos_root = os.path.expanduser(options.chromeos_root)
171*760c253cSXin Li
172*760c253cSXin Li    if options.board is None:
173*760c253cSXin Li        board = cmd_executer.CrosLearnBoard(
174*760c253cSXin Li            options.chromeos_root, options.remote
175*760c253cSXin Li        )
176*760c253cSXin Li    else:
177*760c253cSXin Li        board = options.board
178*760c253cSXin Li
179*760c253cSXin Li    if options.image is None:
180*760c253cSXin Li        images_dir = misc.GetImageDir(options.chromeos_root, board)
181*760c253cSXin Li        image = os.path.join(images_dir, "latest", "chromiumos_test_image.bin")
182*760c253cSXin Li        if not os.path.exists(image):
183*760c253cSXin Li            image = os.path.join(images_dir, "latest", "chromiumos_image.bin")
184*760c253cSXin Li        is_xbuddy_image = False
185*760c253cSXin Li    else:
186*760c253cSXin Li        image = options.image
187*760c253cSXin Li        is_xbuddy_image = image.startswith("xbuddy://")
188*760c253cSXin Li        if not is_xbuddy_image:
189*760c253cSXin Li            image = os.path.expanduser(image)
190*760c253cSXin Li
191*760c253cSXin Li    if not is_xbuddy_image:
192*760c253cSXin Li        image = os.path.realpath(image)
193*760c253cSXin Li
194*760c253cSXin Li    if not os.path.exists(image) and not is_xbuddy_image:
195*760c253cSXin Li        Usage(parser, "Image file: " + image + " does not exist!")
196*760c253cSXin Li
197*760c253cSXin Li    try:
198*760c253cSXin Li        should_unlock = False
199*760c253cSXin Li        if not options.no_lock:
200*760c253cSXin Li            try:
201*760c253cSXin Li                _ = locks.AcquireLock(
202*760c253cSXin Li                    list(options.remote.split()), options.chromeos_root
203*760c253cSXin Li                )
204*760c253cSXin Li                should_unlock = True
205*760c253cSXin Li            except Exception as e:
206*760c253cSXin Li                raise RuntimeError("Error acquiring machine: %s" % str(e))
207*760c253cSXin Li
208*760c253cSXin Li        reimage = False
209*760c253cSXin Li        local_image = False
210*760c253cSXin Li        if not is_xbuddy_image:
211*760c253cSXin Li            local_image = True
212*760c253cSXin Li            image_checksum = FileUtils().Md5File(image, log_level=log_level)
213*760c253cSXin Li
214*760c253cSXin Li            command = "cat " + checksum_file
215*760c253cSXin Li            ret, device_checksum, _ = cmd_executer.CrosRunCommandWOutput(
216*760c253cSXin Li                command,
217*760c253cSXin Li                chromeos_root=options.chromeos_root,
218*760c253cSXin Li                machine=options.remote,
219*760c253cSXin Li            )
220*760c253cSXin Li
221*760c253cSXin Li            device_checksum = device_checksum.strip()
222*760c253cSXin Li            image_checksum = str(image_checksum)
223*760c253cSXin Li
224*760c253cSXin Li            l.LogOutput("Image checksum: " + image_checksum)
225*760c253cSXin Li            l.LogOutput("Device checksum: " + device_checksum)
226*760c253cSXin Li
227*760c253cSXin Li            if image_checksum != device_checksum:
228*760c253cSXin Li                [found, located_image] = LocateOrCopyImage(
229*760c253cSXin Li                    options.chromeos_root, image, board=board
230*760c253cSXin Li                )
231*760c253cSXin Li
232*760c253cSXin Li                reimage = True
233*760c253cSXin Li                l.LogOutput("Checksums do not match. Re-imaging...")
234*760c253cSXin Li
235*760c253cSXin Li                chroot_image = FindChromeOSImage(
236*760c253cSXin Li                    located_image, options.chromeos_root
237*760c253cSXin Li                )
238*760c253cSXin Li
239*760c253cSXin Li                is_test_image = IsImageModdedForTest(
240*760c253cSXin Li                    options.chromeos_root, chroot_image, log_level
241*760c253cSXin Li                )
242*760c253cSXin Li
243*760c253cSXin Li                if not is_test_image and not options.force:
244*760c253cSXin Li                    logger.GetLogger().LogFatal(
245*760c253cSXin Li                        "Have to pass --force to image a " "non-test image!"
246*760c253cSXin Li                    )
247*760c253cSXin Li        else:
248*760c253cSXin Li            reimage = True
249*760c253cSXin Li            found = True
250*760c253cSXin Li            l.LogOutput("Using non-local image; Re-imaging...")
251*760c253cSXin Li
252*760c253cSXin Li        if reimage:
253*760c253cSXin Li            # If the device has /tmp mounted as noexec, image_to_live.sh can fail.
254*760c253cSXin Li            command = "mount -o remount,rw,exec /tmp"
255*760c253cSXin Li            cmd_executer.CrosRunCommand(
256*760c253cSXin Li                command,
257*760c253cSXin Li                chromeos_root=options.chromeos_root,
258*760c253cSXin Li                machine=options.remote,
259*760c253cSXin Li            )
260*760c253cSXin Li
261*760c253cSXin Li            # Check to see if cros flash will work for the remote machine.
262*760c253cSXin Li            CheckForCrosFlash(options.chromeos_root, options.remote, log_level)
263*760c253cSXin Li
264*760c253cSXin Li            # Disable the annoying chromebook beeps after reboot.
265*760c253cSXin Li            DisableCrosBeeps(options.chromeos_root, options.remote, log_level)
266*760c253cSXin Li
267*760c253cSXin Li            cros_flash_args = [
268*760c253cSXin Li                "cros",
269*760c253cSXin Li                "flash",
270*760c253cSXin Li                "--board=%s" % board,
271*760c253cSXin Li            ]
272*760c253cSXin Li            if not options.keep_stateful:
273*760c253cSXin Li                cros_flash_args.append("--clobber-stateful")
274*760c253cSXin Li            # New arguments should be added here.
275*760c253cSXin Li
276*760c253cSXin Li            # The last two arguments are positional and have to be at the end.
277*760c253cSXin Li            cros_flash_args.append(options.remote)
278*760c253cSXin Li            if local_image:
279*760c253cSXin Li                cros_flash_args.append(chroot_image)
280*760c253cSXin Li            else:
281*760c253cSXin Li                cros_flash_args.append(image)
282*760c253cSXin Li
283*760c253cSXin Li            command = " ".join(cros_flash_args)
284*760c253cSXin Li
285*760c253cSXin Li            # Workaround for crosbug.com/35684.
286*760c253cSXin Li            os.chmod(misc.GetChromeOSKeyFile(options.chromeos_root), 0o600)
287*760c253cSXin Li
288*760c253cSXin Li            if log_level == "average":
289*760c253cSXin Li                cmd_executer.SetLogLevel("verbose")
290*760c253cSXin Li            retries = 0
291*760c253cSXin Li            while True:
292*760c253cSXin Li                if log_level == "quiet":
293*760c253cSXin Li                    l.LogOutput("CMD : %s" % command)
294*760c253cSXin Li                ret = cmd_executer.ChrootRunCommand(
295*760c253cSXin Li                    options.chromeos_root, command, command_timeout=1800
296*760c253cSXin Li                )
297*760c253cSXin Li                if ret == 0 or retries >= 2:
298*760c253cSXin Li                    break
299*760c253cSXin Li                retries += 1
300*760c253cSXin Li                if log_level == "quiet":
301*760c253cSXin Li                    l.LogOutput("Imaging failed. Retry # %d." % retries)
302*760c253cSXin Li
303*760c253cSXin Li            if log_level == "average":
304*760c253cSXin Li                cmd_executer.SetLogLevel(log_level)
305*760c253cSXin Li
306*760c253cSXin Li            logger.GetLogger().LogFatalIf(ret, "Image command failed")
307*760c253cSXin Li
308*760c253cSXin Li            # Unfortunately cros_image_to_target.py sometimes returns early when the
309*760c253cSXin Li            # machine isn't fully up yet.
310*760c253cSXin Li            ret = EnsureMachineUp(
311*760c253cSXin Li                options.chromeos_root, options.remote, log_level
312*760c253cSXin Li            )
313*760c253cSXin Li
314*760c253cSXin Li            # If this is a non-local image, then the ret returned from
315*760c253cSXin Li            # EnsureMachineUp is the one that will be returned by this function;
316*760c253cSXin Li            # in that case, make sure the value in 'ret' is appropriate.
317*760c253cSXin Li            if not local_image and ret:
318*760c253cSXin Li                ret = 0
319*760c253cSXin Li            else:
320*760c253cSXin Li                ret = 1
321*760c253cSXin Li
322*760c253cSXin Li            if local_image:
323*760c253cSXin Li                if log_level == "average":
324*760c253cSXin Li                    l.LogOutput("Verifying image.")
325*760c253cSXin Li                command = "echo %s > %s && chmod -w %s" % (
326*760c253cSXin Li                    image_checksum,
327*760c253cSXin Li                    checksum_file,
328*760c253cSXin Li                    checksum_file,
329*760c253cSXin Li                )
330*760c253cSXin Li                ret = cmd_executer.CrosRunCommand(
331*760c253cSXin Li                    command,
332*760c253cSXin Li                    chromeos_root=options.chromeos_root,
333*760c253cSXin Li                    machine=options.remote,
334*760c253cSXin Li                )
335*760c253cSXin Li                logger.GetLogger().LogFatalIf(ret, "Writing checksum failed.")
336*760c253cSXin Li
337*760c253cSXin Li                successfully_imaged = VerifyChromeChecksum(
338*760c253cSXin Li                    options.chromeos_root,
339*760c253cSXin Li                    chroot_image,
340*760c253cSXin Li                    options.remote,
341*760c253cSXin Li                    log_level,
342*760c253cSXin Li                )
343*760c253cSXin Li                logger.GetLogger().LogFatalIf(
344*760c253cSXin Li                    not successfully_imaged, "Image verification failed!"
345*760c253cSXin Li                )
346*760c253cSXin Li                TryRemountPartitionAsRW(
347*760c253cSXin Li                    options.chromeos_root, options.remote, log_level
348*760c253cSXin Li                )
349*760c253cSXin Li
350*760c253cSXin Li            if not found:
351*760c253cSXin Li                temp_dir = os.path.dirname(located_image)
352*760c253cSXin Li                l.LogOutput("Deleting temp image dir: %s" % temp_dir)
353*760c253cSXin Li                shutil.rmtree(temp_dir)
354*760c253cSXin Li            l.LogOutput("Image updated.")
355*760c253cSXin Li        else:
356*760c253cSXin Li            l.LogOutput("Checksums match, skip image update and reboot.")
357*760c253cSXin Li            command = "reboot && exit"
358*760c253cSXin Li            _ = cmd_executer.CrosRunCommand(
359*760c253cSXin Li                command,
360*760c253cSXin Li                chromeos_root=options.chromeos_root,
361*760c253cSXin Li                machine=options.remote,
362*760c253cSXin Li            )
363*760c253cSXin Li            # Wait 30s after reboot.
364*760c253cSXin Li            time.sleep(30)
365*760c253cSXin Li
366*760c253cSXin Li    finally:
367*760c253cSXin Li        if should_unlock:
368*760c253cSXin Li            locks.ReleaseLock(
369*760c253cSXin Li                list(options.remote.split()), options.chromeos_root
370*760c253cSXin Li            )
371*760c253cSXin Li
372*760c253cSXin Li    return ret
373*760c253cSXin Li
374*760c253cSXin Li
375*760c253cSXin Lidef LocateOrCopyImage(chromeos_root, image, board=None):
376*760c253cSXin Li    l = logger.GetLogger()
377*760c253cSXin Li    if board is None:
378*760c253cSXin Li        board_glob = "*"
379*760c253cSXin Li    else:
380*760c253cSXin Li        board_glob = board
381*760c253cSXin Li
382*760c253cSXin Li    chromeos_root_realpath = os.path.realpath(chromeos_root)
383*760c253cSXin Li    image = os.path.realpath(image)
384*760c253cSXin Li
385*760c253cSXin Li    if image.startswith("%s/" % chromeos_root_realpath):
386*760c253cSXin Li        return [True, image]
387*760c253cSXin Li
388*760c253cSXin Li    # First search within the existing build dirs for any matching files.
389*760c253cSXin Li    images_glob = "%s/src/build/images/%s/*/*.bin" % (
390*760c253cSXin Li        chromeos_root_realpath,
391*760c253cSXin Li        board_glob,
392*760c253cSXin Li    )
393*760c253cSXin Li    images_list = glob.glob(images_glob)
394*760c253cSXin Li    for potential_image in images_list:
395*760c253cSXin Li        if filecmp.cmp(potential_image, image):
396*760c253cSXin Li            l.LogOutput(
397*760c253cSXin Li                "Found matching image %s in chromeos_root." % potential_image
398*760c253cSXin Li            )
399*760c253cSXin Li            return [True, potential_image]
400*760c253cSXin Li    # We did not find an image. Copy it in the src dir and return the copied
401*760c253cSXin Li    # file.
402*760c253cSXin Li    if board is None:
403*760c253cSXin Li        board = ""
404*760c253cSXin Li    base_dir = "%s/src/build/images/%s" % (chromeos_root_realpath, board)
405*760c253cSXin Li    if not os.path.isdir(base_dir):
406*760c253cSXin Li        os.makedirs(base_dir)
407*760c253cSXin Li    temp_dir = tempfile.mkdtemp(prefix="%s/tmp" % base_dir)
408*760c253cSXin Li    new_image = "%s/%s" % (temp_dir, os.path.basename(image))
409*760c253cSXin Li    l.LogOutput(
410*760c253cSXin Li        "No matching image found. Copying %s to %s" % (image, new_image)
411*760c253cSXin Li    )
412*760c253cSXin Li    shutil.copyfile(image, new_image)
413*760c253cSXin Li    return [False, new_image]
414*760c253cSXin Li
415*760c253cSXin Li
416*760c253cSXin Lidef GetImageMountCommand(image, rootfs_mp, stateful_mp):
417*760c253cSXin Li    image_dir = os.path.dirname(image)
418*760c253cSXin Li    image_file = os.path.basename(image)
419*760c253cSXin Li    mount_command = (
420*760c253cSXin Li        "cd /mnt/host/source/src/scripts &&"
421*760c253cSXin Li        "./mount_gpt_image.sh --from=%s --image=%s"
422*760c253cSXin Li        " --safe --read_only"
423*760c253cSXin Li        " --rootfs_mountpt=%s"
424*760c253cSXin Li        " --stateful_mountpt=%s"
425*760c253cSXin Li        % (image_dir, image_file, rootfs_mp, stateful_mp)
426*760c253cSXin Li    )
427*760c253cSXin Li    return mount_command
428*760c253cSXin Li
429*760c253cSXin Li
430*760c253cSXin Lidef MountImage(
431*760c253cSXin Li    chromeos_root,
432*760c253cSXin Li    image,
433*760c253cSXin Li    rootfs_mp,
434*760c253cSXin Li    stateful_mp,
435*760c253cSXin Li    log_level,
436*760c253cSXin Li    unmount=False,
437*760c253cSXin Li    extra_commands="",
438*760c253cSXin Li):
439*760c253cSXin Li    cmd_executer = command_executer.GetCommandExecuter(log_level=log_level)
440*760c253cSXin Li    command = GetImageMountCommand(image, rootfs_mp, stateful_mp)
441*760c253cSXin Li    if unmount:
442*760c253cSXin Li        command = "%s --unmount" % command
443*760c253cSXin Li    if extra_commands:
444*760c253cSXin Li        command = "%s ; %s" % (command, extra_commands)
445*760c253cSXin Li    ret, out, _ = cmd_executer.ChrootRunCommandWOutput(chromeos_root, command)
446*760c253cSXin Li    logger.GetLogger().LogFatalIf(ret, "Mount/unmount command failed!")
447*760c253cSXin Li    return out
448*760c253cSXin Li
449*760c253cSXin Li
450*760c253cSXin Lidef IsImageModdedForTest(chromeos_root, image, log_level):
451*760c253cSXin Li    if log_level != "verbose":
452*760c253cSXin Li        log_level = "quiet"
453*760c253cSXin Li    command = "mktemp -d"
454*760c253cSXin Li    cmd_executer = command_executer.GetCommandExecuter(log_level=log_level)
455*760c253cSXin Li    _, rootfs_mp, _ = cmd_executer.ChrootRunCommandWOutput(
456*760c253cSXin Li        chromeos_root, command
457*760c253cSXin Li    )
458*760c253cSXin Li    _, stateful_mp, _ = cmd_executer.ChrootRunCommandWOutput(
459*760c253cSXin Li        chromeos_root, command
460*760c253cSXin Li    )
461*760c253cSXin Li    rootfs_mp = rootfs_mp.strip()
462*760c253cSXin Li    stateful_mp = stateful_mp.strip()
463*760c253cSXin Li    lsb_release_file = os.path.join(rootfs_mp, "etc/lsb-release")
464*760c253cSXin Li    extra = "grep CHROMEOS_RELEASE_TRACK %s | grep -i test" % lsb_release_file
465*760c253cSXin Li    output = MountImage(
466*760c253cSXin Li        chromeos_root,
467*760c253cSXin Li        image,
468*760c253cSXin Li        rootfs_mp,
469*760c253cSXin Li        stateful_mp,
470*760c253cSXin Li        log_level,
471*760c253cSXin Li        extra_commands=extra,
472*760c253cSXin Li    )
473*760c253cSXin Li    is_test_image = re.search("test", output, re.IGNORECASE)
474*760c253cSXin Li    MountImage(
475*760c253cSXin Li        chromeos_root, image, rootfs_mp, stateful_mp, log_level, unmount=True
476*760c253cSXin Li    )
477*760c253cSXin Li    return is_test_image
478*760c253cSXin Li
479*760c253cSXin Li
480*760c253cSXin Lidef VerifyChromeChecksum(chromeos_root, image, remote, log_level):
481*760c253cSXin Li    command = "mktemp -d"
482*760c253cSXin Li    cmd_executer = command_executer.GetCommandExecuter(log_level=log_level)
483*760c253cSXin Li    _, rootfs_mp, _ = cmd_executer.ChrootRunCommandWOutput(
484*760c253cSXin Li        chromeos_root, command
485*760c253cSXin Li    )
486*760c253cSXin Li    _, stateful_mp, _ = cmd_executer.ChrootRunCommandWOutput(
487*760c253cSXin Li        chromeos_root, command
488*760c253cSXin Li    )
489*760c253cSXin Li    rootfs_mp = rootfs_mp.strip()
490*760c253cSXin Li    stateful_mp = stateful_mp.strip()
491*760c253cSXin Li    chrome_file = "%s/opt/google/chrome/chrome" % rootfs_mp
492*760c253cSXin Li    extra = "md5sum %s" % chrome_file
493*760c253cSXin Li    out = MountImage(
494*760c253cSXin Li        chromeos_root,
495*760c253cSXin Li        image,
496*760c253cSXin Li        rootfs_mp,
497*760c253cSXin Li        stateful_mp,
498*760c253cSXin Li        log_level,
499*760c253cSXin Li        extra_commands=extra,
500*760c253cSXin Li    )
501*760c253cSXin Li    image_chrome_checksum = out.strip().split()[0]
502*760c253cSXin Li    MountImage(
503*760c253cSXin Li        chromeos_root, image, rootfs_mp, stateful_mp, log_level, unmount=True
504*760c253cSXin Li    )
505*760c253cSXin Li
506*760c253cSXin Li    command = "md5sum /opt/google/chrome/chrome"
507*760c253cSXin Li    [_, o, _] = cmd_executer.CrosRunCommandWOutput(
508*760c253cSXin Li        command, chromeos_root=chromeos_root, machine=remote
509*760c253cSXin Li    )
510*760c253cSXin Li    device_chrome_checksum = o.split()[0]
511*760c253cSXin Li    return image_chrome_checksum.strip() == device_chrome_checksum.strip()
512*760c253cSXin Li
513*760c253cSXin Li
514*760c253cSXin Li# Remount partition as writable.
515*760c253cSXin Li# TODO: auto-detect if an image is built using --noenable_rootfs_verification.
516*760c253cSXin Lidef TryRemountPartitionAsRW(chromeos_root, remote, log_level):
517*760c253cSXin Li    l = logger.GetLogger()
518*760c253cSXin Li    cmd_executer = command_executer.GetCommandExecuter(log_level=log_level)
519*760c253cSXin Li    command = "sudo mount -o remount,rw /"
520*760c253cSXin Li    ret = cmd_executer.CrosRunCommand(
521*760c253cSXin Li        command,
522*760c253cSXin Li        chromeos_root=chromeos_root,
523*760c253cSXin Li        machine=remote,
524*760c253cSXin Li        terminated_timeout=10,
525*760c253cSXin Li    )
526*760c253cSXin Li    if ret:
527*760c253cSXin Li        ## Safely ignore.
528*760c253cSXin Li        l.LogWarning(
529*760c253cSXin Li            "Failed to remount partition as rw, "
530*760c253cSXin Li            "probably the image was not built with "
531*760c253cSXin Li            '"--noenable_rootfs_verification", '
532*760c253cSXin Li            "you can safely ignore this."
533*760c253cSXin Li        )
534*760c253cSXin Li    else:
535*760c253cSXin Li        l.LogOutput("Re-mounted partition as writable.")
536*760c253cSXin Li
537*760c253cSXin Li
538*760c253cSXin Lidef EnsureMachineUp(chromeos_root, remote, log_level):
539*760c253cSXin Li    l = logger.GetLogger()
540*760c253cSXin Li    cmd_executer = command_executer.GetCommandExecuter(log_level=log_level)
541*760c253cSXin Li    timeout = 600
542*760c253cSXin Li    magic = "abcdefghijklmnopqrstuvwxyz"
543*760c253cSXin Li    command = "echo %s" % magic
544*760c253cSXin Li    start_time = time.time()
545*760c253cSXin Li    while True:
546*760c253cSXin Li        current_time = time.time()
547*760c253cSXin Li        if current_time - start_time > timeout:
548*760c253cSXin Li            l.LogError(
549*760c253cSXin Li                "Timeout of %ss reached. Machine still not up. Aborting."
550*760c253cSXin Li                % timeout
551*760c253cSXin Li            )
552*760c253cSXin Li            return False
553*760c253cSXin Li        ret = cmd_executer.CrosRunCommand(
554*760c253cSXin Li            command, chromeos_root=chromeos_root, machine=remote
555*760c253cSXin Li        )
556*760c253cSXin Li        if not ret:
557*760c253cSXin Li            return True
558*760c253cSXin Li
559*760c253cSXin Li
560*760c253cSXin Liif __name__ == "__main__":
561*760c253cSXin Li    retval = DoImage(sys.argv)
562*760c253cSXin Li    sys.exit(retval)
563