1#
2# Copyright (c) 2013, Intel Corporation.
3#
4# SPDX-License-Identifier: GPL-2.0-only
5#
6# DESCRIPTION
7# This implements the 'direct' imager plugin class for 'wic'
8#
9# AUTHORS
10# Tom Zanussi <tom.zanussi (at] linux.intel.com>
11#
12
13import logging
14import os
15import random
16import shutil
17import tempfile
18import uuid
19
20from time import strftime
21
22from oe.path import copyhardlinktree
23
24from wic import WicError
25from wic.filemap import sparse_copy
26from wic.ksparser import KickStart, KickStartError
27from wic.pluginbase import PluginMgr, ImagerPlugin
28from wic.misc import get_bitbake_var, exec_cmd, exec_native_cmd
29
30logger = logging.getLogger('wic')
31
32class DirectPlugin(ImagerPlugin):
33    """
34    Install a system into a file containing a partitioned disk image.
35
36    An image file is formatted with a partition table, each partition
37    created from a rootfs or other OpenEmbedded build artifact and dd'ed
38    into the virtual disk. The disk image can subsequently be dd'ed onto
39    media and used on actual hardware.
40    """
41    name = 'direct'
42
43    def __init__(self, wks_file, rootfs_dir, bootimg_dir, kernel_dir,
44                 native_sysroot, oe_builddir, options):
45        try:
46            self.ks = KickStart(wks_file)
47        except KickStartError as err:
48            raise WicError(str(err))
49
50        # parse possible 'rootfs=name' items
51        self.rootfs_dir = dict(rdir.split('=') for rdir in rootfs_dir.split(' '))
52        self.bootimg_dir = bootimg_dir
53        self.kernel_dir = kernel_dir
54        self.native_sysroot = native_sysroot
55        self.oe_builddir = oe_builddir
56
57        self.debug = options.debug
58        self.outdir = options.outdir
59        self.compressor = options.compressor
60        self.bmap = options.bmap
61        self.no_fstab_update = options.no_fstab_update
62        self.updated_fstab_path = None
63
64        self.name = "%s-%s" % (os.path.splitext(os.path.basename(wks_file))[0],
65                               strftime("%Y%m%d%H%M"))
66        self.workdir = self.setup_workdir(options.workdir)
67        self._image = None
68        self.ptable_format = self.ks.bootloader.ptable
69        self.parts = self.ks.partitions
70
71        # as a convenience, set source to the boot partition source
72        # instead of forcing it to be set via bootloader --source
73        for part in self.parts:
74            if not self.ks.bootloader.source and part.mountpoint == "/boot":
75                self.ks.bootloader.source = part.source
76                break
77
78        image_path = self._full_path(self.workdir, self.parts[0].disk, "direct")
79        self._image = PartitionedImage(image_path, self.ptable_format,
80                                       self.parts, self.native_sysroot,
81                                       options.extra_space)
82
83    def setup_workdir(self, workdir):
84        if workdir:
85            if os.path.exists(workdir):
86                raise WicError("Internal workdir '%s' specified in wic arguments already exists!" % (workdir))
87
88            os.makedirs(workdir)
89            return workdir
90        else:
91            return tempfile.mkdtemp(dir=self.outdir, prefix='tmp.wic.')
92
93    def do_create(self):
94        """
95        Plugin entry point.
96        """
97        try:
98            self.create()
99            self.assemble()
100            self.finalize()
101            self.print_info()
102        finally:
103            self.cleanup()
104
105    def update_fstab(self, image_rootfs):
106        """Assume partition order same as in wks"""
107        if not image_rootfs:
108            return
109
110        fstab_path = image_rootfs + "/etc/fstab"
111        if not os.path.isfile(fstab_path):
112            return
113
114        with open(fstab_path) as fstab:
115            fstab_lines = fstab.readlines()
116
117        updated = False
118        for part in self.parts:
119            if not part.realnum or not part.mountpoint \
120               or part.mountpoint == "/" or not part.mountpoint.startswith('/'):
121                continue
122
123            if part.use_uuid:
124                if part.fsuuid:
125                    # FAT UUID is different from others
126                    if len(part.fsuuid) == 10:
127                        device_name = "UUID=%s-%s" % \
128                                       (part.fsuuid[2:6], part.fsuuid[6:])
129                    else:
130                        device_name = "UUID=%s" % part.fsuuid
131                else:
132                    device_name = "PARTUUID=%s" % part.uuid
133            elif part.use_label:
134                device_name = "LABEL=%s" % part.label
135            else:
136                # mmc device partitions are named mmcblk0p1, mmcblk0p2..
137                prefix = 'p' if  part.disk.startswith('mmcblk') else ''
138                device_name = "/dev/%s%s%d" % (part.disk, prefix, part.realnum)
139
140            opts = part.fsopts if part.fsopts else "defaults"
141            line = "\t".join([device_name, part.mountpoint, part.fstype,
142                              opts, "0", "0"]) + "\n"
143
144            fstab_lines.append(line)
145            updated = True
146
147        if updated:
148            self.updated_fstab_path = os.path.join(self.workdir, "fstab")
149            with open(self.updated_fstab_path, "w") as f:
150                f.writelines(fstab_lines)
151
152    def _full_path(self, path, name, extention):
153        """ Construct full file path to a file we generate. """
154        return os.path.join(path, "%s-%s.%s" % (self.name, name, extention))
155
156    #
157    # Actual implemention
158    #
159    def create(self):
160        """
161        For 'wic', we already have our build artifacts - we just create
162        filesystems from the artifacts directly and combine them into
163        a partitioned image.
164        """
165        if not self.no_fstab_update:
166            self.update_fstab(self.rootfs_dir.get("ROOTFS_DIR"))
167
168        for part in self.parts:
169            # get rootfs size from bitbake variable if it's not set in .ks file
170            if not part.size:
171                # and if rootfs name is specified for the partition
172                image_name = self.rootfs_dir.get(part.rootfs_dir)
173                if image_name and os.path.sep not in image_name:
174                    # Bitbake variable ROOTFS_SIZE is calculated in
175                    # Image._get_rootfs_size method from meta/lib/oe/image.py
176                    # using IMAGE_ROOTFS_SIZE, IMAGE_ROOTFS_ALIGNMENT,
177                    # IMAGE_OVERHEAD_FACTOR and IMAGE_ROOTFS_EXTRA_SPACE
178                    rsize_bb = get_bitbake_var('ROOTFS_SIZE', image_name)
179                    if rsize_bb:
180                        part.size = int(round(float(rsize_bb)))
181
182        self._image.prepare(self)
183        self._image.layout_partitions()
184        self._image.create()
185
186    def assemble(self):
187        """
188        Assemble partitions into disk image
189        """
190        self._image.assemble()
191
192    def finalize(self):
193        """
194        Finalize the disk image.
195
196        For example, prepare the image to be bootable by e.g.
197        creating and installing a bootloader configuration.
198        """
199        source_plugin = self.ks.bootloader.source
200        disk_name = self.parts[0].disk
201        if source_plugin:
202            plugin = PluginMgr.get_plugins('source')[source_plugin]
203            plugin.do_install_disk(self._image, disk_name, self, self.workdir,
204                                   self.oe_builddir, self.bootimg_dir,
205                                   self.kernel_dir, self.native_sysroot)
206
207        full_path = self._image.path
208        # Generate .bmap
209        if self.bmap:
210            logger.debug("Generating bmap file for %s", disk_name)
211            python = os.path.join(self.native_sysroot, 'usr/bin/python3-native/python3')
212            bmaptool = os.path.join(self.native_sysroot, 'usr/bin/bmaptool')
213            exec_native_cmd("%s %s create %s -o %s.bmap" % \
214                            (python, bmaptool, full_path, full_path), self.native_sysroot)
215        # Compress the image
216        if self.compressor:
217            logger.debug("Compressing disk %s with %s", disk_name, self.compressor)
218            exec_cmd("%s %s" % (self.compressor, full_path))
219
220    def print_info(self):
221        """
222        Print the image(s) and artifacts used, for the user.
223        """
224        msg = "The new image(s) can be found here:\n"
225
226        extension = "direct" + {"gzip": ".gz",
227                                "bzip2": ".bz2",
228                                "xz": ".xz",
229                                None: ""}.get(self.compressor)
230        full_path = self._full_path(self.outdir, self.parts[0].disk, extension)
231        msg += '  %s\n\n' % full_path
232
233        msg += 'The following build artifacts were used to create the image(s):\n'
234        for part in self.parts:
235            if part.rootfs_dir is None:
236                continue
237            if part.mountpoint == '/':
238                suffix = ':'
239            else:
240                suffix = '["%s"]:' % (part.mountpoint or part.label)
241            rootdir = part.rootfs_dir
242            msg += '  ROOTFS_DIR%s%s\n' % (suffix.ljust(20), rootdir)
243
244        msg += '  BOOTIMG_DIR:                  %s\n' % self.bootimg_dir
245        msg += '  KERNEL_DIR:                   %s\n' % self.kernel_dir
246        msg += '  NATIVE_SYSROOT:               %s\n' % self.native_sysroot
247
248        logger.info(msg)
249
250    @property
251    def rootdev(self):
252        """
253        Get root device name to use as a 'root' parameter
254        in kernel command line.
255
256        Assume partition order same as in wks
257        """
258        for part in self.parts:
259            if part.mountpoint == "/":
260                if part.uuid:
261                    return "PARTUUID=%s" % part.uuid
262                elif part.label and self.ptable_format != 'msdos':
263                    return "PARTLABEL=%s" % part.label
264                else:
265                    suffix = 'p' if part.disk.startswith('mmcblk') else ''
266                    return "/dev/%s%s%-d" % (part.disk, suffix, part.realnum)
267
268    def cleanup(self):
269        if self._image:
270            self._image.cleanup()
271
272        # Move results to the output dir
273        if not os.path.exists(self.outdir):
274            os.makedirs(self.outdir)
275
276        for fname in os.listdir(self.workdir):
277            path = os.path.join(self.workdir, fname)
278            if os.path.isfile(path):
279                shutil.move(path, os.path.join(self.outdir, fname))
280
281        # remove work directory when it is not in debugging mode
282        if not self.debug:
283            shutil.rmtree(self.workdir, ignore_errors=True)
284
285# Overhead of the MBR partitioning scheme (just one sector)
286MBR_OVERHEAD = 1
287
288# Overhead of the GPT partitioning scheme
289GPT_OVERHEAD = 34
290
291# Size of a sector in bytes
292SECTOR_SIZE = 512
293
294class PartitionedImage():
295    """
296    Partitioned image in a file.
297    """
298
299    def __init__(self, path, ptable_format, partitions, native_sysroot=None, extra_space=0):
300        self.path = path  # Path to the image file
301        self.numpart = 0  # Number of allocated partitions
302        self.realpart = 0 # Number of partitions in the partition table
303        self.primary_part_num = 0  # Number of primary partitions (msdos)
304        self.extendedpart = 0      # Create extended partition before this logical partition (msdos)
305        self.extended_size_sec = 0 # Size of exteded partition (msdos)
306        self.logical_part_cnt = 0  # Number of total logical paritions (msdos)
307        self.offset = 0   # Offset of next partition (in sectors)
308        self.min_size = 0 # Minimum required disk size to fit
309                          # all partitions (in bytes)
310        self.ptable_format = ptable_format  # Partition table format
311        # Disk system identifier
312        self.identifier = random.SystemRandom().randint(1, 0xffffffff)
313
314        self.partitions = partitions
315        self.partimages = []
316        # Size of a sector used in calculations
317        self.sector_size = SECTOR_SIZE
318        self.native_sysroot = native_sysroot
319        num_real_partitions = len([p for p in self.partitions if not p.no_table])
320        self.extra_space = extra_space
321
322        # calculate the real partition number, accounting for partitions not
323        # in the partition table and logical partitions
324        realnum = 0
325        for part in self.partitions:
326            if part.no_table:
327                part.realnum = 0
328            else:
329                realnum += 1
330                if self.ptable_format == 'msdos' and realnum > 3 and num_real_partitions > 4:
331                    part.realnum = realnum + 1
332                    continue
333                part.realnum = realnum
334
335        # generate parition and filesystem UUIDs
336        for part in self.partitions:
337            if not part.uuid and part.use_uuid:
338                if self.ptable_format == 'gpt':
339                    part.uuid = str(uuid.uuid4())
340                else: # msdos partition table
341                    part.uuid = '%08x-%02d' % (self.identifier, part.realnum)
342            if not part.fsuuid:
343                if part.fstype == 'vfat' or part.fstype == 'msdos':
344                    part.fsuuid = '0x' + str(uuid.uuid4())[:8].upper()
345                else:
346                    part.fsuuid = str(uuid.uuid4())
347            else:
348                #make sure the fsuuid for vfat/msdos align with format 0xYYYYYYYY
349                if part.fstype == 'vfat' or part.fstype == 'msdos':
350                    if part.fsuuid.upper().startswith("0X"):
351                        part.fsuuid = '0x' + part.fsuuid.upper()[2:].rjust(8,"0")
352                    else:
353                        part.fsuuid = '0x' + part.fsuuid.upper().rjust(8,"0")
354
355    def prepare(self, imager):
356        """Prepare an image. Call prepare method of all image partitions."""
357        for part in self.partitions:
358            # need to create the filesystems in order to get their
359            # sizes before we can add them and do the layout.
360            part.prepare(imager, imager.workdir, imager.oe_builddir,
361                         imager.rootfs_dir, imager.bootimg_dir,
362                         imager.kernel_dir, imager.native_sysroot,
363                         imager.updated_fstab_path)
364
365            # Converting kB to sectors for parted
366            part.size_sec = part.disk_size * 1024 // self.sector_size
367
368    def layout_partitions(self):
369        """ Layout the partitions, meaning calculate the position of every
370        partition on the disk. The 'ptable_format' parameter defines the
371        partition table format and may be "msdos". """
372
373        logger.debug("Assigning %s partitions to disks", self.ptable_format)
374
375        # The number of primary and logical partitions. Extended partition and
376        # partitions not listed in the table are not included.
377        num_real_partitions = len([p for p in self.partitions if not p.no_table])
378
379        # Go through partitions in the order they are added in .ks file
380        for num in range(len(self.partitions)):
381            part = self.partitions[num]
382
383            if self.ptable_format == 'msdos' and part.part_name:
384                raise WicError("setting custom partition name is not " \
385                               "implemented for msdos partitions")
386
387            if self.ptable_format == 'msdos' and part.part_type:
388                # The --part-type can also be implemented for MBR partitions,
389                # in which case it would map to the 1-byte "partition type"
390                # filed at offset 3 of the partition entry.
391                raise WicError("setting custom partition type is not " \
392                               "implemented for msdos partitions")
393
394            # Get the disk where the partition is located
395            self.numpart += 1
396            if not part.no_table:
397                self.realpart += 1
398
399            if self.numpart == 1:
400                if self.ptable_format == "msdos":
401                    overhead = MBR_OVERHEAD
402                elif self.ptable_format == "gpt":
403                    overhead = GPT_OVERHEAD
404
405                # Skip one sector required for the partitioning scheme overhead
406                self.offset += overhead
407
408            if self.ptable_format == "msdos":
409                if self.primary_part_num > 3 or \
410                   (self.extendedpart == 0 and self.primary_part_num >= 3 and num_real_partitions > 4):
411                    part.type = 'logical'
412                # Reserve a sector for EBR for every logical partition
413                # before alignment is performed.
414                if part.type == 'logical':
415                    self.offset += 2
416
417            align_sectors = 0
418            if part.align:
419                # If not first partition and we do have alignment set we need
420                # to align the partition.
421                # FIXME: This leaves a empty spaces to the disk. To fill the
422                # gaps we could enlargea the previous partition?
423
424                # Calc how much the alignment is off.
425                align_sectors = self.offset % (part.align * 1024 // self.sector_size)
426
427                if align_sectors:
428                    # If partition is not aligned as required, we need
429                    # to move forward to the next alignment point
430                    align_sectors = (part.align * 1024 // self.sector_size) - align_sectors
431
432                    logger.debug("Realignment for %s%s with %s sectors, original"
433                                 " offset %s, target alignment is %sK.",
434                                 part.disk, self.numpart, align_sectors,
435                                 self.offset, part.align)
436
437                    # increase the offset so we actually start the partition on right alignment
438                    self.offset += align_sectors
439
440            if part.offset is not None:
441                offset = part.offset // self.sector_size
442
443                if offset * self.sector_size != part.offset:
444                    raise WicError("Could not place %s%s at offset %d with sector size %d" % (part.disk, self.numpart, part.offset, self.sector_size))
445
446                delta = offset - self.offset
447                if delta < 0:
448                    raise WicError("Could not place %s%s at offset %d: next free sector is %d (delta: %d)" % (part.disk, self.numpart, part.offset, self.offset, delta))
449
450                logger.debug("Skipping %d sectors to place %s%s at offset %dK",
451                             delta, part.disk, self.numpart, part.offset)
452
453                self.offset = offset
454
455            part.start = self.offset
456            self.offset += part.size_sec
457
458            if not part.no_table:
459                part.num = self.realpart
460            else:
461                part.num = 0
462
463            if self.ptable_format == "msdos" and not part.no_table:
464                if part.type == 'logical':
465                    self.logical_part_cnt += 1
466                    part.num = self.logical_part_cnt + 4
467                    if self.extendedpart == 0:
468                        # Create extended partition as a primary partition
469                        self.primary_part_num += 1
470                        self.extendedpart = part.num
471                    else:
472                        self.extended_size_sec += align_sectors
473                    self.extended_size_sec += part.size_sec + 2
474                else:
475                    self.primary_part_num += 1
476                    part.num = self.primary_part_num
477
478            logger.debug("Assigned %s to %s%d, sectors range %d-%d size %d "
479                         "sectors (%d bytes).", part.mountpoint, part.disk,
480                         part.num, part.start, self.offset - 1, part.size_sec,
481                         part.size_sec * self.sector_size)
482
483        # Once all the partitions have been layed out, we can calculate the
484        # minumim disk size
485        self.min_size = self.offset
486        if self.ptable_format == "gpt":
487            self.min_size += GPT_OVERHEAD
488
489        self.min_size *= self.sector_size
490        self.min_size += self.extra_space
491
492    def _create_partition(self, device, parttype, fstype, start, size):
493        """ Create a partition on an image described by the 'device' object. """
494
495        # Start is included to the size so we need to substract one from the end.
496        end = start + size - 1
497        logger.debug("Added '%s' partition, sectors %d-%d, size %d sectors",
498                     parttype, start, end, size)
499
500        cmd = "parted -s %s unit s mkpart %s" % (device, parttype)
501        if fstype:
502            cmd += " %s" % fstype
503        cmd += " %d %d" % (start, end)
504
505        return exec_native_cmd(cmd, self.native_sysroot)
506
507    def create(self):
508        logger.debug("Creating sparse file %s", self.path)
509        with open(self.path, 'w') as sparse:
510            os.ftruncate(sparse.fileno(), self.min_size)
511
512        logger.debug("Initializing partition table for %s", self.path)
513        exec_native_cmd("parted -s %s mklabel %s" %
514                        (self.path, self.ptable_format), self.native_sysroot)
515
516        logger.debug("Set disk identifier %x", self.identifier)
517        with open(self.path, 'r+b') as img:
518            img.seek(0x1B8)
519            img.write(self.identifier.to_bytes(4, 'little'))
520
521        logger.debug("Creating partitions")
522
523        for part in self.partitions:
524            if part.num == 0:
525                continue
526
527            if self.ptable_format == "msdos" and part.num == self.extendedpart:
528                # Create an extended partition (note: extended
529                # partition is described in MBR and contains all
530                # logical partitions). The logical partitions save a
531                # sector for an EBR just before the start of a
532                # partition. The extended partition must start one
533                # sector before the start of the first logical
534                # partition. This way the first EBR is inside of the
535                # extended partition. Since the extended partitions
536                # starts a sector before the first logical partition,
537                # add a sector at the back, so that there is enough
538                # room for all logical partitions.
539                self._create_partition(self.path, "extended",
540                                       None, part.start - 2,
541                                       self.extended_size_sec)
542
543            if part.fstype == "swap":
544                parted_fs_type = "linux-swap"
545            elif part.fstype == "vfat":
546                parted_fs_type = "fat32"
547            elif part.fstype == "msdos":
548                parted_fs_type = "fat16"
549                if not part.system_id:
550                    part.system_id = '0x6' # FAT16
551            else:
552                # Type for ext2/ext3/ext4/btrfs
553                parted_fs_type = "ext2"
554
555            # Boot ROM of OMAP boards require vfat boot partition to have an
556            # even number of sectors.
557            if part.mountpoint == "/boot" and part.fstype in ["vfat", "msdos"] \
558               and part.size_sec % 2:
559                logger.debug("Subtracting one sector from '%s' partition to "
560                             "get even number of sectors for the partition",
561                             part.mountpoint)
562                part.size_sec -= 1
563
564            self._create_partition(self.path, part.type,
565                                   parted_fs_type, part.start, part.size_sec)
566
567            if part.part_name:
568                logger.debug("partition %d: set name to %s",
569                             part.num, part.part_name)
570                exec_native_cmd("sgdisk --change-name=%d:%s %s" % \
571                                         (part.num, part.part_name,
572                                          self.path), self.native_sysroot)
573
574            if part.part_type:
575                logger.debug("partition %d: set type UID to %s",
576                             part.num, part.part_type)
577                exec_native_cmd("sgdisk --typecode=%d:%s %s" % \
578                                         (part.num, part.part_type,
579                                          self.path), self.native_sysroot)
580
581            if part.uuid and self.ptable_format == "gpt":
582                logger.debug("partition %d: set UUID to %s",
583                             part.num, part.uuid)
584                exec_native_cmd("sgdisk --partition-guid=%d:%s %s" % \
585                                (part.num, part.uuid, self.path),
586                                self.native_sysroot)
587
588            if part.label and self.ptable_format == "gpt":
589                logger.debug("partition %d: set name to %s",
590                             part.num, part.label)
591                exec_native_cmd("parted -s %s name %d %s" % \
592                                (self.path, part.num, part.label),
593                                self.native_sysroot)
594
595            if part.active:
596                flag_name = "legacy_boot" if self.ptable_format == 'gpt' else "boot"
597                logger.debug("Set '%s' flag for partition '%s' on disk '%s'",
598                             flag_name, part.num, self.path)
599                exec_native_cmd("parted -s %s set %d %s on" % \
600                                (self.path, part.num, flag_name),
601                                self.native_sysroot)
602            if part.system_id:
603                exec_native_cmd("sfdisk --part-type %s %s %s" % \
604                                (self.path, part.num, part.system_id),
605                                self.native_sysroot)
606
607    def cleanup(self):
608        pass
609
610    def assemble(self):
611        logger.debug("Installing partitions")
612
613        for part in self.partitions:
614            source = part.source_file
615            if source:
616                # install source_file contents into a partition
617                sparse_copy(source, self.path, seek=part.start * self.sector_size)
618
619                logger.debug("Installed %s in partition %d, sectors %d-%d, "
620                             "size %d sectors", source, part.num, part.start,
621                             part.start + part.size_sec - 1, part.size_sec)
622
623                partimage = self.path + '.p%d' % part.num
624                os.rename(source, partimage)
625                self.partimages.append(partimage)
626