# Source code for nobodd.prep

    # nobodd: a boot configuration tool for the Raspberry Pi
    #
    # Copyright (c) 2023-2026 Dave Jones <dave.jones@canonical.com>
    # Copyright (c) 2023-2026 Canonical Ltd.
    #
    # SPDX-License-Identifier: GPL-3.0

    """
    Customizes an OS image to prepare it for netbooting via TFTP. Specifically,
    this expands the image to a specified size (the assumption being the image is a
    copy of a minimally sized template image), then updates the kernel command line
    on the boot partition to point to an NBD server.
    """

    import os
    import sys
    import mmap
    import socket
    import logging
    import argparse
    from pathlib import Path
    from shutil import copyfileobj
    from importlib import resources
    from importlib.metadata import version

    from . import lang
    from .tools import open_file
    from .disk import DiskImage
    from .fs import FatFileSystem
    from .sh import fat_types
    from .config import (
        CONFIG_LOCATIONS,
        ConfigArgumentParser,
        size,
        serial,
        Board,
    )

    def get_parser():
        """
        Returns the command line parser for the application, pre-configured with
        defaults from the application's configuration file(s). See
        :func:`~nobodd.config.ConfigArgumentParser` for more information.
        """
        parser = ConfigArgumentParser(
            description=__doc__,
            template=resources.files('nobodd') / 'default.conf')
        parser.add_argument(
            '--version', action='version', version=version('nobodd'))
        parser.add_argument(
            '-v', '--verbose', dest='log_level',
            action='store_const', const=logging.INFO,
            help=lang._("print more output"))
        parser.add_argument(
            '-q', '--quiet', dest='log_level',
            action='store_const', const=logging.CRITICAL,
            help=lang._("print no output"))

        parser.add_argument(
            'image', type=Path,
            help=lang._("the target image to customize"))
        parser.add_argument(
            '-s', '--size', type=size, default='16GB',
            help=lang._("the size to expand the image to; default: %(default)s"))
        parser.add_argument(
            '--nbd-host', type=str, metavar='HOST', default=socket.getfqdn(),
            help=lang._(
                "the hostname of the nbd server to connect to for the root "
                "device; defaults to the local machine's FQDN"))
        parser.add_argument(
            '--nbd-name', type=str, metavar='NAME', default=None,
            help=lang._(
                "the name of the nbd share to use as the root device; defaults "
                "to the stem of the *image* name"))
        parser.add_argument(
            '--cmdline', type=str, metavar='NAME', default='cmdline.txt',
            help=lang._(
                "the name of the file containing the kernel command line on the "
                "boot partition; default: %(default)s"))
        parser.add_argument(
            '--boot-partition', type=int, metavar='NUM', default=None,
            help=lang._(
                "which partition is the boot partition within the image; "
                "default is the first FAT partition (identified by partition "
                "type) found in the image"))
        parser.add_argument(
            '--root-partition', type=int, metavar='NUM', default=None,
            help=lang._(
                "which partition is the root partition within the image "
                "default is the first non-FAT partition (identified by partition "
                "type) found in the image"))
        parser.add_argument(
            '-C', '--copy', type=Path, metavar='PATH', action='append', default=[],
            help=lang._(
                "copy the specified file or directory into the boot partition. "
                "This may be given multiple times to specify multiple items to "
                "copy"))
        parser.add_argument(
            '-R', '--remove', type=Path, metavar='PATH', action='append', default=[],
            help=lang._(
                "remove the specified file or directory from the boot "
                "partition. This may be given multiple times to specify multiple "
                "items to delete"))
        parser.add_argument(
            '--serial', type=serial, metavar='HEX', default=None,
            help=lang._(
                "defines the serial number of the Raspberry Pi that will be "
                "served this image. When this option is given, a board "
                "configuration compatible with nobodd-tftpd may be output with "
                "--tftpd-conf"))
        parser.add_argument(
            '--tftpd-conf', metavar='FILE', default=None,
            help=lang._(
                "if specified, write a board configuration compatible with "
                "nobodd-tftpd to the specified file; requires --serial to be "
                "given"))
        parser.add_argument(
            '--nbd-conf', metavar='FILE', default=None,
            help=lang._(
                "if specified, write a share configuration compatible with "
                "nbd-server to the specified file"))

        defaults = parser.read_configs(CONFIG_LOCATIONS)
        parser.set_defaults(log_level=logging.WARNING)
        parser.set_defaults_from(defaults)
        return parser

    def prepare_image(conf):
        """
        Given the script's configuration in *conf*, an :class:`argparse.Namespace`,
        resize the target image, and re-write the kernel command line within its
        boot partition to point to the configured NBD server and share.
        """
        with conf.image.open('ab') as f:
            size = f.seek(0, os.SEEK_END)
            if size < conf.size:
                conf.logger.info(
                    lang._('Resizing %s to %d bytes'), conf.image, conf.size)
                f.seek(conf.size)
                f.truncate()
            else:
                conf.logger.info(
                    lang._('Skipping resize; %s is already %d bytes or larger'),
                    conf.image, conf.size)
        with (
            DiskImage(conf.image, access=mmap.ACCESS_WRITE) as img,
            FatFileSystem(img.partitions[conf.boot_partition].data) as fs,
        ):
            remove_items(fs, conf)
            copy_items(fs, conf)
            rewrite_cmdline(fs, conf)

    def remove_items(fs, conf):
        """
        In *fs*, a :class:`~nobodd.fs.FatFileSystem`, remove all items in the
        :class:`list` *conf.remove*, where *conf* is the script's configuration.

        If any item is a directory, it and all files under it will be removed
        recursively. If an item in *to_remove* does not exist, a warning will be
        printed, but no error is raised.
        """
        for item in conf.remove:
            item = fs.root / str(item)
            if item.exists():
                conf.logger.info(
                    lang._('Removing %s from partition %d'),
                    item, conf.boot_partition)
                if item.is_dir():
                    dirs = []
                    for subitem in item.rglob('*'):
                        if subitem.is_dir():
                            dirs.append(subitem)
                        else:
                            subitem.unlink()
                    for subitem in dirs:
                        subitem.rmdir()
                    item.rmdir()
                else:
                    item.unlink()
            else:
                conf.logger.warning(
                    lang._("No such file/dir %s in partition %d"),
                    item, conf.boot_partition)

    def copy_items(fs, conf):
        """
        Copy all :class:`~pathlib.Path` items in the :class:`list` *conf.copy* into
        *fs*, a :class:`~nobodd.fs.FatFileSystem`, where *conf* is the script's
        configuration.

        If an item is a directory, it and all files under it will be copied
        recursively. If an item is a hard-link or a sym-link it will be copied as a
        regular file (since FAT does not support links). If an item does not exist,
        an :exc:`OSError` will be raised. This is in contrast to :func:`to_remove`
        since it is assumed that control over the source file-system is under the
        caller's control, which is not the case in :func:`to_remove`.
        """
        for item in conf.copy:
            conf.logger.info(
                lang._('Copying %s into partition %d'),
                item, conf.boot_partition)
            if item.is_dir():
                copy_root = fs.root / item.name
                copy_root.mkdir(exist_ok=True)
                for subitem in item.rglob('*'):
                    name = subitem.relative_to(item)
                    if subitem.is_dir():
                        (copy_root / str(name)).mkdir(exist_ok=True)
                    else:
                        with (
                            subitem.open('rb') as source,
                            (copy_root / str(name)).open('wb') as target,
                        ):
                            copyfileobj(source, target)
            else:
                with (
                    item.open('rb') as source,
                    (fs.root / item.name).open('wb') as target,
                ):
                    copyfileobj(source, target)

    def rewrite_cmdline(fs, conf):
        """
        Given the script's configuration *conf*, find the file *conf.cmdline*
        containing the kernel command-line in the :class:`~nobodd.fs.FatFileSystem`
        *fs*, and re-write it to point the NBD share specified.
        """
        cmdline = fs.root / conf.cmdline
        conf.logger.info(
            lang._('Re-writing %s in partition %d'),
            conf.cmdline, conf.boot_partition)
        params = cmdline.read_text()
        try:
            params = params[:params.index('\n')]
        except ValueError:
            pass # no newline in the file
        params = [
            param
            for param in params.split()
            if not param.startswith('root=')
        ]
        params[:0] = [
            'ip=dhcp',
            f'nbdroot={conf.nbd_host}/{conf.nbd_name}',
            f'root=/dev/nbd0p{conf.root_partition}',
        ]
        cmdline.write_text(' '.join(params))

    def detect_partitions(conf):
        """
        Given the script's configuration in *conf*, an :class:`argparse.Namespace`,
        open the target image, and attempt to detect the root and/or boot
        partition.
        """
        conf.logger.info(lang._('Detecting partitions'))
        with (
            conf.image.open('rb') as img_file,
            DiskImage(img_file) as img,
        ):
            for num, fat_type in fat_types(img):
                if fat_type.startswith('fat') and conf.boot_partition is None:
                    conf.boot_partition = num
                    conf.logger.info(
                        lang._('Boot partition is %d (%s)'),
                        conf.boot_partition, fat_type)
                elif fat_type == 'notfat' and conf.root_partition is None:
                    conf.root_partition = num
                    conf.logger.info(
                        lang._('Root partition is %d'),
                        conf.root_partition)
                if conf.boot_partition is not None:
                    if conf.root_partition is not None:
                        break
        if conf.boot_partition is None:
            raise ValueError(lang._('Unable to detect boot partition'))
        if conf.root_partition is None:
            raise ValueError(lang._('Unable to detect root partition'))

    def main(args=None):
        """
        The main entry point for the :program:`nobodd-prep` application. Takes
        *args*, the sequence of command line arguments to parse. Returns the exit
        code of the application (0 for a normal exit, and non-zero otherwise).

        If ``DEBUG=1`` is found in the application's environment, top-level
        exceptions will be printed with a full back-trace. ``DEBUG=2`` will launch
        PDB in post-mortem mode.
        """
        try:
            debug = int(os.environ['DEBUG'])
        except (KeyError, ValueError):
            debug = 0
        lang.init()

        try:
            conf = get_parser().parse_args(args)
            conf.image = conf.image.resolve()
            conf.logger = logging.getLogger('prep')
            conf.logger.addHandler(logging.StreamHandler(sys.stderr))
            conf.logger.setLevel(logging.DEBUG if debug else conf.log_level)
            if conf.boot_partition is None or conf.root_partition is None:
                detect_partitions(conf)
            if conf.nbd_name is None:
                conf.nbd_name = conf.image.stem

            with (
                open_file(conf.tftpd_conf) as tftpd_conf,
                open_file(conf.nbd_conf) as nbd_conf,
            ):
                prepare_image(conf)
                if tftpd_conf is not None and conf.serial is not None:
                    board = Board(conf.serial, conf.image, conf.boot_partition, None)
                    tftpd_conf.write(str(board))
                    tftpd_conf.write('\n')
                if nbd_conf is not None:
                    nbd_conf.write(f"[{conf.nbd_name}]\n")
                    nbd_conf.write(f"exportname = {conf.image}\n")
        except Exception as e:
            if not debug:
                print(str(e), file=sys.stderr)
                return 1
            elif debug == 1:
                raise
            else:
                import pdb
                pdb.post_mortem()
        else:
            return 0