kleaf: init_ddk: Add support to download files

* This change introduces a new function to download files, by
  contstructing a valid URL using information like build_id,
  build_target and remote_filename.

* Due to hermeticity reasons, the python toolchain lacks valid
  certificates to verify the downloads, for that reason, the downloads
  rely on the host toolchain just for these.

* To avoid spam during these calls, and in case of any error only the
  last line of the stacktrace is printed in case of any exception.

Bug: 328770706
Change-Id: I80fb8867a6d009043d0d2bb376bf0c2557774e17
Signed-off-by: Ulises Mendez Martinez <umendez@google.com>

4 files changed, 102 insertions, 0 deletions

diff --git a/BUILD.bazel b/BUILD.bazel
index 053d67b..6634b29 100644
--- a/BUILD.bazel
+++ b/BUILD.bazel
@@ -464,6 +464,7 @@ py_binary(
 py_binary(
     name = "init_ddk",
     srcs = ["init/init_ddk.py"],
+    data = ["init/init_download.py"],
     main = "init/init_ddk.py",
     visibility = ["//visibility:private"],
 )
diff --git a/init/init_ddk.py b/init/init_ddk.py
index 2f736d1..1972501 100644
--- a/init/init_ddk.py
+++ b/init/init_ddk.py
@@ -22,9 +22,11 @@ import json
 import logging
 import pathlib
 import shutil
+import subprocess
 import sys
 import tempfile
 import textwrap
+import urllib
 
 _TOOLS_BAZEL = "tools/bazel"
 _DEVICE_BAZELRC = "device.bazelrc"
@@ -173,6 +175,48 @@ class KleafProjectSetter:
             """),
         )
 
+    def _get_url(self, remote_filename: str) -> str:
+        """Returns a valid url when it can be formed with target and id."""
+        url = self.url_fmt.format(
+            build_id=self.build_id,
+            build_target=self.build_target,
+            filename=urllib.parse.quote(remote_filename, safe=""),  # / -> %2F
+        )
+        url_with_fake_id = self.url_fmt.format(
+            build_id="__FAKE_BUILD_NUMBER_PLACEHOLDER__",
+            build_target=self.build_target,
+            filename=urllib.parse.quote(remote_filename, safe=""),  # / -> %2F
+        )
+        if not self.build_id and url != url_with_fake_id:
+            return None
+        return url
+
+    def _download(self, remote_filename: str, out_file_name: str):
+        if not self.url_fmt:
+            logging.error(
+                "Unable to download file %s because --url_fmt was not set.",
+                remote_filename,
+            )
+            return
+        url = self._get_url(remote_filename)
+        if not url:
+            logging.error(
+                "Unable to download %s file because --build_id is missing.",
+                remote_filename,
+            )
+            return
+        # Workaround: Rely on host keychain to download files.
+        # This is needed otheriwese downloads fail when running this script
+        #   using the hermetic Python toolchain.
+        subprocess.check_call(
+            [
+                "python3",
+                pathlib.Path(__file__).parent / "init_download.py",
+                url,
+                out_file_name,
+            ],
+        )
+
     def _handle_ddk_workspace(self):
         if not self.ddk_workspace:
             return
@@ -217,6 +261,7 @@ if __name__ == "__main__":
         "--build_id",
         type=str,
         help="the build id to download the build for, e.g. 6148204",
+        default=None,
     )
     parser.add_argument(
         "--build_target",
diff --git a/init/init_ddk_test.py b/init/init_ddk_test.py
index a13b2c7..fe54e96 100644
--- a/init/init_ddk_test.py
+++ b/init/init_ddk_test.py
@@ -195,6 +195,29 @@ class KleafProjectSetterTest(parameterized.TestCase):
                 expected=str(prebuilts_dir_abs),
             )
 
+    def test_download_works_for_local_file(self):
+        """Tests that local files can be downloaded."""
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            tmp_dir = pathlib.Path(tmp_dir)
+            remote_file = tmp_dir / "remote_file"
+            remote_file.write_text("Hello World!")
+            out_file = tmp_dir / "out_file"
+            url_fmt = f"file://{tmp_dir}/{{filename}}"
+            init_ddk.KleafProjectSetter(
+                build_id=None,
+                build_target=None,
+                ddk_workspace=None,
+                kleaf_repo=None,
+                local=False,
+                prebuilts_dir=None,
+                url_fmt=url_fmt,
+            )._download(
+                remote_filename="remote_file",
+                out_file_name=out_file,
+            )
+            self.assertTrue(out_file.exists())
+            self.assertEqual(out_file.read_text(), "Hello World!")
+
 
 # This could be run as: tools/bazel test //build/kernel:init_ddk_test --test_output=all
 if __name__ == "__main__":
diff --git a/init/init_download.py b/init/init_download.py
new file mode 100644
index 0000000..c125935
--- /dev/null
+++ b/init/init_download.py
@@ -0,0 +1,33 @@
+#!/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.
+
+"""Small utility to download files."""
+
+import shutil
+import sys
+import traceback
+import urllib.request
+
+try:
+    with urllib.request.urlopen(sys.argv[1]) as input_file, open(
+        sys.argv[2], "wb"
+    ) as output_file:
+        shutil.copyfileobj(input_file, output_file)
+except Exception as exc:  # pylint: disable=broad-exception-caught
+    formatted_lines = traceback.format_exc().splitlines()
+    if formatted_lines:
+        print(formatted_lines[-1], file=sys.stderr)
+    sys.exit(1)