From 1c830bfbaa92f23bde31e6b40e72c13e61085ba7 Mon Sep 17 00:00:00 2001 From: Tao Bao Date: Mon, 25 Dec 2017 10:43:47 -0800 Subject: [PATCH] releasetools: Add common.MakeTempDir(). This simplifies the work to track mkdtemp()'d dirs. This CL also cleans up the temporary file creation in build_image.py. Test: `m dist` Test: python -m unittest test_common Change-Id: Ied7d86126714b9b26e5df7bcca1dfdea21cfcf60 --- tools/releasetools/add_img_to_target_files.py | 10 ++--- tools/releasetools/build_image.py | 45 +++++++------------ tools/releasetools/common.py | 29 ++++++++---- tools/releasetools/test_common.py | 18 ++++---- 4 files changed, 47 insertions(+), 55 deletions(-) diff --git a/tools/releasetools/add_img_to_target_files.py b/tools/releasetools/add_img_to_target_files.py index 8b55a45cc..9601d881f 100755 --- a/tools/releasetools/add_img_to_target_files.py +++ b/tools/releasetools/add_img_to_target_files.py @@ -307,8 +307,7 @@ def AddUserdata(output_zip, prefix="IMAGES/"): if OPTIONS.info_dict.get("userdata_img_with_data") == "true": user_dir = os.path.join(OPTIONS.input_tmp, "DATA") else: - user_dir = tempfile.mkdtemp() - OPTIONS.tempfiles.append(user_dir) + user_dir = common.MakeTempDir() fstab = OPTIONS.info_dict["fstab"] if fstab: @@ -363,9 +362,7 @@ def AddVBMeta(output_zip, partitions, prefix="IMAGES/"): cmd = [avbtool, "make_vbmeta_image", "--output", img.name] common.AppendAVBSigningArgs(cmd, "vbmeta") - public_key_dir = tempfile.mkdtemp(prefix="avbpubkey-") - OPTIONS.tempfiles.append(public_key_dir) - + public_key_dir = common.MakeTempDir(prefix="avbpubkey-") for partition, path in partitions.items(): assert partition in common.AVB_PARTITIONS, 'Unknown partition: %s' % ( partition,) @@ -453,8 +450,7 @@ def AddCache(output_zip, prefix="IMAGES/"): timestamp = (datetime.datetime(2009, 1, 1) - epoch).total_seconds() image_props["timestamp"] = int(timestamp) - user_dir = tempfile.mkdtemp() - OPTIONS.tempfiles.append(user_dir) + user_dir = common.MakeTempDir() fstab = OPTIONS.info_dict["fstab"] if fstab: diff --git a/tools/releasetools/build_image.py b/tools/releasetools/build_image.py index 2a92d86e5..11a0055a6 100755 --- a/tools/releasetools/build_image.py +++ b/tools/releasetools/build_image.py @@ -323,7 +323,7 @@ def MakeVerityEnabledImage(out_file, fec_supported, prop_dict): signer_args = OPTIONS.verity_signer_args # make a tempdir - tempdir_name = tempfile.mkdtemp(suffix="_verity_images") + tempdir_name = common.MakeTempDir(suffix="_verity_images") # get partial image paths verity_image_path = os.path.join(tempdir_name, "verity.img") @@ -332,7 +332,6 @@ def MakeVerityEnabledImage(out_file, fec_supported, prop_dict): # build the verity tree and get the root hash and salt if not BuildVerityTree(out_file, verity_image_path, prop_dict): - shutil.rmtree(tempdir_name, ignore_errors=True) return False # build the metadata blocks @@ -342,7 +341,6 @@ def MakeVerityEnabledImage(out_file, fec_supported, prop_dict): if not BuildVerityMetadata(image_size, verity_metadata_path, root_hash, salt, block_dev, signer_path, signer_key, signer_args, verity_disable): - shutil.rmtree(tempdir_name, ignore_errors=True) return False # build the full verified image @@ -358,21 +356,16 @@ def MakeVerityEnabledImage(out_file, fec_supported, prop_dict): verity_fec_path, padding_size, fec_supported): - shutil.rmtree(tempdir_name, ignore_errors=True) return False - shutil.rmtree(tempdir_name, ignore_errors=True) return True def ConvertBlockMapToBaseFs(block_map_file): - fd, base_fs_file = tempfile.mkstemp(prefix="script_gen_", - suffix=".base_fs") - os.close(fd) + base_fs_file = common.MakeTempFile(prefix="script_gen_", suffix=".base_fs") convert_command = ["blk_alloc_to_base_fs", block_map_file, base_fs_file] (_, exit_code) = RunCommand(convert_command) if exit_code != 0: - os.remove(base_fs_file) return None return base_fs_file @@ -426,17 +419,15 @@ def BuildImage(in_dir, prop_dict, out_file, target_out=None): # /system and the ramdisk, and can be mounted at the root of the file system. origin_in = in_dir fs_config = prop_dict.get("fs_config") - base_fs_file = None if (prop_dict.get("system_root_image") == "true" and prop_dict["mount_point"] == "system"): - in_dir = tempfile.mkdtemp() + in_dir = common.MakeTempDir() # Change the mount point to "/" prop_dict["mount_point"] = "/" if fs_config: # We need to merge the fs_config files of system and ramdisk. - fd, merged_fs_config = tempfile.mkstemp(prefix="root_fs_config", - suffix=".txt") - os.close(fd) + merged_fs_config = common.MakeTempFile(prefix="root_fs_config", + suffix=".txt") with open(merged_fs_config, "w") as fw: if "ramdisk_fs_config" in prop_dict: with open(prop_dict["ramdisk_fs_config"]) as fr: @@ -577,19 +568,10 @@ def BuildImage(in_dir, prop_dict, out_file, target_out=None): shutil.copytree(origin_in, staging_system, symlinks=True) ext4fs_output = None - try: - if fs_type.startswith("ext4"): - (ext4fs_output, exit_code) = RunCommand(build_command) - else: - (_, exit_code) = RunCommand(build_command) - finally: - if in_dir != origin_in: - # Clean up temporary directories and files. - shutil.rmtree(in_dir, ignore_errors=True) - if fs_config: - os.remove(fs_config) - if base_fs_file is not None: - os.remove(base_fs_file) + if fs_type.startswith("ext4"): + (ext4fs_output, exit_code) = RunCommand(build_command) + else: + (_, exit_code) = RunCommand(build_command) if exit_code != 0: print("Error: '%s' failed with exit code %d" % (build_command, exit_code)) return False @@ -808,15 +790,18 @@ def main(argv): mount_point = "oem" else: print >> sys.stderr, "error: unknown image file name ", image_filename - exit(1) + sys.exit(1) image_properties = ImagePropFromGlobalDict(glob_dict, mount_point) if not BuildImage(in_dir, image_properties, out_file, target_out): print >> sys.stderr, "error: failed to build %s from %s" % (out_file, in_dir) - exit(1) + sys.exit(1) if __name__ == '__main__': - main(sys.argv[1:]) + try: + main(sys.argv[1:]) + finally: + common.Cleanup() diff --git a/tools/releasetools/common.py b/tools/releasetools/common.py index 12e757dab..03e808f95 100644 --- a/tools/releasetools/common.py +++ b/tools/releasetools/common.py @@ -574,18 +574,16 @@ def Gunzip(in_filename, out_filename): def UnzipTemp(filename, pattern=None): - """Unzip the given archive into a temporary directory and return the name. + """Unzips the given archive into a temporary directory and returns the name. - If filename is of the form "foo.zip+bar.zip", unzip foo.zip into a - temp dir, then unzip bar.zip into that_dir/BOOTABLE_IMAGES. + If filename is of the form "foo.zip+bar.zip", unzip foo.zip into a temp dir, + then unzip bar.zip into that_dir/BOOTABLE_IMAGES. - Returns (tempdir, zipobj) where zipobj is a zipfile.ZipFile (of the - main file), open for reading. + Returns: + (tempdir, zipobj): tempdir is the name of the temprary directory; zipobj is + a zipfile.ZipFile (of the main file), open for reading. """ - tmp = tempfile.mkdtemp(prefix="targetfiles-") - OPTIONS.tempfiles.append(tmp) - def unzip_to_dir(filename, dirname): cmd = ["unzip", "-o", "-q", filename, "-d", dirname] if pattern is not None: @@ -596,6 +594,7 @@ def UnzipTemp(filename, pattern=None): raise ExternalError("failed to unzip input target-files \"%s\"" % (filename,)) + tmp = MakeTempDir(prefix="targetfiles-") m = re.match(r"^(.*[.]zip)\+(.*[.]zip)$", filename, re.IGNORECASE) if m: unzip_to_dir(m.group(1), tmp) @@ -955,12 +954,24 @@ def MakeTempFile(prefix='tmp', suffix=''): return fn +def MakeTempDir(prefix='tmp', suffix=''): + """Makes a temporary dir that will be cleaned up with a call to Cleanup(). + + Returns: + The absolute pathname of the new directory. + """ + dir_name = tempfile.mkdtemp(suffix=suffix, prefix=prefix) + OPTIONS.tempfiles.append(dir_name) + return dir_name + + def Cleanup(): for i in OPTIONS.tempfiles: if os.path.isdir(i): - shutil.rmtree(i) + shutil.rmtree(i, ignore_errors=True) else: os.remove(i) + del OPTIONS.tempfiles[:] class PasswordManager(object): diff --git a/tools/releasetools/test_common.py b/tools/releasetools/test_common.py index bb9393776..ed454ca72 100644 --- a/tools/releasetools/test_common.py +++ b/tools/releasetools/test_common.py @@ -14,12 +14,10 @@ # limitations under the License. # import os -import shutil import tempfile import time import unittest import zipfile - from hashlib import sha1 import common @@ -29,6 +27,7 @@ KiB = 1024 MiB = 1024 * KiB GiB = 1024 * MiB + def get_2gb_string(): size = int(2 * GiB + 1) block_size = 4 * KiB @@ -355,17 +354,18 @@ class CommonZipTest(unittest.TestCase): class InstallRecoveryScriptFormatTest(unittest.TestCase): - """Check the format of install-recovery.sh + """Checks the format of install-recovery.sh. - Its format should match between common.py and validate_target_files.py.""" + Its format should match between common.py and validate_target_files.py. + """ def setUp(self): - self._tempdir = tempfile.mkdtemp() + self._tempdir = common.MakeTempDir() # Create a dummy dict that contains the fstab info for boot&recovery. self._info = {"fstab" : {}} - dummy_fstab = \ - ["/dev/soc.0/by-name/boot /boot emmc defaults defaults", - "/dev/soc.0/by-name/recovery /recovery emmc defaults defaults"] + dummy_fstab = [ + "/dev/soc.0/by-name/boot /boot emmc defaults defaults", + "/dev/soc.0/by-name/recovery /recovery emmc defaults defaults"] self._info["fstab"] = common.LoadRecoveryFSTab("\n".join, 2, dummy_fstab) # Construct the gzipped recovery.img and boot.img self.recovery_data = bytearray([ @@ -414,4 +414,4 @@ class InstallRecoveryScriptFormatTest(unittest.TestCase): self._info) def tearDown(self): - shutil.rmtree(self._tempdir) + common.Cleanup()