Minor updates to startup scripts

Change-Id: I397dfd95f3410d9285a50e221e560d3a0f7bcdac

tools/startup/instrument.py

diff --git a/tools/startup/instrument.py b/tools/startup/instrument.py
new file mode 100755
index 0000000..ecefe55
--- /dev/null
+++ b/tools/startup/instrument.py
@@ -0,0 +1,134 @@
+#!/usr/bin/env python3
+# Copyright (c) 2023, the R8 project authors. Please see the AUTHORS file
+# for details. All rights reserved. Use of this source code is governed by a
+# BSD-style license that can be found in the LICENSE file.
+
+import argparse
+import os
+import shutil
+import subprocess
+import sys
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import apk_masseur
+import apk_utils
+import extractmarker
+import toolhelper
+import utils
+import zip_utils
+
+def parse_options(argv):
+  result = argparse.ArgumentParser(
+      description='Instrument the dex files of a given apk to print what is '
+                  'executed.')
+  result.add_argument('--apk',
+                      help='Path to the .apk',
+                      required=True)
+  result.add_argument('--dex-files',
+                      action='append',
+                      help='Name of dex files to instrument')
+  result.add_argument('--discard',
+                      action='append',
+                      help='Name of dex files to discard')
+  result.add_argument('--out',
+                      help='Destination of resulting apk',
+                      required=True)
+  options, args = result.parse_known_args(argv)
+  return options, args
+
+def add_instrumented_dex(dex_file, instrumented_dex_index, instrumented_dir):
+  dex_name = get_dex_name(instrumented_dex_index)
+  destination = os.path.join(instrumented_dir, dex_name)
+  shutil.move(dex_file, destination)
+
+def get_dex_name(dex_index):
+  assert dex_index > 0
+  return 'classes.dex' if dex_index == 1 else ('classes%s.dex' % dex_index)
+
+def instrument_dex_file(dex_file, include_instrumentation_server, options, tmp_dir):
+  d8_cmd = [
+      'java',
+      '-cp', utils.R8_JAR,
+      '-Dcom.android.tools.r8.startup.instrumentation.instrument=1',
+      '-Dcom.android.tools.r8.startup.instrumentation.instrumentationtag=R8']
+  if not include_instrumentation_server:
+    # We avoid injecting the InstrumentationServer by specifying it should only
+    # be added if foo.bar.Baz is in the program.
+    d8_cmd.append(
+        '-Dcom.android.tools.r8.startup.instrumentation.instrumentationserversyntheticcontext=foo.bar.Baz')
+  d8_cmd.extend([
+      'com.android.tools.r8.D8',
+      '--min-api', str(apk_utils.get_min_api(options.apk)),
+      '--output', tmp_dir,
+      '--release',
+      dex_file])
+  subprocess.check_call(d8_cmd)
+  instrumented_dex_files = []
+  instrumented_dex_index = 1
+  while True:
+    instrumented_dex_name = get_dex_name(instrumented_dex_index)
+    instrumented_dex_file = os.path.join(tmp_dir, instrumented_dex_name)
+    if not os.path.exists(instrumented_dex_file):
+      break
+    instrumented_dex_files.append(instrumented_dex_file)
+    instrumented_dex_index = instrumented_dex_index + 1
+  assert len(instrumented_dex_files) > 0
+  return instrumented_dex_files
+
+def should_discard_dex_file(dex_name, options):
+  return options.discard is not None and dex_name in options.discard
+
+def should_instrument_dex_file(dex_name, options):
+  return options.dex_files is not None and dex_name in options.dex_files
+
+def main(argv):
+  options, args = parse_options(argv)
+  with utils.TempDir() as tmp_dir:
+    # Extract the dex files of the apk.
+    uninstrumented_dir = os.path.join(tmp_dir, 'uninstrumented')
+    os.mkdir(uninstrumented_dir)
+
+    dex_predicate = \
+        lambda name : name.startswith('classes') and name.endswith('.dex')
+    zip_utils.extract_all_that_matches(
+        options.apk, uninstrumented_dir, dex_predicate)
+
+    # Instrument each dex one by one.
+    instrumented_dir = os.path.join(tmp_dir, 'instrumented')
+    os.mkdir(instrumented_dir)
+
+    include_instrumentation_server = True
+    instrumented_dex_index = 1
+    uninstrumented_dex_index = 1
+    while True:
+      dex_name = get_dex_name(uninstrumented_dex_index)
+      dex_file = os.path.join(uninstrumented_dir, dex_name)
+      if not os.path.exists(dex_file):
+        break
+      if not should_discard_dex_file(dex_name, options):
+        if should_instrument_dex_file(dex_name, options):
+          with utils.TempDir() as tmp_instrumentation_dir:
+            instrumented_dex_files = \
+                instrument_dex_file(
+                    dex_file,
+                    include_instrumentation_server,
+                    options,
+                    tmp_instrumentation_dir)
+            for instrumented_dex_file in instrumented_dex_files:
+              add_instrumented_dex(
+                  instrumented_dex_file, instrumented_dex_index, instrumented_dir)
+              instrumented_dex_index = instrumented_dex_index + 1
+            include_instrumentation_server = False
+        else:
+          add_instrumented_dex(dex_file, instrumented_dex_index, instrumented_dir)
+          instrumented_dex_index = instrumented_dex_index + 1
+      uninstrumented_dex_index = uninstrumented_dex_index + 1
+
+    assert instrumented_dex_index > 1
+
+    # Masseur APK.
+    apk_masseur.masseur(options.apk, dex=instrumented_dir, out=options.out)
+
+if __name__ == '__main__':
+  sys.exit(main(sys.argv[1:]))