Open Menu / test / py / img / common.py
Loading...
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
# SPDX-License-Identifier: GPL-2.0+
# Copyright (c) 2016, NVIDIA CORPORATION. All rights reserved.

"""Common utilities for image creation"""

import gzip
import os

import utils
from fs_helper import DiskHelper, FsHelper


def mkdir_cond(dirname):
    """Create a directory if it doesn't already exist

    Args:
        dirname (str): Name of directory to create
    """
    if not os.path.exists(dirname):
        os.mkdir(dirname)


def copy_partition(ubman, fsfile, outname):
    """Copy a partition into a disk image

    Args:
        ubman (ConsoleBase): U-Boot fixture
        fsfile (str): Name of partition file
        outname (str): Name of full-disk file to update
    """
    utils.run_and_log(ubman,
                      f'dd if={fsfile} of={outname} bs=1M seek=1 conv=notrunc')


def setup_extlinux_image(config, log, devnum, basename, vmlinux, initrd, dtbdir,
                         script, part2_size=1, use_fde=0, luks_kdf='pbkdf2'):
    """Create a 20MB disk image with a single FAT partition

    Args:
        config (ArbitraryAttributeContainer): Configuration
        log (multiplexed_log.Logfile): Log to write to
        devnum (int): Device number to use, e.g. 1
        basename (str): Base name to use in the filename, e.g. 'mmc'
        vmlinux (str): Kernel filename
        initrd (str): Ramdisk filename
        dtbdir (str or None): Devicetree filename
        script (str): Script to place in the extlinux.conf file
        part2_size (int): Size of second partition in MB (default: 1)
        use_fde (int): LUKS version for full-disk encryption (0=none, 1=LUKS1, 2=LUKS2)
        luks_kdf (str): Key derivation function for LUKS2: 'pbkdf2' or 'argon2id'.
            Defaults to 'pbkdf2'. Ignored for LUKS1.
    """
    fsh = FsHelper(config, 'vfat', 18, prefix=basename)
    fsh.setup()

    ext = os.path.join(fsh.srcdir, 'extlinux')
    mkdir_cond(ext)

    conf = os.path.join(ext, 'extlinux.conf')
    with open(conf, 'w', encoding='ascii') as fd:
        print(script, file=fd)

    inf = os.path.join(config.persistent_data_dir, 'inf')
    with open(inf, 'wb') as fd:
        fd.write(gzip.compress(b'vmlinux'))
    mkimage = config.build_dir + '/tools/mkimage'
    utils.run_and_log_no_ubman(
        log, f'{mkimage} -f auto -d {inf} {os.path.join(fsh.srcdir, vmlinux)}')

    with open(os.path.join(fsh.srcdir, initrd), 'w', encoding='ascii') as fd:
        print('initrd', file=fd)

    if dtbdir:
        mkdir_cond(os.path.join(fsh.srcdir, dtbdir))

        dtb_file = os.path.join(fsh.srcdir, f'{dtbdir}/sandbox.dtb')
        utils.run_and_log_no_ubman(
            log, f'dtc -o {dtb_file}', stdin=b'/dts-v1/; / {};')

    fsh.mk_fs()

    img = DiskHelper(config, devnum, basename, True)
    img.add_fs(fsh, DiskHelper.VFAT, bootable=True)

    ext4 = FsHelper(config, 'ext4', max(1, part2_size - 30), prefix=basename,
                    part_mb=part2_size,
                    passphrase='test' if use_fde else None,
                    luks_version=use_fde if use_fde else 2,
                    luks_kdf=luks_kdf)
    ext4.setup()

    bindir = os.path.join(ext4.srcdir, 'bin')
    mkdir_cond(bindir)
    with open(os.path.join(bindir, 'bash'), 'w', encoding='ascii') as fd:
        print('bash', file=fd)

    ext4.mk_fs()

    img.add_fs(ext4, DiskHelper.EXT4)
    img.create()
    fsh.cleanup()