1 files changed, 192 insertions, 0 deletions
diff --git a/rust_tools/copy_rust_bootstrap.py b/rust_tools/copy_rust_bootstrap.py
new file mode 100755
index 00000000..5da8007f
--- /dev/null
+++ b/rust_tools/copy_rust_bootstrap.py
@@ -0,0 +1,192 @@
+#!/usr/bin/env python3
+# Copyright 2022 The ChromiumOS Authors.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Copies rust-bootstrap artifacts from an SDK build to localmirror.
+
+We use localmirror to host these artifacts, but they've changed a bit over
+time, so simply `gsutil.py cp $FROM $TO` doesn't work. This script allows the
+convenience of the old `cp` command.
+"""
+
+import argparse
+import logging
+import os
+from pathlib import Path
+import shutil
+import subprocess
+import sys
+import tempfile
+from typing import List
+
+
+_LOCALMIRROR_ROOT = 'gs://chromeos-localmirror/distfiles/'
+
+
+def _is_in_chroot() -> bool:
+  return Path('/etc/cros_chroot_version').exists()
+
+
+def _ensure_pbzip2_is_installed():
+  if shutil.which('pbzip2'):
+    return
+
+  logging.info('Auto-installing pbzip2...')
+  subprocess.run(['sudo', 'emerge', '-G', 'pbzip2'], check=True)
+
+
+def _determine_target_path(sdk_path: str) -> str:
+  """Determine where `sdk_path` should sit in localmirror."""
+  gs_prefix = 'gs://'
+  if not sdk_path.startswith(gs_prefix):
+    raise ValueError(f'Invalid GS path: {sdk_path!r}')
+
+  file_name = Path(sdk_path[len(gs_prefix):]).name
+  return _LOCALMIRROR_ROOT + file_name
+
+
+def _download(remote_path: str, local_file: Path):
+  """Downloads the given gs:// path to the given local file."""
+  logging.info('Downloading %s -> %s', remote_path, local_file)
+  subprocess.run(
+      ['gsutil.py', 'cp', remote_path,
+       str(local_file)],
+      check=True,
+  )
+
+
+def _debinpkgify(binpkg_file: Path) -> Path:
+  """Converts a binpkg into the files it installs.
+
+  Note that this function makes temporary files in the same directory as
+  `binpkg_file`. It makes no attempt to clean them up.
+  """
+  logging.info('Converting %s from a binpkg...', binpkg_file)
+
+  # The SDK builder produces binary packages:
+  # https://wiki.gentoo.org/wiki/Binary_package_guide
+  #
+  # Which means that `binpkg_file` is in the XPAK format. We want to split
+  # that out, and recompress it from zstd (which is the compression format
+  # that CrOS uses) to bzip2 (which is what we've historically used, and
+  # which is what our ebuild expects).
+  tmpdir = binpkg_file.parent
+
+  def _mkstemp(suffix=None) -> str:
+    fd, file_path = tempfile.mkstemp(dir=tmpdir, suffix=suffix)
+    os.close(fd)
+    return Path(file_path)
+
+  # First, split the actual artifacts that land in the chroot out to
+  # `temp_file`.
+  artifacts_file = _mkstemp()
+  logging.info('Extracting artifacts from %s into %s...', binpkg_file,
+               artifacts_file)
+  with artifacts_file.open('wb') as f:
+    subprocess.run(
+        [
+            'qtbz2',
+            '-s',
+            '-t',
+            '-O',
+            str(binpkg_file),
+        ],
+        check=True,
+        stdout=f,
+    )
+
+  decompressed_artifacts_file = _mkstemp()
+  decompressed_artifacts_file.unlink()
+  logging.info('Decompressing artifacts from %s to %s...', artifacts_file,
+               decompressed_artifacts_file)
+  subprocess.run(
+      [
+          'zstd',
+          '-d',
+          str(artifacts_file),
+          '-o',
+          str(decompressed_artifacts_file),
+      ],
+      check=True,
+  )
+
+  # Finally, recompress it as a tbz2.
+  tbz2_file = _mkstemp('.tbz2')
+  logging.info(
+      'Recompressing artifacts from %s to %s (this may take a while)...',
+      decompressed_artifacts_file, tbz2_file)
+  with tbz2_file.open('wb') as f:
+    subprocess.run(
+        [
+            'pbzip2',
+            '-9',
+            '-c',
+            str(decompressed_artifacts_file),
+        ],
+        check=True,
+        stdout=f,
+    )
+  return tbz2_file
+
+
+def _upload(local_file: Path, remote_path: str, force: bool):
+  """Uploads the local file to the given gs:// path."""
+  logging.info('Uploading %s -> %s', local_file, remote_path)
+  cmd_base = ['gsutil.py', 'cp', '-a', 'public-read']
+  if not force:
+    cmd_base.append('-n')
+  subprocess.run(
+      cmd_base + [str(local_file), remote_path],
+      check=True,
+      stdin=subprocess.DEVNULL,
+  )
+
+
+def main(argv: List[str]):
+  logging.basicConfig(
+      format='>> %(asctime)s: %(levelname)s: %(filename)s:%(lineno)d: '
+      '%(message)s',
+      level=logging.INFO,
+  )
+
+  parser = argparse.ArgumentParser(
+      description=__doc__,
+      formatter_class=argparse.RawDescriptionHelpFormatter,
+  )
+
+  parser.add_argument(
+      'sdk_artifact',
+      help='Path to the SDK rust-bootstrap artifact to copy. e.g., '
+      'gs://chromeos-prebuilt/host/amd64/amd64-host/'
+      'chroot-2022.07.12.134334/packages/dev-lang/'
+      'rust-bootstrap-1.59.0.tbz2.')
+  parser.add_argument(
+      '-n',
+      '--dry-run',
+      action='store_true',
+      help='Do everything except actually uploading the artifact.')
+  parser.add_argument(
+      '--force',
+      action='store_true',
+      help='Upload the artifact even if one exists in localmirror already.')
+  opts = parser.parse_args(argv)
+
+  if not _is_in_chroot():
+    parser.error('Run me from within the chroot.')
+  _ensure_pbzip2_is_installed()
+
+  target_path = _determine_target_path(opts.sdk_artifact)
+  with tempfile.TemporaryDirectory() as tempdir:
+    download_path = Path(tempdir) / 'sdk_artifact'
+    _download(opts.sdk_artifact, download_path)
+    file_to_upload = _debinpkgify(download_path)
+    if opts.dry_run:
+      logging.info('--dry-run specified; skipping upload of %s to %s',
+                   file_to_upload, target_path)
+    else:
+      _upload(file_to_upload, target_path, opts.force)
+
+
+if __name__ == '__main__':
+  main(sys.argv[1:])