unpack.py: Allow custom input and output directories

Adds two command line arguments to unpack.py:
* --artifacts-dir to specify a different input directory
* --output-dir to specify an output directory

Bug: 230135749
Change-Id: Iac957c45a6ca763966d3b0d9c960de0e573818c3

1 files changed, 53 insertions, 38 deletions

diff --git a/unpack.py b/unpack.py
index 19ded66..4e5578b 100755
--- a/unpack.py
+++ b/unpack.py
@@ -3,33 +3,50 @@
 
 """Unpacks a prebuilt to look like an Android build directory."""
 
+import argparse
 import glob
 import itertools
 import os
 import shutil
 import stat
 import sys
+import tempfile
 import zipfile
 
 
+SCRIPT_DIR = os.path.abspath(os.path.dirname(__file__))
+
 def main():
     """Unpacks a prebuilt to look like an Android build directory.
 
-    Assumes that it is run from trusty/prebuilts/aosp.
-    Remove the pristine directory before use.
-
     Inputs, (specifically adb and qemu_trusty_arm64-target-files-buildid.zip)
-    must be placed in the pristine directory before running the script.
+    must be placed in the pristine directory before running the script, or
+    passed in with the --artifacts-dir argument.
     """
-    pristine_dir = "pristine"
-    adb_source_path = pristine_dir + "/adb"
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-a", "--artifacts-dir",
+        type=str,
+        default=os.path.join(SCRIPT_DIR, "pristine"),
+        help="Path to the directory containing the input artifacts",
+    )
+    parser.add_argument(
+        "-o", "--output-dir",
+        type=str,
+        default=os.path.join(SCRIPT_DIR, "android"),
+        help="Path to the output directory",
+    )
+    args = parser.parse_args()
+
+    adb_source_path = os.path.join(args.artifacts_dir, "adb")
     # Make sure the relevant files are here - currently just adb + target files.
     if not os.path.isfile(adb_source_path):
-        print("ADB missing from prebuilt")
+        print(f"ADB missing from {adb_source_path}")
         sys.exit(1)
 
     # Make sure we have exactly one target zip
-    target_files_search = pristine_dir + "/qemu_trusty_arm64-target_files-*.zip"
+    target_files_search = os.path.join(args.artifacts_dir,
+                                       "qemu_trusty_arm64-target_files-*.zip")
     target_files_zip_glob = glob.glob(target_files_search)
     if not target_files_zip_glob:
         print("No Android target_files found")
@@ -39,37 +56,35 @@ def main():
                "old inputs before updating the prebuilt.")
         sys.exit(1)
 
+    if os.path.exists(args.output_dir):
+        print(f"Output directory {args.output_dir} already exists.")
+        print("Please remove it manually before running this script.")
+        sys.exit(1)
+
     # Unpack the target_files archive
-    target_files_path = "target"
-    shutil.rmtree(target_files_path, ignore_errors=True)
-    zipfile.ZipFile(target_files_zip_glob[0]).extractall(target_files_path)
-
-    # Build our directory structure
-    base = "android"
-    # Clean the unpacked build directory if it already exists
-    shutil.rmtree(base, ignore_errors=True)
-    host_bin_path = base + "/out/host/linux-x86/bin"
-    image_path = base + "/out/target/product/trusty"
-    test_path = base + "/out/target/product/trusty/data"
-    os.makedirs(host_bin_path)
-    os.makedirs(image_path)
-    os.makedirs(test_path)
-
-    # Install adb to the host bin directory
-    shutil.copy(adb_source_path, host_bin_path)
-
-    # Install images
-    image_source_path = target_files_path + "/IMAGES"
-    shutil.move(image_source_path + "/system.img", image_path)
-    shutil.move(image_source_path + "/vendor.img", image_path)
-    shutil.move(image_source_path + "/userdata.img", image_path)
-
-    # Install test apps
-    test_source_path = target_files_path + "/DATA/nativetest64"
-    shutil.move(test_source_path, test_path)
-
-    # Clean up unpacked but unneeded files
-    shutil.rmtree(target_files_path)
+    with tempfile.TemporaryDirectory() as target_files_path:
+        zipfile.ZipFile(target_files_zip_glob[0]).extractall(target_files_path)
+
+        # Build our directory structure
+        host_bin_path = os.path.join(args.output_dir, "out", "host", "linux-x86", "bin")
+        image_path = os.path.join(args.output_dir, "out", "target", "product", "trusty")
+        test_path = os.path.join(image_path, "data")
+        os.makedirs(host_bin_path)
+        os.makedirs(image_path)
+        os.makedirs(test_path)
+
+        # Install adb to the host bin directory
+        shutil.copy(adb_source_path, host_bin_path)
+
+        # Install images
+        image_source_path = os.path.join(target_files_path, "IMAGES")
+        shutil.move(os.path.join(image_source_path, "system.img"), image_path)
+        shutil.move(os.path.join(image_source_path, "vendor.img"), image_path)
+        shutil.move(os.path.join(image_source_path, "userdata.img"), image_path)
+
+        # Install test apps
+        test_source_path = os.path.join(target_files_path, "DATA", "nativetest64")
+        shutil.move(test_source_path, test_path)
 
     # Fix up permissions
     exec_perms = (stat.S_IXUSR | stat.S_IRUSR | stat.S_IXGRP | stat.S_IRGRP