ddk_bootstrap: Initial version of init.py

* As it stands now, it calculates where to download the init_ddk script.

Bug: 328770706
Change-Id: I622d9b764c5c094e05190940f4a93101cf223f12

3 files changed, 283 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..43ae0e2
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+__pycache__/
+*.py[cod]
diff --git a/init.py b/init.py
new file mode 100644
index 0000000..ac05bfc
--- /dev/null
+++ b/init.py
@@ -0,0 +1,173 @@
+#!/usr/bin/env python3
+#
+# Copyright (C) 2024 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Downloads and run the appropriate DDK init script."""
+
+import argparse
+import io
+import json
+import logging
+import shutil
+import subprocess
+import sys
+import tempfile
+from typing import BinaryIO
+import urllib.error
+import urllib.request
+
+# Googlers: go/android-build-api-getting-started
+_API_ENDPOINT_PREFIX = (
+    "https://androidbuildinternal.googleapis.com/android/internal/build/v3/"
+)
+_ARTIFACT_URL_FMT = (
+    _API_ENDPOINT_PREFIX
+    # pylint: disable-next=line-too-long
+    + "builds/{build_id}/{build_target}/attempts/latest/artifacts/{filename}/url?redirect=true"
+)
+_BUILD_IDS_URL_FMT = (
+    _API_ENDPOINT_PREFIX
+    # pylint: disable-next=line-too-long
+    + "builds?branch={branch}&target={build_target}&buildAttemptStatus=complete&buildType=submitted&maxResults=1&fields=builds.buildId"
+)
+_DEFAULT_BUILD_TARGET = "kernel_aarch64"
+
+
+class KleafBootstrapError(RuntimeError):
+    pass
+
+
+class KleafBootstrap:
+    """Calculates the necessary work needed to setup a DDK workspace."""
+
+    def __init__(self, known_args: argparse.Namespace, unknown_args: list[str]):
+        self.branch: str | None = known_args.branch
+        self.build_id: str | None = known_args.build_id
+        self.url_fmt: str = known_args.url_fmt
+        self.build_target: str | None = known_args.build_target
+        self.unknown_args = unknown_args
+
+    def run(self):
+        if not self.branch and not self.build_id:
+            logging.error("Either --branch or --build_id must be specified")
+            sys.exit(1)
+
+        if not self.build_id:
+            self._set_build_id()
+
+        assert self.build_id, "build id is not set!"
+
+        with tempfile.NamedTemporaryFile(
+            prefix="init_ddk_", suffix=".zip", mode="w+b"
+        ) as init_ddk:
+            self._download_artifact("init_ddk.zip", init_ddk)
+
+            args = ["python3", init_ddk.name]
+            # Do not add --branch, because its meaning may change during the
+            # process of this script.
+            if self.build_id:
+                args += ["--build_id", self.build_id]
+            if self.build_target:
+                args += ["--build_target", self.build_target]
+            args += ["--url_fmt", self.url_fmt]
+            args += self.unknown_args
+            logging.debug("Running %s", args)
+            subprocess.check_call(args)
+
+    def _set_build_id(self) -> str:
+        assert self.branch, "branch is not set!"
+        build_ids_fp = io.BytesIO()
+        url = _BUILD_IDS_URL_FMT.format(
+            branch=self.branch,
+            build_target=self.build_target,
+        )
+        self._download(url, "build_id", build_ids_fp)
+        build_ids_fp.seek(0)
+        try:
+            build_ids_res = json.load(
+                io.TextIOWrapper(build_ids_fp, encoding="utf-8")
+            )
+        except json.JSONDecodeError as exc:
+            raise KleafBootstrapError(
+                "Unable to get build_id: not json") from exc
+
+        try:
+            self.build_id = build_ids_res["builds"][0]["buildId"]
+        except (KeyError, IndexError) as exc:
+            raise KleafBootstrapError(
+                "Unable to get build_id: json not in expected format") from exc
+
+        if not isinstance(self.build_id, str):
+            raise KleafBootstrapError(
+                "Unable to get build_id: json not in expected format: "
+                "build id is not string")
+
+    @staticmethod
+    def _download(url, remote_filename, out_f: BinaryIO):
+        try:
+            with urllib.request.urlopen(url) as in_f:
+                logging.debug("Scheduling download for %s", remote_filename)
+                shutil.copyfileobj(in_f, out_f)
+        except urllib.error.URLError as exc:
+            raise KleafBootstrapError(f"Fail to download {url}") from exc
+
+    def _download_artifact(self, remote_filename, out_f: BinaryIO):
+        url = self.url_fmt.format(
+            build_id=self.build_id,
+            build_target=self.build_target,
+            filename=urllib.parse.quote(remote_filename, safe=""),  # / -> %2F
+        )
+        self._download(url, remote_filename, out_f)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawTextHelpFormatter
+    )
+    parser.add_argument(
+        "--branch",
+        help=(
+            "Android Kernel branch from CI. e.g."
+            " aosp_kernel-common-android-mainline."
+        ),
+        type=str,
+        default=None,
+    )
+    parser.add_argument(
+        "--url_fmt",
+        help="URL format endpoint for CI downloads.",
+        default=_ARTIFACT_URL_FMT,
+    )
+    parser.add_argument(
+        "--build_id",
+        type=str,
+        help="the build id to download the build for, e.g. 6148204",
+    )
+    parser.add_argument(
+        "--build_target",
+        type=str,
+        help='the build target to download, e.g. "kernel_aarch64"',
+        default=_DEFAULT_BUILD_TARGET,
+    )
+    known_args, unknown_args = parser.parse_known_args()
+    logging.basicConfig(
+        level=logging.DEBUG, format="%(levelname)s: %(message)s"
+    )
+
+    try:
+        KleafBootstrap(known_args=known_args, unknown_args=unknown_args).run()
+    except KleafBootstrapError as e:
+        logging.error(e, exc_info=e)
+        sys.exit(1)
diff --git a/init_test.py b/init_test.py
new file mode 100644
index 0000000..088c45a
--- /dev/null
+++ b/init_test.py
@@ -0,0 +1,108 @@
+# Copyright (C) 2024 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for init.py"""
+
+import argparse
+import io
+import json
+import unittest
+import urllib.parse
+from typing import Generator
+
+from init import (KleafBootstrap, _API_ENDPOINT_PREFIX,
+                  _ARTIFACT_URL_FMT, _DEFAULT_BUILD_TARGET)
+
+# pylint: disable=protected-access
+
+
+def _get_branches() -> Generator[str, None, None]:
+    common_url = _API_ENDPOINT_PREFIX + "branches"
+
+    page_token = None
+    while True:
+        parsed = urllib.parse.urlparse(common_url)
+        query_dict = urllib.parse.parse_qs(parsed.query)
+        if page_token:
+            query_dict["pageToken"] = page_token
+        query_dict["fields"] = "branches.name,nextPageToken"
+        parsed = parsed._replace(
+            query=urllib.parse.urlencode(query_dict, doseq=True))
+        url = parsed.geturl()
+
+        buf = io.BytesIO()
+        KleafBootstrap._download(url, "branches", buf)
+        buf.seek(0)
+        json_obj = json.load(buf)
+        page_token = json_obj.get("nextPageToken")
+        for branch in json_obj.get("branches", []):
+            yield branch["name"]
+
+        if not page_token:
+            return
+
+
+def _get_supported_branches() -> Generator[str, None, None]:
+    for branch in _get_branches():
+        if not branch.startswith("aosp_kernel-common-android"):
+            continue
+        if branch == "aosp_kernel-common-android-mainline":
+            yield branch
+            continue
+        android_release = branch.removeprefix(
+            "aosp_kernel-common-android").split("-", maxsplit=1)[0]
+        if not android_release.isdecimal():
+            continue
+        android_release = int(android_release)
+        if android_release < 15:
+            continue
+        yield branch
+
+
+class KleafBootstrapTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        print("INFO: Calculating supported branches")
+        cls.supported_branches = _get_supported_branches()
+        return super().setUpClass()
+
+    def test_infer_build_id(self):
+        for branch in self.supported_branches:
+            with self.subTest(branch=branch):
+                obj = KleafBootstrap(argparse.Namespace(
+                    branch=branch,
+                    build_target=_DEFAULT_BUILD_TARGET,
+                    build_id=None,
+                    url_fmt=_ARTIFACT_URL_FMT,
+                ), [])
+                obj._set_build_id()
+                print(f"For branch {branch}, checking {obj.build_id}")
+                buf = io.BytesIO()
+                obj._download_artifact("BUILD_INFO", buf)
+
+    @unittest.skip("init_ddk.zip is not published yet")
+    def test_run_with_help(self):
+        for branch in self.supported_branches:
+            with self.subTest(branch=branch):
+                obj = KleafBootstrap(argparse.Namespace(
+                    branch=branch,
+                    build_target=_DEFAULT_BUILD_TARGET,
+                    build_id=None,
+                    url_fmt=_ARTIFACT_URL_FMT,
+                ), ["-h"])
+                obj.run()
+
+
+if __name__ == "__main__":
+    unittest.main()