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

@@ -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:

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()

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):

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()