linux-image-manager/lim/image/transfer.py

"""Download the selected image and transfer it onto the target device."""

import shutil
from pathlib import Path

from lim import device as device_module
from lim import runner, ui
from lim.errors import LimError
from lim.image import loopimg
from lim.image.plan import DEFAULT_BOOT_SIZE, ImagePlan
from lim.image.session import ImageSession


def download_image(plan: ImagePlan, *, force_prompt: bool = True) -> None:
    if force_prompt and ui.confirm("Should the image download be forced?"):
        if plan.image_path.is_file():
            ui.info(f"Removing image {plan.image_path}.")
            plan.image_path.unlink()
        else:
            ui.info(f"Forcing download wasn't neccessary. File {plan.image_path} doesn't exist.")

    ui.info("Start Download procedure...")
    if plan.image_path.is_file():
        ui.info("Image exist local. Download skipped.")
        return
    ui.info(f'Image "{plan.image_name}" doesn\'t exist under local path "{plan.image_path}".')
    ui.info(f'Image "{plan.image_name}" gets downloaded from "{plan.download_url}"...')
    runner.run(
        ["wget", plan.download_url, "-O", str(plan.image_path)],
        error_msg=f'Download from "{plan.download_url}" failed.',
    )


def arch_partition_input(boot_size: str) -> str:
    """Fdisk answers: FAT32 boot partition of boot_size plus root partition."""
    return (
        "o\n"  # clear out any partitions on the drive
        "p\n"  # list partitions (should be empty)
        "n\np\n1\n\n"  # new primary partition 1, default start sector
        f"{boot_size}\n"
        "t\nc\n"  # set partition 1 to type W95 FAT32 (LBA)
        "n\np\n2\n\n\n"  # new primary partition 2 over the remaining space
        "w\n"  # write partition table
    )


def decompress_command(image_path: Path) -> list[str]:
    """Build the command that streams the raw image to stdout for dd."""
    suffix = image_path.suffix
    if suffix == ".zip":
        return ["unzip", "-p", str(image_path)]
    if suffix == ".gz":
        return ["gunzip", "-c", str(image_path)]
    if suffix == ".iso":
        return ["pv", str(image_path)]
    if suffix == ".xz":
        return ["unxz", "-c", str(image_path)]
    raise LimError(f'Image transfer for "{image_path.name}" is not supported yet!')


def _luks_format_root(plan: ImagePlan, session: ImageSession) -> None:
    luks_format = [
        "cryptsetup",
        "-v",
        "luksFormat",
        "-c",
        "aes-xts-plain64",
        "-s",
        "512",
        "-h",
        "sha512",
        "--use-random",
        "-i",
        "1000",
    ]
    if plan.luks_memory_cost:
        ui.info(
            f"Formating {session.root_partition_path} with LUKS "
            f"with --pbkdf-memory set to {plan.luks_memory_cost}"
        )
        luks_format += ["--pbkdf-memory", plan.luks_memory_cost]
    else:
        ui.info(f"Formating {session.root_partition_path} with LUKS")
    runner.run([*luks_format, session.root_partition_path], sudo=True)


def transfer_arch_image(plan: ImagePlan, session: ImageSession) -> None:
    boot_size = plan.boot_size or DEFAULT_BOOT_SIZE
    ui.info(f"The boot partition will be set to {boot_size}.")
    ui.info("Creating partitions...")
    runner.run(
        ["fdisk", session.device.path],
        input_text=arch_partition_input(boot_size),
        sudo=True,
    )

    ui.info("Format boot partition...")
    runner.run(["mkfs.vfat", session.boot_partition_path], sudo=True)

    if plan.encrypt_system:
        _luks_format_root(plan, session)
        session.decrypt_root()

    ui.info("Format root partition...")
    runner.run([f"mkfs.{plan.root_filesystem}", "-f", session.root_mapper_path], sudo=True)
    session.mount_partitions()

    ui.info("Root files will be transfered to device...")
    runner.run(
        ["bsdtar", "-xpf", str(plan.image_path), "-C", str(session.root_mount_path)],
        sudo=True,
    )
    runner.sync_disks()

    ui.info("Boot files will be transfered to device...")
    boot_source = session.root_mount_path / "boot"
    for entry in sorted(boot_source.iterdir()):
        runner.run(["mv", "-v", str(entry), str(session.boot_mount_path)], sudo=True)


def _create_encrypted_layout(plan: ImagePlan, session: ImageSession) -> None:
    """Partition the target, LUKS-format the root, format and mount both."""
    boot_size = plan.boot_size or DEFAULT_BOOT_SIZE
    ui.info(f"Creating partitions (boot {boot_size})...")
    runner.run(
        ["fdisk", session.device.path],
        input_text=arch_partition_input(boot_size),
        sudo=True,
    )
    runner.run(["mkfs.vfat", session.boot_partition_path], sudo=True)
    _luks_format_root(plan, session)
    session.decrypt_root()
    runner.run([f"mkfs.{plan.root_filesystem}", "-F", session.root_mapper_path], sudo=True)
    session.mount_partitions()


def _copy_tree(source: str, target: str, flags: list[str]) -> None:
    """Copy a directory's contents, showing progress when rsync is available."""
    if shutil.which("rsync"):
        runner.run(["rsync", *flags, "--info=progress2", f"{source}/", f"{target}/"], sudo=True)
    else:
        runner.run(["cp", "-a", f"{source}/.", target], sudo=True)


def _fix_boot_fstab(session: ImageSession) -> None:
    """Repoint the stock boot fstab line at the freshly created boot partition.

    The copied image mounts /boot(/firmware) by the stock PARTUUID, which no
    longer exists after repartitioning; swap in the new boot UUID, mount point
    unchanged. No-op when the image carries no boot line.
    """
    fstab = session.root_mount_path / "etc/fstab"
    if not fstab.is_file():
        return
    lines = []
    for line in fstab.read_text().splitlines():
        fields = line.split()
        if not line.lstrip().startswith("#") and fields[1:2] and fields[1].startswith("/boot"):
            fields[0] = f"UUID={session.boot_partition_uuid}"
            lines.append("  ".join(fields))
        else:
            lines.append(line)
    fstab.write_text("".join(f"{entry}\n" for entry in lines))


def transfer_disk_image(plan: ImagePlan, session: ImageSession) -> None:
    """Copy a full disk image (boot + root partitions) into a fresh encrypted layout.

    Distributions that ship a partitioned .img rather than a rootfs tarball
    (Raspberry Pi OS, moode, RetroPie, Manjaro ARM) are loop-mounted and their
    first two partitions copied into a new LUKS container.
    """
    raw = plan.image_path.with_suffix("")
    ui.info(f"Decompressing {plan.image_path} to {raw}...")
    runner.pipeline(
        decompress_command(plan.image_path),
        ["dd", f"of={raw}", "bs=4M", "conv=fsync", "status=progress"],
        error_msg=f"Decompressing {plan.image_path} failed.",
    )
    loop = loopimg.attach(raw)
    source_boot = session.working_folder / "source-boot"
    source_root = session.working_folder / "source-root"
    try:
        source_boot.mkdir()
        source_root.mkdir()
        runner.run(["mount", "-o", "ro", loopimg.partition(loop, 1), str(source_boot)], sudo=True)
        runner.run(["mount", "-o", "ro", loopimg.partition(loop, 2), str(source_root)], sudo=True)

        _create_encrypted_layout(plan, session)

        ui.info("Copying the root filesystem into the encrypted container...")
        _copy_tree(str(source_root), str(session.root_mount_path), ["-aHAX"])
        ui.info("Copying the boot partition...")
        _copy_tree(str(source_boot), str(session.boot_mount_path), ["-a"])
        runner.sync_disks()
        _fix_boot_fstab(session)
    finally:
        runner.run(["umount", "-v", str(source_boot)], sudo=True, check=False)
        runner.run(["umount", "-v", str(source_root)], sudo=True, check=False)
        loopimg.detach(loop)
        for path in (source_boot, source_root):
            runner.run(["rmdir", str(path)], sudo=True, check=False)


def transfer_image(plan: ImagePlan, session: ImageSession, *, interactive: bool = True) -> None:
    if interactive:
        if ui.confirm(f"Should the partition table of {session.device.path} be deleted?"):
            ui.info("Deleting...")
            runner.run(["wipefs", "-a", session.device.path], sudo=True)
        else:
            ui.info("Skipping partition table deletion...")
        device_module.overwrite_device(session.device)

    ui.info("Starting image transfer...")
    if plan.distribution == "arch":
        transfer_arch_image(plan, session)
        return
    if plan.encrypt_system:
        # Non-tarball distros ship a full .img; copy it into a LUKS container.
        transfer_disk_image(plan, session)
        return
    blocksize = session.device.optimal_blocksize
    ui.info(f"Transfering {plan.image_path.suffix} file...")
    runner.pipeline(
        decompress_command(plan.image_path),
        [
            "dd",
            f"of={session.device.path}",
            f"bs={blocksize}",
            "conv=fsync",
            "status=progress",
        ],
        sudo_last=True,
        error_msg=f"DD {plan.image_path} to {session.device.path} failed.",
    )
    runner.sync_disks()