3pp/ubuntu_24.04_binaries/download_extract.py

#!/usr/bin/env vpython3
# Copyright 2026 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Pulls the most recent Ubuntu ISOs and repackages them for CIPD.

The kernel and initrd are extracted from the ISO and packaged together for
PXE installs for bare metal machines.
"""

import argparse
import os
import re
import shutil
import subprocess
import sys
import tempfile

from bs4 import BeautifulSoup
import requests

RELEASE = 'noble'
RELEASE_NUM = '24.04'


def get_arch():
  _3pp_platform = os.environ['_3PP_PLATFORM']
  if _3pp_platform == 'linux-amd64':
    return 'amd64'
  elif _3pp_platform == 'linux-arm64':
    return 'arm64'
  else:
    raise RuntimeError(f'Unsupported target platform {_3pp_platform}')


def get_iso_info(arch):
  if arch == 'amd64':
    url = f'https://releases.ubuntu.com/{RELEASE}/'
    page = requests.get(url)
    page.raise_for_status()
    soup = BeautifulSoup(page.text, 'html.parser')
    link = soup.find(
        'a', href=re.compile(f'ubuntu-{RELEASE_NUM}.*live-server-amd64.iso'))
    if not link:
      raise RuntimeError(f'Could not find download link for {arch} on {url}')
    iso_name = link['href']
    iso_url = f'{url}{iso_name}'
  elif arch == 'arm64':
    url = f'https://cdimage.ubuntu.com/releases/{RELEASE_NUM}/release/'
    page = requests.get(url)
    page.raise_for_status()
    soup = BeautifulSoup(page.text, 'html.parser')
    link = soup.find(
        'a', href=re.compile(f'ubuntu-{RELEASE_NUM}.*live-server-arm64.iso'))
    if not link:
      raise RuntimeError(f'Could not find download link for {arch} on {url}')
    iso_name = link['href']
    iso_url = f'{url}{iso_name}'
  else:
    raise ValueError(f'Unknown arch: {arch}')

  # Extract version from filename eg: ubuntu-24.04.5-live-server-amd64.iso
  match = re.search(r'(\d{2}\.\d{2}\.\d+)', iso_name)
  if not match:
    raise RuntimeError(f'Could not extract version from ISO name: {iso_name}')
  version = match.group(1)

  return iso_name, iso_url, version


def cmd_get_latest_version():
  arch = get_arch()
  _, _, version = get_iso_info(arch)
  print(version)


def cmd_checkout(checkout_path):
  arch = get_arch()
  iso_name, iso_url, version = get_iso_info(arch)
  requested_version = os.environ.get('_3PP_VERSION')
  if requested_version != version:
    raise RuntimeError(
        f'Requested version {requested_version} does not match latest {version}'
    )

  # Download the ISO
  print(f'Downloading {iso_url} to {iso_name}...')
  with requests.get(iso_url, stream=True) as r:
    r.raise_for_status()
    with open(iso_name, 'wb') as f:
      for chunk in r.iter_content(chunk_size=8192):
        f.write(chunk)
  print(f'Downloaded {iso_name}.')

  if not os.path.exists(iso_name):
    raise RuntimeError(f'Expected to find downloaded ISO: {iso_name}')

  with tempfile.TemporaryDirectory() as tmpdir:
    print(f'Extracting {iso_name} to {tmpdir}...')
    _7z_path = os.environ.get('_7z')
    if not _7z_path:
      raise RuntimeError('Error: _7z environment variable not set.')
    _7z_exe = os.path.join(_7z_path, '7zz')

    if not os.path.exists(_7z_exe) or not os.access(_7z_exe, os.X_OK):
      raise RuntimeError(f"Error: '{_7z_exe}' not found or is not executable. "
                         'Please ensure it is installed and accessible.')

    result = subprocess.run(
        [_7z_exe, 'x',
         os.path.abspath(iso_name), '-y', '-o' + tmpdir],
        check=True,
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
      print(f'7zz command returned non-zero exit status {result.returncode}')
      print(f'Stdout: {result.stdout}')
      print(f'Stderr: {result.stderr}', file=sys.stderr)
      print('Continuing despite 7zz errors...')

    vmlinuz_path = os.path.join(tmpdir, 'casper', 'vmlinuz')
    initrd_path = os.path.join(tmpdir, 'casper', 'initrd')

    if not os.path.exists(vmlinuz_path):
      raise RuntimeError(f'Could not find vmlinuz in {tmpdir}/casper')
    if not os.path.exists(initrd_path):
      raise RuntimeError(f'Could not find initrd in {tmpdir}/casper')

    # Move the extracted files to the checkout path for the install script
    os.makedirs(checkout_path, exist_ok=True)
    shutil.move(vmlinuz_path, os.path.join(checkout_path, 'vmlinuz'))
    shutil.move(initrd_path, os.path.join(checkout_path, 'initrd'))

    hwe_vmlinuz_path = os.path.join(tmpdir, 'casper', 'hwe-vmlinuz')
    hwe_initrd_path = os.path.join(tmpdir, 'casper', 'hwe-initrd')

    if os.path.exists(hwe_vmlinuz_path):
      shutil.move(hwe_vmlinuz_path, os.path.join(checkout_path, 'hwe-vmlinuz'))
    # No warning if not found, it's optional

    if os.path.exists(hwe_initrd_path):
      shutil.move(hwe_initrd_path, os.path.join(checkout_path, 'hwe-initrd'))
    # No warning if not found, it's optional

    shutil.move(iso_name, os.path.join(checkout_path, 'ubuntu.iso'))

  print('Checkout complete.')


def _parse_args() -> argparse.Namespace:
  parser = argparse.ArgumentParser()
  subparsers = parser.add_subparsers(dest='action', required=True)

  get_latest_version_parser = subparsers.add_parser('latest')
  get_latest_version_parser.set_defaults(
      func=lambda args: cmd_get_latest_version())

  checkout_parser = subparsers.add_parser('checkout')
  checkout_parser.add_argument('checkout_path')
  checkout_parser.set_defaults(
      func=lambda args: cmd_checkout(args.checkout_path))

  return parser.parse_args()


def main() -> int:
  args = _parse_args()
  args.func(args)
  return 0


if __name__ == '__main__':
  sys.exit(main())