ros-rt-img

#!/usr/bin/env python3
import argparse
import logging
import subprocess
import os
import shutil
import sys

from image_builder.builder import Builder, LOOP_DEVICE_FILENAME, DEFAULT_CHROOT_PATH, CMAKE_TOOLCHAIN_FILE

logging.basicConfig(format="[%(asctime)s][%(levelname)s] %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.DEBUG)


def parse_arguments():
  parser = argparse.ArgumentParser(
    description="ROS + RT Raspberry Pi 4 image builder"
  )

  parser.add_argument(
    "-o",
    "--out-dir",
    default="out",
    help="The directory in which to place the built image. Default: out",
  )

  parser.add_argument(
    "--cache-dir",
    default="cache",
    help="The directory in which to place to cached artifacts, like the downloaded base image. Default: cache",
  )

  parser.add_argument(
    "--chroot-path",
    default=DEFAULT_CHROOT_PATH,
    help="The path the OS image will be mounted and chrooted at. Default: /tmp/rpi4-image-build"
  )

  subparsers = parser.add_subparsers(dest="action")
  subparsers.required = True

  build_parser = subparsers.add_parser("build", help="Builds the image")

  build_parser.add_argument(
    "--pause-after",
    default=None,
    choices=[
      "download_and_extract_image_if_necessary",
      "setup_loop_device_and_mount_partitions",
      "prepare_chroot",
      "copy_files_to_chroot",
      "run_phase1_host_scripts",
      "run_phase1_target_scripts",
      "run_phase2_host_scripts",
      "run_phase2_target_scripts",
      "cleanup_chroot",
      "umount_everything",
    ],
    help="Pause after a particular build stage. Default: None (will not pause)",
  )

  build_parser.add_argument(
    "profiles",
    type=str,
    nargs="*",
    default=["focal-rt", "focal-rt-galactic"],
    help="The profiles to be built. The order is how they will overlay each other. Default: focal-rt, focal-rt-galactic"
  )

  subparsers.add_parser("teardown", help="Reset the ongoing session (if any) and unmount the mounted image and teardown the loop device (if required)")
  subparsers.add_parser("chroot", help="Enters the chroot_path if possible")

  mount_parser = subparsers.add_parser("mount", help="Mount a built image for interactive exploration and/or cross-compilation purposes. only works with Raspberry Pi Ubuntu images.")
  mount_parser.add_argument("image", type=str, help="The path to the image to be mounted")

  subparsers.add_parser("umount", help="unmount an image mounted for interaction via the mount sub-command")

  subparsers.add_parser("toolchain", help="prints the path to the cmake toolchain location")

  # Convert the action into a callable function.
  args = parser.parse_args()
  args.action = globals()[args.action]
  return args


def build(args):
  builder = Builder(
    profile_dirs=args.profiles,
    cache_dir=args.cache_dir,
    out_dir=args.out_dir,
    chroot_path=args.chroot_path,
    pause_after=args.pause_after
  )

  builder.build()


def teardown(args):
  # TODO: refactor Builder such that it supports lazy initialization of the
  # profile directories, so that code from here and from clean can be reused.

  session_file = os.path.join(args.cache_dir, "session.txt")
  session_loop_device_file = os.path.join(args.cache_dir, LOOP_DEVICE_FILENAME)

  if os.path.isfile(session_loop_device_file):
    with open(session_loop_device_file) as f:
      loop_device = f.read().strip()
  else:
    loop_device = None

  logging.info(f"Unmounting {args.chroot_path}")
  subprocess.run(["umount", "-R", args.chroot_path])

  if loop_device is not None:
    logging.info(f"detaching loop device {loop_device}")
    subprocess.run(["losetup", "-d", loop_device])
    os.remove(session_loop_device_file)
  else:
    logging.info(f"loop device not found in {session_loop_device_file}, so not detaching...")

  if os.path.isfile(session_file):
    logging.info(f"removing session file {session_file}")
    os.remove(session_file)
  else:
    logging.info(f"session file {session_file} not found, not removing...")


def chroot(args):
  if not os.path.isdir(args.chroot_path):
    logging.error(f"cannot chroot into {args.chroot_path} as it doesn't exist")
    sys.exit(1)

  subprocess.run(["systemd-nspawn", "-D", args.chroot_path, "bash"])


# TODO: there are some duplication with Builder.prepare_chroot and Builder.setup_loop_device_and_mount_partitions, but different enough that it is duplicated.
def mount(args):
  loop_device_cache_file = os.path.join(args.cache_dir, LOOP_DEVICE_FILENAME)
  if os.path.exists(loop_device_cache_file):
    logging.error("{} already exists, make sure no other build and mounts are in progress!".format(loop_device_cache_file))
    sys.exit(1)

  loop_device = subprocess.check_output(["losetup", "-P", "--show", "-f", args.image]).decode("utf-8").strip()
  with open(loop_device_cache_file, "w") as f:
    f.write(loop_device)

  logging.info("setting up loop device on {}".format(loop_device))

  os.chmod(loop_device_cache_file, 0o666) # since this script is root, set it to world writable so it can be more easily deleted

  # TODO: this is kind of a hack, as the mount point for each partition is
  # specified in the image_mounts in the config.ini. It is hard-coded for now
  # as we assume the raspberry pi Ubuntu layout.
  for i, mount_point in reversed(list(enumerate(["/boot/firmware", "/"]))):
    i += 1
    device_name = f"{loop_device}p{i}"
    mount_point = os.path.join(args.chroot_path, mount_point.lstrip("/"))

    logging.info("mounting {} to {}".format(device_name, mount_point))
    subprocess.run(["mount", device_name, mount_point], check=True)

  logging.info("copying resolv.conf and qemu-user-static")

  shutil.move(
    os.path.join(args.chroot_path, "etc", "resolv.conf"),
    os.path.join(args.chroot_path, "etc", "resolv.conf.bak")
  )

  # TODO: also a hack, as qemu_user_static_path in the config.ini, but that's also a hack so...
  qemu_user_static_path = "/usr/bin/qemu-aarch64-static"
  shutil.copy(
    qemu_user_static_path,
    os.path.join(args.chroot_path, qemu_user_static_path.lstrip("/"))
  )

  shutil.copy(
    "/etc/resolv.conf",
    os.path.join(args.chroot_path, "etc", "resolv.conf")
  )


def umount(args):
  os.remove(os.path.join(args.chroot_path, "etc", "resolv.conf"))
  shutil.move(
    os.path.join(args.chroot_path, "etc", "resolv.conf.bak"),
    os.path.join(args.chroot_path, "etc", "resolv.conf"),
  )

  os.remove(os.path.join(args.chroot_path, "usr/bin/qemu-aarch64-static"))

  subprocess.run(["umount", "-R", args.chroot_path], check=True)

  loop_device_cache_file = os.path.join(args.cache_dir, LOOP_DEVICE_FILENAME)
  if not os.path.exists(loop_device_cache_file):
    logging.error("failed to detach loop device due to missing loop-device cache file, you'll have to manually call losetup -d")
  else:
    with open(loop_device_cache_file) as f:
      loop_device = f.read().strip()

    subprocess.run(["losetup", "-d", loop_device], check=True)
    os.remove(loop_device_cache_file)


def toolchain(args):
  print(CMAKE_TOOLCHAIN_FILE)


def main():
  args = parse_arguments()
  args.action(args)


if __name__ == "__main__":
  main()