parchbootstrap/image/cli.py

import click
import logging
import os

from signal import pause
from typing import Optional

from config.state import config, Profile
from constants import BASE_LOCAL_PACKAGES, BASE_PACKAGES, LUKS_MAPPER_DEFAULT
from devices.device import get_profile_device
from exec.file import makedir
from flavours.flavour import get_profile_flavour
from packages.build import build_enable_qemu_binfmt, build_packages, filter_pkgbuilds
from wrapper import enforce_wrap

from .boot import cmd_boot
from .cryptsetup import encryption_option, get_cryptmapper_path, luks_close, luks_create, luks_open
from .flash import cmd_flash
from .image import (
    IMG_FILE_BOOT_DEFAULT_SIZE,
    create_boot_fs,
    create_img_file,
    create_root_fs,
    dd_image,
    get_device_chroot,
    get_image_path,
    install_rootfs,
    losetup_setup_image,
    mount_chroot,
    partprobe,
    partition_device,
)


@click.group(name='image')
def cmd_image():
    """Build, flash and boot device images"""


for cmd in [cmd_boot, cmd_flash]:
    cmd_image.add_command(cmd)

sectorsize_option = click.option(
    '-b',
    '--sector-size',
    help="Override the device's sector size",
    type=int,
    default=None,
)


@cmd_image.command(name='build')
@click.argument('profile_name', required=False)
@click.option(
    '--local-repos/--no-local-repos',
    '-l/-L',
    help='Whether to use local package repos at all or only use HTTPS repos.',
    default=True,
    show_default=True,
    is_flag=True,
)
@click.option(
    '--build-pkgs/--no-build-pkgs',
    '-p/-P',
    help='Whether to build missing/outdated local packages if local repos are enabled.',
    default=True,
    show_default=True,
    is_flag=True,
)
@click.option(
    '--no-download-pkgs',
    help='Disable trying to download packages instead of building if building is enabled.',
    default=False,
    is_flag=True,
)
@click.option(
    '--block-target',
    help='Override the block device file to write the final image to',
    type=click.Path(),
    default=None,
)
@click.option(
    '--skip-part-images',
    help='Skip creating image files for the partitions and directly work on the target block device.',
    default=False,
    is_flag=True,
)
@encryption_option
@sectorsize_option
def cmd_build(
    profile_name: Optional[str] = None,
    local_repos: bool = True,
    build_pkgs: bool = True,
    no_download_pkgs=False,
    block_target: Optional[str] = None,
    sector_size: Optional[int] = None,
    skip_part_images: bool = False,
    encryption: Optional[bool] = None,
    encryption_password: Optional[str] = None,
    encryption_mapper: str = LUKS_MAPPER_DEFAULT,
):
    """
    Build a device image.

    Unless overriden, required packages will be built or preferably downloaded from HTTPS repos.
    """

    config.enforce_profile_device_set()
    config.enforce_profile_flavour_set()
    enforce_wrap()
    device = get_profile_device(profile_name)
    arch = device.arch
    # check_programs_wrap(['makepkg', 'pacman', 'pacstrap'])
    profile: Profile = config.get_profile(profile_name)
    flavour = get_profile_flavour(profile_name)
    rootfs_size_mb = flavour.parse_flavourinfo().rootfs_size * 1000 + int(profile.size_extra_mb)
    bootfs_size_str = IMG_FILE_BOOT_DEFAULT_SIZE
    bootfs_size_mb = -1
    if bootfs_size_str.endswith('M'):
        bootfs_size_mb = int(bootfs_size_str.rstrip('M'))
    elif bootfs_size_str.endswith('G'):
        bootfs_size_mb = int(bootfs_size_str.rstrip('G')) * 1024
    elif not bootfs_size_str.isdecimal():
        raise Exception(f"Couldn't part bootfs target size as megabytes: {bootfs_size_str}")
    else:
        bootfs_size_mb = int(bootfs_size_str)

    if encryption is None:
        encryption = profile.encryption

    packages = BASE_LOCAL_PACKAGES + [device.package.name, flavour.pkgbuild.name]
    packages_extra = BASE_PACKAGES + profile.pkgs_include

    if encryption:
        packages_extra += ['cryptsetup', 'util-linux']  # TODO: select osk-sdl here somehow

    if arch != config.runtime.arch:
        build_enable_qemu_binfmt(arch)

    if local_repos and build_pkgs:
        logging.info("Making sure all packages are built")
        # enforce that local base packages are built
        pkgbuilds = set(filter_pkgbuilds(packages, arch=arch, allow_empty_results=False, use_paths=False))
        # extra packages might be a mix of package names that are in our PKGBUILDs and packages from the base distro
        pkgbuilds |= set(filter_pkgbuilds(packages_extra, arch=arch, allow_empty_results=True, use_paths=False))
        build_packages(pkgbuilds, arch, try_download=not no_download_pkgs)

    sector_size = sector_size or device.get_image_sectorsize()

    image_path = block_target or get_image_path(device, flavour.name)

    makedir(os.path.dirname(image_path))

    logging.info(f'Creating new file at {image_path}')
    create_img_file(image_path, f"{rootfs_size_mb + bootfs_size_mb}M")

    loop_device = losetup_setup_image(image_path, sector_size or device.get_image_sectorsize_default())

    partition_device(loop_device, boot_partition_size=bootfs_size_str)
    partprobe(loop_device)

    boot_dev: str
    root_dev: str
    root_dev_raw: str
    loop_boot = loop_device + 'p1'
    loop_root = loop_device + 'p2'
    if skip_part_images:
        boot_dev = loop_boot
        root_dev = loop_root
    else:
        logging.info('Creating per-partition image files')
        boot_dev = create_img_file(get_image_path(device, flavour, 'boot'), f'{bootfs_size_mb}M')
        root_dev = create_img_file(get_image_path(device, flavour, 'root'), f'{rootfs_size_mb - 200}M')

    root_dev_raw = root_dev

    if encryption:
        encryption_password = encryption_password or profile.encryption_password
        if not encryption_password:
            encryption_password = click.prompt(
                "Please enter your encryption password (input hidden)",
                hide_input=True,
                confirmation_prompt=True,
            )
        luks_create(root_dev, password=encryption_password)
        luks_open(root_dev, mapper_name=encryption_mapper, password=encryption_password)
        root_dev = get_cryptmapper_path(encryption_mapper)

    assert os.path.exists(root_dev)
    create_root_fs(root_dev, sector_size)
    create_boot_fs(boot_dev, sector_size)

    install_rootfs(
        root_dev,
        boot_dev,
        device,
        flavour,
        arch,
        list(set(packages) | set(packages_extra)),
        local_repos,
        profile,
        encrypted=bool(encryption),
    )
    if encryption:
        luks_close(mapper_name=encryption_mapper)

    if not skip_part_images:
        logging.info('Copying partition image files into full image:')
        logging.info(f'Block-copying /boot to {image_path}')
        dd_image(input=boot_dev, output=loop_boot)
        logging.info(f'Block-copying rootfs to {image_path}')
        dd_image(input=root_dev_raw, output=loop_root)

    logging.info(f'Done! Image saved to {image_path}')


@cmd_image.command(name='inspect')
@click.option('--shell', '-s', is_flag=True)
@click.option('--use-local-repos', '-l', is_flag=True)
@sectorsize_option
@encryption_option
@click.argument('profile', required=False)
def cmd_inspect(
    profile: Optional[str] = None,
    shell: bool = False,
    sector_size: Optional[int] = None,
    use_local_repos: bool = False,
    encryption: Optional[bool] = None,
):
    """Loop-mount the device image for inspection."""
    config.enforce_profile_device_set()
    config.enforce_profile_flavour_set()
    enforce_wrap()
    profile_conf = config.get_profile(profile)
    device = get_profile_device(profile)
    arch = device.arch
    flavour = get_profile_flavour(profile).name
    sector_size = sector_size or device.get_image_sectorsize_default()

    chroot = get_device_chroot(device.name, flavour, arch, packages=[], use_local_repos=use_local_repos)
    image_path = get_image_path(device, flavour)
    loop_device = losetup_setup_image(image_path, sector_size)
    partprobe(loop_device)
    mount_chroot(loop_device + 'p2', loop_device + 'p1', chroot, password=profile_conf.encryption_password)

    logging.info(f'Inspect the rootfs image at {chroot.path}')

    if shell:
        chroot.initialized = True
        chroot.activate()
        if arch != config.runtime.arch:
            logging.info('Installing requisites for foreign-arch shell')
            build_enable_qemu_binfmt(arch)
        logging.info('Starting inspection shell')
        chroot.run_cmd('/bin/bash')
    else:
        pause()