image: add LUKS support and --[no-]encryption CLI flag to build & inspect subcommands

image/image.py

@@ -7,15 +7,21 @@ import subprocess
 from subprocess import CompletedProcess
 from typing import Optional, Union
 
-from config.state import config, Profile
+from chroot.build import BuildChroot, get_build_chroot
 from chroot.device import DeviceChroot, get_device_chroot
-from constants import Arch, POST_INSTALL_CMDS
+from config.state import config, Profile
+from constants import Arch, LUKS_MAPPER_DEFAULT, POST_INSTALL_CMDS
 from distro.distro import get_base_distro, get_kupfer_https
 from devices.device import Device
 from exec.cmd import run_root_cmd, generate_cmd_su
 from exec.file import get_temp_dir, root_write_file, root_makedir
 from flavours.flavour import Flavour
 from net.ssh import copy_ssh_keys
+from utils import programs_available
+
+from .cryptsetup import is_luks, get_luks_offset, luks_close, luks_open
+
+MAPPER_DIR = '/dev/mapper/'
 
 # image files need to be slightly smaller than partitions to fit
 IMG_FILE_ROOT_DEFAULT_SIZE = "1800M"
@@ -72,24 +78,58 @@ def align_bytes(size_bytes: int, alignment: int = 4096) -> int:
     return size_bytes
 
 
-def shrink_fs(loop_device: str, file: str, sector_size: int):
+def shrink_fs(
+    loop_device: str,
+    file: str,
+    sector_size: int,
+    encrypted: Optional[bool] = None,
+    encryption_password: Optional[str] = None,
+    crypt_mapper=LUKS_MAPPER_DEFAULT,
+):
     partprobe(loop_device)
-    logging.debug(f"Checking filesystem at {loop_device}p2")
-    result = run_root_cmd(['e2fsck', '-fy', f'{loop_device}p2'])
+    root_partition = f'{loop_device}p2'
+    root_partition_fs = root_partition
+    if not (encrypted is False):
+        root_partition_fs, native_chroot, encrypted = resolve_rootfs_crypt(
+            root_partition,
+            fail_on_unencrypted=bool(encrypted),
+            crypt_mapper=crypt_mapper,
+            password=encryption_password,
+        )
+    logging.debug(f"Checking filesystem at {root_partition_fs}")
+    result = run_root_cmd(['e2fsck', '-fy', root_partition_fs])
     if result.returncode > 2:
         # https://man7.org/linux/man-pages/man8/e2fsck.8.html#EXIT_CODE
-        raise Exception(f'Failed to e2fsck {loop_device}p2 with exit code {result.returncode}')
+        raise Exception(f'Failed to e2fsck {root_partition_fs} with exit code {result.returncode}')
 
-    logging.info(f'Shrinking filesystem at {loop_device}p2')
-    result = run_root_cmd(['resize2fs', '-M', f'{loop_device}p2'])
+    logging.info(f'Shrinking filesystem at {root_partition_fs}')
+    result = run_root_cmd(['resize2fs', '-M', root_partition_fs])
     if result.returncode != 0:
-        raise Exception(f'Failed to resize2fs {loop_device}p2')
+        raise Exception(f'Failed to resize2fs {root_partition_fs}')
 
-    logging.debug(f'Reading size of shrunken filesystem on {loop_device}p2')
-    fs_blocks, fs_block_size = get_fs_size(f'{loop_device}p2')
+    logging.debug(f'Reading size of shrunken filesystem on {root_partition_fs}')
+    fs_blocks, fs_block_size = get_fs_size(root_partition_fs)
     sectors = bytes_to_sectors(fs_blocks * fs_block_size, sector_size)
+    logging.debug(f"shrunken FS length is {fs_blocks} blocks * {fs_block_size} bytes = {sectors} bytes")
 
-    logging.info(f'Shrinking partition at {loop_device}p2 to {sectors} sectors ({sectors * sector_size} bytes)')
+    _, image_size = find_end_sector(loop_device, root_partition, sector_size)
+    if image_size == -1:
+        raise Exception(f'Failed to find pre-repartition size of {loop_device}')
+
+    if encrypted:
+        if sectors > image_size:
+            raise Exception("Shrunk FS size allegedly larger than the image itself; this is probably "
+                            f"a kupferbootstrap parsing bug. shrunk partition end={sectors}, image size={image_size}, {sector_size=}")
+        old_sectors = sectors
+        luks_offset, luks_sector_size = get_luks_offset(crypt_mapper, native_chroot)
+        #luks_offset_bytes = align_bytes((luks_offset + 1) * luks_sector_size, sector_size)
+        luks_offset_normalized = bytes_to_sectors(luks_offset * luks_sector_size, sector_size)
+        logging.debug(f"Discovered LUKS attrs: {luks_offset=}, {luks_sector_size=}, {luks_offset_normalized=}")
+        luks_close(crypt_mapper, native_chroot)
+        sectors += luks_offset_normalized + 1024
+        logging.debug(f"Increasing sectors from {old_sectors} to {sectors} ({sectors - old_sectors}) to leave space for the LUKS header")
+
+    logging.info(f'Shrinking partition at {root_partition} to {sectors} {sector_size}b sectors ({sectors * sector_size} bytes)')
     child_proccess = subprocess.Popen(
         generate_cmd_su(['fdisk', '-b', str(sector_size), loop_device], switch_user='root'),  # type: ignore
         stdin=subprocess.PIPE,
@@ -113,27 +153,18 @@ def shrink_fs(loop_device: str, file: str, sector_size: int):
         # For some reason re-reading the partition table fails, but that is not a problem
         partprobe(loop_device)
     if returncode > 1:
-        raise Exception(f'Failed to shrink partition size of {loop_device}p2 with fdisk')
+        raise Exception(f'Failed to shrink partition size of {root_partition} with fdisk')
 
     partprobe(loop_device).check_returncode()
 
-    logging.debug(f'Finding end sector of partition at {loop_device}p2')
-    result = run_root_cmd(['fdisk', '-b', str(sector_size), '-l', loop_device], capture_output=True)
-    if result.returncode != 0:
-        print(result.stdout)
-        print(result.stderr)
-        raise Exception(f'Failed to fdisk -l {loop_device}')
+    end_sector, _ = find_end_sector(loop_device, root_partition, sector_size)
+    if end_sector == -1:
+        raise Exception(f'Failed to find end sector of {root_partition}')
 
-    end_sector = 0
-    for line in result.stdout.decode('utf-8').split('\n'):
-        if line.startswith(f'{loop_device}p2'):
-            parts = list(filter(lambda part: part != '', line.split(' ')))
-            end_sector = int(parts[2])
-
-    if end_sector == 0:
-        raise Exception(f'Failed to find end sector of {loop_device}p2')
-
-    end_size = align_bytes((end_sector + 1) * sector_size, 4096)
+    if end_sector > image_size:
+        logging.warning(f"Clipping sectors ({end_sector}) to {image_size=}")
+        end_sector = image_size
+    end_size = align_bytes((end_sector + 1024) * sector_size, 4096)
 
     logging.debug(f'({end_sector} + 1) sectors * {sector_size} bytes/sector = {end_size} bytes')
     logging.info(f'Truncating {file} to {end_size} bytes')
@@ -143,6 +174,26 @@ def shrink_fs(loop_device: str, file: str, sector_size: int):
     partprobe(loop_device)
 
 
+def find_end_sector(device: str, partition: str, sector_size: int) -> tuple[int, int]:
+    """Return (last_sector_index, sector_count) of a partition on a device, returns (-1, -1) if not found"""
+    logging.debug(f'Finding end sector of partition at {partition}')
+    result = run_root_cmd(['fdisk', '-b', str(sector_size), '-l', device], capture_output=True)
+    if result.returncode != 0:
+        print(result.stdout)
+        print(result.stderr)
+        raise Exception(f'Failed to fdisk -l {device}')
+
+    end_sector = -1
+    num_sectors = -1
+    for line in result.stdout.decode('utf-8').split('\n'):
+        if line.startswith(partition):
+            parts = list(filter(lambda part: part != '', line.split(' ')))
+            end_sector = int(parts[2])
+            num_sectors = int(parts[3])
+
+    return end_sector, num_sectors
+
+
 def losetup_destroy(loop_device):
     logging.debug(f'Destroying loop device {loop_device}')
     run_root_cmd(
@@ -210,16 +261,59 @@ def losetup_rootfs_image(image_path: str, sector_size: int) -> str:
     return loop_device
 
 
-def mount_chroot(rootfs_source: str, boot_src: str, chroot: DeviceChroot):
-    logging.debug(f'Mounting {rootfs_source} at {chroot.path}')
+def resolve_rootfs_crypt(
+    rootfs_source: str,
+    password: Optional[str] = None,
+    crypt_mapper: str = LUKS_MAPPER_DEFAULT,
+    native_chroot: Optional[BuildChroot] = None,
+    fail_on_unencrypted: bool = True,
+) -> tuple[str, Optional[BuildChroot], bool]:
+    assert config.runtime.arch
+    is_encrypted = False
+    if not (native_chroot or programs_available(['blkid'])):
+        native_chroot = get_build_chroot(config.runtime.arch, packages=['base', 'util-linux'])
+    if is_luks(rootfs_source, native_chroot=native_chroot):
+        if not (native_chroot or programs_available(['cryptsetup'])):
+            native_chroot = get_build_chroot(config.runtime.arch, packages=['base', 'cryptsetup', 'util-linux'])
+        luks_open(rootfs_source, crypt_mapper, password=password, native_chroot=native_chroot)
+        rootfs_source = f'{MAPPER_DIR}{crypt_mapper}'
+        is_encrypted = True
+    elif fail_on_unencrypted:
+        hint = ''
+        if rootfs_source.startswith(MAPPER_DIR):
+            hint = (f' HINT: path starts with {MAPPER_DIR!r}, probably already a decrypted volume.'
+                    ' This is likely a kupferbootstrap bug.')
+        raise Exception(f"Error: {rootfs_source!r} is not an encrypted LUKS volume.{hint}")
+    return rootfs_source, native_chroot, is_encrypted
 
-    chroot.mount_rootfs(rootfs_source)
-    assert (os.path.ismount(chroot.path))
 
-    root_makedir(chroot.get_path('boot'))
+def mount_chroot(
+    rootfs_source: str,
+    boot_src: str,
+    device_chroot: DeviceChroot,
+    encrypted: Optional[bool] = None,
+    password: Optional[str] = None,
+    native_chroot: Optional[BuildChroot] = None,
+    crypt_mapper: str = LUKS_MAPPER_DEFAULT,
+):
+    if encrypted is not False:
+        rootfs_source, native_chroot, encrypted = resolve_rootfs_crypt(
+            rootfs_source,
+            native_chroot=native_chroot,
+            crypt_mapper=crypt_mapper,
+            fail_on_unencrypted=bool(encrypted),
+            password=password,
+        )
 
-    logging.debug(f'Mounting {boot_src} at {chroot.path}/boot')
-    chroot.mount(boot_src, '/boot', options=['defaults'])
+    logging.debug(f'Mounting {rootfs_source} at {device_chroot.path}')
+
+    device_chroot.mount_rootfs(rootfs_source)
+    assert (os.path.ismount(device_chroot.path))
+
+    root_makedir(device_chroot.get_path('boot'))
+
+    logging.debug(f'Mounting {boot_src} at {device_chroot.path}/boot')
+    device_chroot.mount(boot_src, '/boot', options=['defaults'])
 
 
 def dump_file_from_image(image_path: str, file_path: str, target_path: Optional[str] = None):
@@ -314,18 +408,28 @@ def install_rootfs(
     packages: list[str],
     use_local_repos: bool,
     profile: Profile,
+    encrypted: bool,
 ):
-    user = profile['username'] or 'kupfer'
+    user = profile.username or 'kupfer'
     chroot = get_device_chroot(device=get_device_name(device), flavour=flavour.name, arch=arch, packages=packages, use_local_repos=use_local_repos)
 
-    mount_chroot(rootfs_device, bootfs_device, chroot)
+    # rootfs_device must be passed the crypt_mapper if encrypted is True
+    if encrypted:
+        assert rootfs_device.startswith(MAPPER_DIR)
+
+    mount_chroot(
+        rootfs_device,
+        bootfs_device,
+        chroot,
+        encrypted=False,  # rootfs_device is already the crypt_mapper
+    )
 
     chroot.mount_pacman_cache()
     chroot.initialize()
     chroot.activate()
     chroot.create_user(
         user=user,
-        password=profile['password'],
+        password=profile.password,
     )
     chroot.add_sudo_config(config_name='wheel', privilegee='%wheel', password_required=True)
     copy_ssh_keys(
@@ -338,7 +442,7 @@ def install_rootfs(
             extra_repos=get_kupfer_https(arch).repos,
             in_chroot=True,
         ),
-        'etc/hostname': profile['hostname'] or 'kupfer',
+        'etc/hostname': profile.hostname or 'kupfer',
     }
     for target, content in files.items():
         root_write_file(os.path.join(chroot.path, target.lstrip('/')), content)