Update prebuilt script to import framework stubs into a new format.

The 'current', and finalized "N" dirs are moving to a new format, with the
subdirectories "core", "public", "system", where all the artifacts with a
particular api visibility level are stored. For the finalized dirs, the api
txt artifacts are also stored in the appropriate subdir.

For now, both formats will co-exist to allow migrating all the tools that
consume these files to be updated incrementally, before deleting the old files.

See this doc for more details:
https://docs.google.com/document/d/10vPR5VmwwbNJ-SugrQ7idtwfCBB7zImQ9aoedmHWj14/edit

Bug: 77525052
Test: ./update_prebuilts.py -p -f 99 4707605

Change-Id: Ib1de6f72c09c6aab3b1c0214d3fb35a83d4725c9
diff --git a/update_prebuilts/update_prebuilts.py b/update_prebuilts/update_prebuilts.py
index 08db17f..9c5171f 100755
--- a/update_prebuilts/update_prebuilts.py
+++ b/update_prebuilts/update_prebuilts.py
@@ -14,6 +14,7 @@
 system_path = 'system_current'
 api_path = 'api'
 system_api_path = 'system-api'
+framework_sdk_target = 'sdk_phone_armv7-sdk_mac'
 support_dir = os.path.join(current_path, 'support')
 androidx_dir = os.path.join(current_path, 'androidx')
 extras_dir = os.path.join(current_path, 'extras')
@@ -513,6 +514,15 @@
     return artifact_path
 
 
+def fetch_artifacts(target, build_id, artifact_dict):
+    for artifact, target_path in artifact_dict.items():
+        artifact_path = fetch_artifact(target, build_id.url_id, artifact)
+        if not artifact_path:
+            return False
+        mv(artifact_path, target_path)
+    return True
+
+
 def extract_artifact(artifact_path):
     # Unzip the repo archive into a separate directory.
     repo_dir = os.path.basename(artifact_path)[:-4]
@@ -627,17 +637,37 @@
                                  os.path.join(extras_dir, 'material-design-x'), extract_res=False)
 
 
-def extract_to(zip_file, paths, filename, parent_path):
-    zip_path = next(filter(lambda path: filename in path, paths))
+def extract_to(zip_file, filename, parent_path):
+    zip_path = next(filter(lambda path: filename in path, zip_file.namelist()))
     src_path = zip_file.extract(zip_path)
     dst_path = path(parent_path, filename)
     mv(src_path, dst_path)
 
 
-def fetch_framework_artifacts(target, build_id, target_path, is_current_sdk):
-    platform = 'darwin' if 'mac' in target else 'linux'
-    artifact_path = fetch_artifact(
-        target, build_id.url_id, 'sdk-repo-%s-platforms-%s.zip' % (platform, build_id.fs_id))
+# This is a dict from an sdk level to an "artifact dict". The artifact dict
+# maps from artifact name to the respective package it stubs.
+# TODO(hansson): standardize the artifact names and remove this dict.
+sdk_artifacts_dict = {
+    'core': {
+        'core.current.stubs.jar': 'android.jar',
+    },
+    'public': {
+        'android.test.base.stubs.jar': 'android.test.base.jar',
+        'android.test.runner.stubs.jar': 'android.test.runner.jar',
+        'android.test.mock.stubs.jar': 'android.test.mock.jar',
+        'org.apache.http.legacy.jar': 'org.apache.http.legacy.jar',
+    },
+    'system': {
+        'android_system.jar': 'android.jar',
+        'android.test.mock.stubs_system.jar': 'android.test.mock.jar',
+    }
+}
+
+
+# TODO(hansson): Remove this method once the tools support the new structure.
+def update_framework_current_legacy(build_id):
+    sdk_artifact = 'sdk-repo-darwin-platforms-%s.zip' % build_id.fs_id
+    artifact_path = fetch_artifact(framework_sdk_target, build_id.url_id, sdk_artifact)
     if not artifact_path:
         return False
 
@@ -649,64 +679,56 @@
             'optional/android.test.runner.jar']
 
         for filename in filenames:
-            extract_to(zipFile, paths, filename, target_path)
+            extract_to(zipFile, filename, current_path)
 
-        if is_current_sdk:
-            # There's no system version of framework.aidl, so use the public one.
-            extract_to(zipFile, paths, 'framework.aidl', system_path)
+        # There's no system version of framework.aidl, so use the public one.
+        extract_to(zipFile, 'framework.aidl', system_path)
 
-            # We don't keep historical artifacts for these.
-            artifact_path = fetch_artifact(target, build_id.fs_id, 'core.current.stubs.jar')
+    artifact_dict = {
+        'core.current.stubs.jar': path(current_path, 'core.jar'),
+        'android_system.jar':  path(system_path, 'android.jar'),
+        'android.test.mock.stubs_system.jar': path(system_path, 'optional/android.test.mock.jar'),
+    }
+    return fetch_artifacts(framework_sdk_target, build_id, artifact_dict)
+
+
+def update_framework(build_id, sdk_dir):
+    for api_level in ['core', 'public', 'system']:
+        target_dir = path(sdk_dir, api_level)
+        artifact_to_filename = sdk_artifacts_dict[api_level]
+        artifact_to_path = {artifact: path(target_dir, filename)
+                            for (artifact, filename) in artifact_to_filename.items()}
+
+        if not fetch_artifacts(framework_sdk_target, build_id, artifact_to_path):
+            return False
+
+        if api_level == 'public':
+            # Fetch a few artifacts from the public sdk.
+            artifact = 'sdk-repo-darwin-platforms-%s.zip' % build_id.fs_id
+            artifact_path = fetch_artifact(framework_sdk_target, build_id.url_id, artifact)
             if not artifact_path:
                 return False
-            mv(artifact_path, path(current_path, 'core.jar'))
+
+            with zipfile.ZipFile(artifact_path) as zipFile:
+                for filename in ['android.jar', 'framework.aidl', 'uiautomator.jar']:
+                    extract_to(zipFile, filename, target_dir)
 
     return True
 
 
-def update_sdk_repo(target, build_id):
-    return fetch_framework_artifacts(target, build_id, current_path, is_current_sdk = True)
+def finalize_sdk(build_id, sdk_version):
+    target_finalize_dir = '%d' % sdk_version
 
-
-def update_system(target, build_id):
-    artifact_path = fetch_artifact(target, build_id.url_id, 'android_system.jar')
-    if not artifact_path:
-        return False
-
-    mv(artifact_path, path(system_path, 'android.jar'))
-
-    artifact_path = fetch_artifact(target, build_id.url_id, 'android.test.mock.stubs_system.jar')
-    if not artifact_path:
-        return False
-
-    mv(artifact_path, path(system_path, 'optional/android.test.mock.jar'))
-
-    return True
-
-
-
-def finalize_sdk(target, build_id, sdk_version):
-    target_finalize_dir = "%d" % sdk_version
-    target_finalize_system_dir = "system_%d" % sdk_version
-
-    artifact_to_path = {
-      'android_system.jar': path(target_finalize_system_dir, 'android.jar'),
-      'public_api.txt': path(api_path, "%d.txt" % sdk_version),
-      'system-api.txt': path(system_api_path, "%d.txt" % sdk_version),
+    extra_finalize_artifacts = {
+      'public_api.txt': path(target_finalize_dir, 'public/api/android.txt'),
+      'system-api.txt': path(target_finalize_dir, 'system/api/android.txt'),
     }
+    return fetch_artifacts(framework_sdk_target, build_id, extra_finalize_artifacts) \
+            and update_framework(build_id, target_finalize_dir)
 
-    for artifact, target_path in artifact_to_path.items():
-        artifact_path = fetch_artifact(target, build_id.url_id, artifact)
-        if not artifact_path:
-            return False
-        mv(artifact_path, target_path)
 
-    if not fetch_framework_artifacts(target, build_id, target_finalize_dir, is_current_sdk = False):
-      return False
-
-    copyfile(path(target_finalize_dir, 'framework.aidl'),
-             path(target_finalize_system_dir, 'framework.aidl'))
-    return True
+def update_framework_current(build_id):
+    return update_framework(build_id, current_path)
 
 
 def update_buildtools(target, arch, build_id):
@@ -884,16 +906,15 @@
             print_e('Failed to update App Toolkit, aborting...')
             sys.exit(1)
     if args.platform:
-        if update_sdk_repo('sdk_phone_armv7-sdk_mac', getBuildId(args)) \
-                and update_system('sdk_phone_armv7-sdk_mac', getBuildId(args)):
+        build_id = getBuildId(args)
+        if update_framework_current_legacy(build_id) and update_framework_current(build_id):
             components = append(components, 'platform SDK')
         else:
             print_e('Failed to update platform SDK, aborting...')
             sys.exit(1)
     if args.finalize_sdk:
-        n = args.finalize_sdk
-        if finalize_sdk('sdk_phone_armv7-sdk_mac', getBuildId(args), n):
-            subprocess.check_call(['git', 'add', "%d" % n, 'system_%d' % n])
+        if finalize_sdk(getBuildId(args), args.finalize_sdk):
+            subprocess.check_call(['git', 'add', "%d" % args.finalize_sdk])
             components = append(components, 'finalized SDK %d' % args.finalize_sdk)
         else:
             print_e('Failed to finalize SDK %d, aborting...' % args.finalize_sdk)