image, flash: refactor get_image_{name,path} to not rely on a chroot instance, introduce per-partition image files

flash.py

@@ -6,9 +6,8 @@ import click
 import tempfile
 
 from constants import FLASH_PARTS, LOCATIONS
-from chroot import get_device_chroot
 from fastboot import fastboot_flash
-from image import partprobe, shrink_fs, losetup_rootfs_image, dump_aboot, dump_lk2nd, dump_qhypstub, get_device_and_flavour, get_image_name, get_image_path
+from image import dd_image, partprobe, shrink_fs, losetup_rootfs_image, dump_aboot, dump_lk2nd, dump_qhypstub, get_device_and_flavour, get_image_name, get_image_path
 from wrapper import enforce_wrap
 
 ABOOT = FLASH_PARTS['ABOOT']
@@ -23,9 +22,8 @@ ROOTFS = FLASH_PARTS['ROOTFS']
 def cmd_flash(what, location):
     enforce_wrap()
     device, flavour = get_device_and_flavour()
-    chroot = get_device_chroot(device, flavour, 'aarch64')
+    device_image_name = get_image_name(device, flavour)
-    device_image_name = get_image_name(chroot)
+    device_image_path = get_image_path(device, flavour)
-    device_image_path = get_image_path(chroot)
 
     # TODO: PARSE DEVICE SECTOR SIZE
     sector_size = 4096
@@ -73,16 +71,8 @@ def cmd_flash(what, location):
         partprobe(loop_device)
         shrink_fs(loop_device, minimal_image_path, sector_size)
 
-        result = subprocess.run([
+        result = dd_image(input=minimal_image_path, output=path)
-            'dd',
+
-            f'if={minimal_image_path}',
-            f'of={path}',
-            'bs=20M',
-            'iflag=direct',
-            'oflag=direct',
-            'status=progress',
-            'conv=sync,noerror',
-        ])
         if result.returncode != 0:
             raise Exception(f'Failed to flash {minimal_image_path} to {path}')
     else:

image.py

@@ -7,7 +7,7 @@ import click
 import logging
 
 from signal import pause
-from subprocess import run
+from subprocess import run, CompletedProcess
 
 from chroot import Chroot, get_device_chroot
 from constants import BASE_PACKAGES, DEVICES, FLAVOURS
@@ -17,6 +17,23 @@ from packages import build_enable_qemu_binfmt, discover_packages, build_packages
 from ssh import copy_ssh_keys
 from wrapper import enforce_wrap
 
+# image files need to be slightly smaller than partitions to fit
+IMG_FILE_ROOT_DEFAULT_SIZE = "1800M"
+IMG_FILE_BOOT_DEFAULT_SIZE = "90M"
+
+
+def dd_image(input: str, output: str, blocksize='1M') -> CompletedProcess:
+    return subprocess.run([
+        'dd',
+        f'if={input}',
+        f'of={output}',
+        f'bs={blocksize}',
+        'iflag=direct',
+        'oflag=direct',
+        'status=progress',
+        'conv=sync,noerror',
+    ])
+
 
 def partprobe(device: str):
     return subprocess.run(['partprobe', device])
@@ -109,12 +126,12 @@ def get_device_and_flavour(profile: str = None) -> tuple[str, str]:
     return (profile['device'], profile['flavour'])
 
 
-def get_image_name(device_chroot: Chroot) -> str:
+def get_image_name(device, flavour, img_type='full') -> str:
-    return f'{device_chroot.name}.img'
+    return f'{device}-{flavour}-{img_type}.img'
 
 
-def get_image_path(device_chroot: Chroot) -> str:
+def get_image_path(device, flavour, img_type='full') -> str:
-    return os.path.join(config.get_path('images'), get_image_name(device_chroot))
+    return os.path.join(config.get_path('images'), get_image_name(device, flavour, img_type))
 
 
 def losetup_rootfs_image(image_path: str, sector_size: int) -> str:
@@ -224,56 +241,19 @@ def dump_qhypstub(image_path: str) -> str:
     return path
 
 
-@click.group(name='image')
+def create_img_file(image_path: str, size_str: str):
-def cmd_image():
-    pass
-
-
-@cmd_image.command(name='build')
-@click.argument('profile_name', required=False)
-@click.option('--build-pkgs/--no-build-pkgs', '-p/-P', default=True, help='Whether to build missing/outdated packages. Defaults to true.')
-def cmd_build(profile_name: str = None, build_pkgs: bool = True):
-    enforce_wrap()
-    profile = config.get_profile(profile_name)
-    device, flavour = get_device_and_flavour(profile_name)
-    post_cmds = FLAVOURS[flavour].get('post_cmds', [])
-
-    user = profile['username'] or 'kupfer'
-    # TODO: PARSE DEVICE ARCH AND SECTOR SIZE
-    arch = 'aarch64'
-    sector_size = 4096
-
-    build_enable_qemu_binfmt(arch)
-
-    packages_dir = config.get_package_dir(arch)
-    if os.path.exists(os.path.join(packages_dir, 'main')):
-        extra_repos = get_kupfer_local(arch).repos
-    else:
-        extra_repos = get_kupfer_https(arch).repos
-    packages = BASE_PACKAGES + DEVICES[device] + FLAVOURS[flavour]['packages'] + profile['pkgs_include']
-
-    if build_pkgs:
-        repo = discover_packages()
-        build_packages(repo, [p for name, p in repo.items() if name in packages], arch)
-
-    chroot = get_device_chroot(device=device, flavour=flavour, arch=arch, packages=packages, extra_repos=extra_repos)
-    image_path = get_image_path(chroot)
-
-    os.makedirs(config.get_path('images'), exist_ok=True)
-    new_image = not os.path.exists(image_path)
-    if new_image:
     result = subprocess.run([
         'truncate',
         '-s',
-            f"{FLAVOURS[flavour].get('size',2)}G",
+        size_str,
         image_path,
     ])
     if result.returncode != 0:
         raise Exception(f'Failed to allocate {image_path}')
+    return image_path
 
-    loop_device = losetup_rootfs_image(image_path, sector_size)
 
-    if new_image:
+def partition_device(device: str):
     boot_partition_size = '100MiB'
     create_partition_table = ['mklabel', 'msdos']
     create_boot_partition = ['mkpart', 'primary', 'ext2', '0%', boot_partition_size]
@@ -282,21 +262,13 @@ def cmd_build(profile_name: str = None, build_pkgs: bool = True):
     result = subprocess.run([
         'parted',
         '--script',
-            loop_device,
+        device,
     ] + create_partition_table + create_boot_partition + create_root_partition + enable_boot)
     if result.returncode != 0:
-            raise Exception(f'Failed to create partitions on {loop_device}')
+        raise Exception(f'Failed to create partitions on {device}')
-        partprobe(loop_device)
-        result = subprocess.run([
-            'mkfs.ext2',
-            '-F',
-            '-L',
-            'kupfer_boot',
-            f'{loop_device}p1',
-        ])
-        if result.returncode != 0:
-            raise Exception(f'Failed to create ext2 filesystem on {loop_device}p1')
 
+
+def create_root_fs(device: str):
     result = subprocess.run([
         'mkfs.ext4',
         '-O',
@@ -306,12 +278,30 @@ def cmd_build(profile_name: str = None, build_pkgs: bool = True):
         'kupfer_root',
         '-N',
         '100000',
-            f'{loop_device}p2',
+        device,
     ])
     if result.returncode != 0:
-            raise Exception(f'Failed to create ext4 filesystem on {loop_device}p2')
+        raise Exception(f'Failed to create ext4 filesystem on {device}')
 
-    mount_chroot(loop_device + 'p2', loop_device + 'p1', chroot)
+
+def create_boot_fs(device: str):
+    result = subprocess.run([
+        'mkfs.ext2',
+        '-F',
+        '-L',
+        'kupfer_boot',
+        device,
+    ])
+    if result.returncode != 0:
+        raise Exception(f'Failed to create ext2 filesystem on {device}')
+
+
+def install_rootfs(rootfs_device: str, bootfs_device: str, device, flavour, arch, packages, extra_repos, profile):
+    user = profile['username'] or 'kupfer'
+    post_cmds = FLAVOURS[flavour].get('post_cmds', [])
+    chroot = get_device_chroot(device=device, flavour=flavour, arch=arch, packages=packages, extra_repos=extra_repos)
+
+    mount_chroot(rootfs_device, bootfs_device, chroot)
 
     chroot.mount_pacman_cache()
     chroot.initialize()
@@ -341,6 +331,85 @@ def cmd_build(profile_name: str = None, build_pkgs: bool = True):
     res = run(['umount', chroot.path])
     logging.debug(f'rc: {res.returncode}')
 
+
+@click.group(name='image')
+def cmd_image():
+    pass
+
+
+@cmd_image.command(name='build')
+@click.argument('profile_name', required=False)
+@click.option('--build-pkgs/--no-build-pkgs', '-p/-P', default=True, help='Whether to build missing/outdated packages. Defaults to true.')
+@click.option('--block-target', default=None, help='Override the block device file to target')
+@click.option('--skip-part-images', default=False, help='Skip creating image files for the partitions and directly work on the target block device.')
+def cmd_build(profile_name: str = None, build_pkgs: bool = True, block_target: str = None, skip_part_images: bool = False):
+    enforce_wrap()
+    profile = config.get_profile(profile_name)
+    device, flavour = get_device_and_flavour(profile_name)
+
+    # TODO: PARSE DEVICE ARCH AND SECTOR SIZE
+    arch = 'aarch64'
+    sector_size = 4096
+    rootfs_size_gb = FLAVOURS[flavour].get('size', 2)
+
+    build_enable_qemu_binfmt(arch)
+
+    packages_dir = config.get_package_dir(arch)
+    if os.path.exists(os.path.join(packages_dir, 'main')):
+        extra_repos = get_kupfer_local(arch).repos
+    else:
+        extra_repos = get_kupfer_https(arch).repos
+    packages = BASE_PACKAGES + DEVICES[device] + FLAVOURS[flavour]['packages'] + profile['pkgs_include']
+
+    if build_pkgs:
+        repo = discover_packages()
+        build_packages(repo, [p for name, p in repo.items() if name in packages], arch)
+
+    image_path = block_target or get_image_path(device, flavour)
+
+    os.makedirs(os.path.dirname(image_path), exist_ok=True)
+    new_image = not os.path.exists(image_path)
+    if new_image:
+        logging.info(f'Creating new file at {image_path}')
+        create_img_file(image_path, f"{rootfs_size_gb}G")
+
+    loop_device = losetup_rootfs_image(image_path, sector_size)
+
+    partition_device(loop_device)
+    partprobe(loop_device)
+
+    boot_dev, root_dev = None, None
+    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'), IMG_FILE_BOOT_DEFAULT_SIZE)
+        root_dev = create_img_file(get_image_path(device, flavour, 'root'), f'{rootfs_size_gb * 1000 - 200}M')
+
+    create_boot_fs(boot_dev)
+    create_root_fs(root_dev)
+
+    install_rootfs(
+        root_dev,
+        boot_dev,
+        device,
+        flavour,
+        arch,
+        packages,
+        extra_repos,
+        profile,
+    )
+
+    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, output=loop_root)
+
     logging.info(f'Done! Image saved to {image_path}')
 
 
@@ -354,7 +423,7 @@ def cmd_inspect(shell: bool = False):
     # TODO: PARSE DEVICE SECTOR SIZE
     sector_size = 4096
     chroot = get_device_chroot(device, flavour, arch)
-    image_path = get_image_path(chroot)
+    image_path = get_image_path(device, flavour)
     loop_device = losetup_rootfs_image(image_path, sector_size)
     partprobe(loop_device)
     mount_chroot(loop_device + 'p2', loop_device + 'p1', chroot)