summaryrefslogtreecommitdiffstats
path: root/scripts/lib/wic/plugins/imager/direct.py
blob: 2441cc33adb25ad72817c4b60142838fe4bf6e04 (plain)
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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
#
# Copyright (c) 2013, Intel Corporation.
#
# SPDX-License-Identifier: GPL-2.0-only
#
# DESCRIPTION
# This implements the 'direct' imager plugin class for 'wic'
#
# AUTHORS
# Tom Zanussi <tom.zanussi (at] linux.intel.com>
#

import logging
import os
import random
import shutil
import tempfile
import uuid

from time import strftime

from oe.path import copyhardlinktree

from wic import WicError
from wic.filemap import sparse_copy
from wic.ksparser import KickStart, KickStartError
from wic.pluginbase import PluginMgr, ImagerPlugin
from wic.misc import get_bitbake_var, exec_cmd, exec_native_cmd

logger = logging.getLogger('wic')

class DirectPlugin(ImagerPlugin):
    """
    Install a system into a file containing a partitioned disk image.

    An image file is formatted with a partition table, each partition
    created from a rootfs or other OpenEmbedded build artifact and dd'ed
    into the virtual disk. The disk image can subsequently be dd'ed onto
    media and used on actual hardware.
    """
    name = 'direct'

    def __init__(self, wks_file, rootfs_dir, bootimg_dir, kernel_dir,
                 native_sysroot, oe_builddir, options):
        try:
            self.ks = KickStart(wks_file)
        except KickStartError as err:
            raise WicError(str(err))

        # parse possible 'rootfs=name' items
        self.rootfs_dir = dict(rdir.split('=') for rdir in rootfs_dir.split(' '))
        self.bootimg_dir = bootimg_dir
        self.kernel_dir = kernel_dir
        self.native_sysroot = native_sysroot
        self.oe_builddir = oe_builddir

        self.outdir = options.outdir
        self.compressor = options.compressor
        self.bmap = options.bmap
        self.no_fstab_update = options.no_fstab_update
        self.original_fstab = None

        self.name = "%s-%s" % (os.path.splitext(os.path.basename(wks_file))[0],
                               strftime("%Y%m%d%H%M"))
        self.workdir = tempfile.mkdtemp(dir=self.outdir, prefix='tmp.wic.')
        self._image = None
        self.ptable_format = self.ks.bootloader.ptable
        self.parts = self.ks.partitions

        # as a convenience, set source to the boot partition source
        # instead of forcing it to be set via bootloader --source
        for part in self.parts:
            if not self.ks.bootloader.source and part.mountpoint == "/boot":
                self.ks.bootloader.source = part.source
                break

        image_path = self._full_path(self.workdir, self.parts[0].disk, "direct")
        self._image = PartitionedImage(image_path, self.ptable_format,
                                       self.parts, self.native_sysroot)

    def do_create(self):
        """
        Plugin entry point.
        """
        try:
            self.create()
            self.assemble()
            self.finalize()
            self.print_info()
        finally:
            self.cleanup()

    def _write_fstab(self, image_rootfs):
        """overriden to generate fstab (temporarily) in rootfs. This is called
        from _create, make sure it doesn't get called from
        BaseImage.create()
        """
        if not image_rootfs:
            return

        fstab_path = image_rootfs + "/etc/fstab"
        if not os.path.isfile(fstab_path):
            return

        with open(fstab_path) as fstab:
            fstab_lines = fstab.readlines()
            self.original_fstab = fstab_lines.copy()

        if self._update_fstab(fstab_lines, self.parts):
            with open(fstab_path, "w") as fstab:
                fstab.writelines(fstab_lines)
        else:
            self.original_fstab = None

    def _update_fstab(self, fstab_lines, parts):
        """Assume partition order same as in wks"""
        updated = False
        for part in parts:
            if not part.realnum or not part.mountpoint \
               or part.mountpoint == "/":
                continue

            if part.use_uuid:
                if part.fsuuid:
                    # FAT UUID is different from others
                    if len(part.fsuuid) == 10:
                        device_name = "UUID=%s-%s" % \
                                       (part.fsuuid[2:6], part.fsuuid[6:])
                    else:
                        device_name = "UUID=%s" % part.fsuuid
                else:
                    device_name = "PARTUUID=%s" % part.uuid
            elif part.use_label:
                device_name = "LABEL=%s" % part.label
            else:
                # mmc device partitions are named mmcblk0p1, mmcblk0p2..
                prefix = 'p' if  part.disk.startswith('mmcblk') else ''
                device_name = "/dev/%s%s%d" % (part.disk, prefix, part.realnum)

            opts = part.fsopts if part.fsopts else "defaults"
            line = "\t".join([device_name, part.mountpoint, part.fstype,
                              opts, "0", "0"]) + "\n"

            fstab_lines.append(line)
            updated = True

        return updated

    def _full_path(self, path, name, extention):
        """ Construct full file path to a file we generate. """
        return os.path.join(path, "%s-%s.%s" % (self.name, name, extention))

    #
    # Actual implemention
    #
    def create(self):
        """
        For 'wic', we already have our build artifacts - we just create
        filesystems from the artifacts directly and combine them into
        a partitioned image.
        """
        if not self.no_fstab_update:
            self._write_fstab(self.rootfs_dir.get("ROOTFS_DIR"))

        for part in self.parts:
            # get rootfs size from bitbake variable if it's not set in .ks file
            if not part.size:
                # and if rootfs name is specified for the partition
                image_name = self.rootfs_dir.get(part.rootfs_dir)
                if image_name and os.path.sep not in image_name:
                    # Bitbake variable ROOTFS_SIZE is calculated in
                    # Image._get_rootfs_size method from meta/lib/oe/image.py
                    # using IMAGE_ROOTFS_SIZE, IMAGE_ROOTFS_ALIGNMENT,
                    # IMAGE_OVERHEAD_FACTOR and IMAGE_ROOTFS_EXTRA_SPACE
                    rsize_bb = get_bitbake_var('ROOTFS_SIZE', image_name)
                    if rsize_bb:
                        part.size = int(round(float(rsize_bb)))

        self._image.prepare(self)
        self._image.layout_partitions()
        self._image.create()

    def assemble(self):
        """
        Assemble partitions into disk image
        """
        self._image.assemble()

    def finalize(self):
        """
        Finalize the disk image.

        For example, prepare the image to be bootable by e.g.
        creating and installing a bootloader configuration.
        """
        source_plugin = self.ks.bootloader.source
        disk_name = self.parts[0].disk
        if source_plugin:
            plugin = PluginMgr.get_plugins('source')[source_plugin]
            plugin.do_install_disk(self._image, disk_name, self, self.workdir,
                                   self.oe_builddir, self.bootimg_dir,
                                   self.kernel_dir, self.native_sysroot)

        full_path = self._image.path
        # Generate .bmap
        if self.bmap:
            logger.debug("Generating bmap file for %s", disk_name)
            python = os.path.join(self.native_sysroot, 'usr/bin/python3-native/python3')
            bmaptool = os.path.join(self.native_sysroot, 'usr/bin/bmaptool')
            exec_native_cmd("%s %s create %s -o %s.bmap" % \
                            (python, bmaptool, full_path, full_path), self.native_sysroot)
        # Compress the image
        if self.compressor:
            logger.debug("Compressing disk %s with %s", disk_name, self.compressor)
            exec_cmd("%s %s" % (self.compressor, full_path))

    def print_info(self):
        """
        Print the image(s) and artifacts used, for the user.
        """
        msg = "The new image(s) can be found here:\n"

        extension = "direct" + {"gzip": ".gz",
                                "bzip2": ".bz2",
                                "xz": ".xz",
                                None: ""}.get(self.compressor)
        full_path = self._full_path(self.outdir, self.parts[0].disk, extension)
        msg += '  %s\n\n' % full_path

        msg += 'The following build artifacts were used to create the image(s):\n'
        for part in self.parts:
            if part.rootfs_dir is None:
                continue
            if part.mountpoint == '/':
                suffix = ':'
            else:
                suffix = '["%s"]:' % (part.mountpoint or part.label)
            rootdir = part.rootfs_dir
            msg += '  ROOTFS_DIR%s%s\n' % (suffix.ljust(20), rootdir)

        msg += '  BOOTIMG_DIR:                  %s\n' % self.bootimg_dir
        msg += '  KERNEL_DIR:                   %s\n' % self.kernel_dir
        msg += '  NATIVE_SYSROOT:               %s\n' % self.native_sysroot

        logger.info(msg)

    @property
    def rootdev(self):
        """
        Get root device name to use as a 'root' parameter
        in kernel command line.

        Assume partition order same as in wks
        """
        for part in self.parts:
            if part.mountpoint == "/":
                if part.uuid:
                    return "PARTUUID=%s" % part.uuid
                else:
                    suffix = 'p' if part.disk.startswith('mmcblk') else ''
                    return "/dev/%s%s%-d" % (part.disk, suffix, part.realnum)

    def cleanup(self):
        if self._image:
            self._image.cleanup()

        # Move results to the output dir
        if not os.path.exists(self.outdir):
            os.makedirs(self.outdir)

        for fname in os.listdir(self.workdir):
            path = os.path.join(self.workdir, fname)
            if os.path.isfile(path):
                shutil.move(path, os.path.join(self.outdir, fname))

        #Restore original fstab
        if self.original_fstab:
            fstab_path = self.rootfs_dir.get("ROOTFS_DIR") + "/etc/fstab"
            with open(fstab_path, "w") as fstab:
                fstab.writelines(self.original_fstab)

        # remove work directory
        shutil.rmtree(self.workdir, ignore_errors=True)

# Overhead of the MBR partitioning scheme (just one sector)
MBR_OVERHEAD = 1

# Overhead of the GPT partitioning scheme
GPT_OVERHEAD = 34

# Size of a sector in bytes
SECTOR_SIZE = 512

class PartitionedImage():
    """
    Partitioned image in a file.
    """

    def __init__(self, path, ptable_format, partitions, native_sysroot=None):
        self.path = path  # Path to the image file
        self.numpart = 0  # Number of allocated partitions
        self.realpart = 0 # Number of partitions in the partition table
        self.primary_part_num = 0  # Number of primary partitions (msdos)
        self.extendedpart = 0      # Create extended partition before this logical partition (msdos)
        self.extended_size_sec = 0 # Size of exteded partition (msdos)
        self.logical_part_cnt = 0  # Number of total logical paritions (msdos)
        self.offset = 0   # Offset of next partition (in sectors)
        self.min_size = 0 # Minimum required disk size to fit
                          # all partitions (in bytes)
        self.ptable_format = ptable_format  # Partition table format
        # Disk system identifier
        self.identifier = random.SystemRandom().randint(1, 0xffffffff)

        self.partitions = partitions
        self.partimages = []
        # Size of a sector used in calculations
        self.sector_size = SECTOR_SIZE
        self.native_sysroot = native_sysroot
        num_real_partitions = len([p for p in self.partitions if not p.no_table])

        # calculate the real partition number, accounting for partitions not
        # in the partition table and logical partitions
        realnum = 0
        for part in self.partitions:
            if part.no_table:
                part.realnum = 0
            else:
                realnum += 1
                if self.ptable_format == 'msdos' and realnum > 3 and num_real_partitions > 4:
                    part.realnum = realnum + 1
                    continue
                part.realnum = realnum

        # generate parition and filesystem UUIDs
        for part in self.partitions:
            if not part.uuid and part.use_uuid:
                if self.ptable_format == 'gpt':
                    part.uuid = str(uuid.uuid4())
                else: # msdos partition table
                    part.uuid = '%08x-%02d' % (self.identifier, part.realnum)
            if not part.fsuuid:
                if part.fstype == 'vfat' or part.fstype == 'msdos':
                    part.fsuuid = '0x' + str(uuid.uuid4())[:8].upper()
                else:
                    part.fsuuid = str(uuid.uuid4())

    def prepare(self, imager):
        """Prepare an image. Call prepare method of all image partitions."""
        for part in self.partitions:
            # need to create the filesystems in order to get their
            # sizes before we can add them and do the layout.
            part.prepare(imager, imager.workdir, imager.oe_builddir,
                         imager.rootfs_dir, imager.bootimg_dir,
                         imager.kernel_dir, imager.native_sysroot)

            # Converting kB to sectors for parted
            part.size_sec = part.disk_size * 1024 // self.sector_size

    def layout_partitions(self):
        """ Layout the partitions, meaning calculate the position of every
        partition on the disk. The 'ptable_format' parameter defines the
        partition table format and may be "msdos". """

        logger.debug("Assigning %s partitions to disks", self.ptable_format)

        # The number of primary and logical partitions. Extended partition and
        # partitions not listed in the table are not included.
        num_real_partitions = len([p for p in self.partitions if not p.no_table])

        # Go through partitions in the order they are added in .ks file
        for num in range(len(self.partitions)):
            part = self.partitions[num]

            if self.ptable_format == 'msdos' and part.part_name:
                raise WicError("setting custom partition name is not " \
                               "implemented for msdos partitions")

            if self.ptable_format == 'msdos' and part.part_type:
                # The --part-type can also be implemented for MBR partitions,
                # in which case it would map to the 1-byte "partition type"
                # filed at offset 3 of the partition entry.
                raise WicError("setting custom partition type is not " \
                               "implemented for msdos partitions")

            # Get the disk where the partition is located
            self.numpart += 1
            if not part.no_table:
                self.realpart += 1

            if self.numpart == 1:
                if self.ptable_format == "msdos":
                    overhead = MBR_OVERHEAD
                elif self.ptable_format == "gpt":
                    overhead = GPT_OVERHEAD

                # Skip one sector required for the partitioning scheme overhead
                self.offset += overhead

            if self.ptable_format == "msdos":
                if self.primary_part_num > 3 or \
                   (self.extendedpart == 0 and self.primary_part_num >= 3 and num_real_partitions > 4):
                    part.type = 'logical'
                # Reserve a sector for EBR for every logical partition
                # before alignment is performed.
                if part.type == 'logical':
                    self.offset += 1

            align_sectors = 0
            if part.align:
                # If not first partition and we do have alignment set we need
                # to align the partition.
                # FIXME: This leaves a empty spaces to the disk. To fill the
                # gaps we could enlargea the previous partition?

                # Calc how much the alignment is off.
                align_sectors = self.offset % (part.align * 1024 // self.sector_size)

                if align_sectors:
                    # If partition is not aligned as required, we need
                    # to move forward to the next alignment point
                    align_sectors = (part.align * 1024 // self.sector_size) - align_sectors

                    logger.debug("Realignment for %s%s with %s sectors, original"
                                 " offset %s, target alignment is %sK.",
                                 part.disk, self.numpart, align_sectors,
                                 self.offset, part.align)

                    # increase the offset so we actually start the partition on right alignment
                    self.offset += align_sectors

            part.start = self.offset
            self.offset += part.size_sec

            if not part.no_table:
                part.num = self.realpart
            else:
                part.num = 0

            if self.ptable_format == "msdos" and not part.no_table:
                if part.type == 'logical':
                    self.logical_part_cnt += 1
                    part.num = self.logical_part_cnt + 4
                    if self.extendedpart == 0:
                        # Create extended partition as a primary partition
                        self.primary_part_num += 1
                        self.extendedpart = part.num
                    else:
                        self.extended_size_sec += align_sectors
                    self.extended_size_sec += part.size_sec + 1
                else:
                    self.primary_part_num += 1
                    part.num = self.primary_part_num

            logger.debug("Assigned %s to %s%d, sectors range %d-%d size %d "
                         "sectors (%d bytes).", part.mountpoint, part.disk,
                         part.num, part.start, self.offset - 1, part.size_sec,
                         part.size_sec * self.sector_size)

        # Once all the partitions have been layed out, we can calculate the
        # minumim disk size
        self.min_size = self.offset
        if self.ptable_format == "gpt":
            self.min_size += GPT_OVERHEAD

        self.min_size *= self.sector_size

    def _create_partition(self, device, parttype, fstype, start, size):
        """ Create a partition on an image described by the 'device' object. """

        # Start is included to the size so we need to substract one from the end.
        end = start + size - 1
        logger.debug("Added '%s' partition, sectors %d-%d, size %d sectors",
                     parttype, start, end, size)

        cmd = "parted -s %s unit s mkpart %s" % (device, parttype)
        if fstype:
            cmd += " %s" % fstype
        cmd += " %d %d" % (start, end)

        return exec_native_cmd(cmd, self.native_sysroot)

    def create(self):
        logger.debug("Creating sparse file %s", self.path)
        with open(self.path, 'w') as sparse:
            os.ftruncate(sparse.fileno(), self.min_size)

        logger.debug("Initializing partition table for %s", self.path)
        exec_native_cmd("parted -s %s mklabel %s" %
                        (self.path, self.ptable_format), self.native_sysroot)

        logger.debug("Set disk identifier %x", self.identifier)
        with open(self.path, 'r+b') as img:
            img.seek(0x1B8)
            img.write(self.identifier.to_bytes(4, 'little'))

        logger.debug("Creating partitions")

        for part in self.partitions:
            if part.num == 0:
                continue

            if self.ptable_format == "msdos" and part.num == self.extendedpart:
                # Create an extended partition (note: extended
                # partition is described in MBR and contains all
                # logical partitions). The logical partitions save a
                # sector for an EBR just before the start of a
                # partition. The extended partition must start one
                # sector before the start of the first logical
                # partition. This way the first EBR is inside of the
                # extended partition. Since the extended partitions
                # starts a sector before the first logical partition,
                # add a sector at the back, so that there is enough
                # room for all logical partitions.
                self._create_partition(self.path, "extended",
                                       None, part.start - 1,
                                       self.extended_size_sec)

            if part.fstype == "swap":
                parted_fs_type = "linux-swap"
            elif part.fstype == "vfat":
                parted_fs_type = "fat32"
            elif part.fstype == "msdos":
                parted_fs_type = "fat16"
                if not part.system_id:
                    part.system_id = '0x6' # FAT16
            else:
                # Type for ext2/ext3/ext4/btrfs
                parted_fs_type = "ext2"

            # Boot ROM of OMAP boards require vfat boot partition to have an
            # even number of sectors.
            if part.mountpoint == "/boot" and part.fstype in ["vfat", "msdos"] \
               and part.size_sec % 2:
                logger.debug("Subtracting one sector from '%s' partition to "
                             "get even number of sectors for the partition",
                             part.mountpoint)
                part.size_sec -= 1

            self._create_partition(self.path, part.type,
                                   parted_fs_type, part.start, part.size_sec)

            if part.part_name:
                logger.debug("partition %d: set name to %s",
                             part.num, part.part_name)
                exec_native_cmd("sgdisk --change-name=%d:%s %s" % \
                                         (part.num, part.part_name,
                                          self.path), self.native_sysroot)

            if part.part_type:
                logger.debug("partition %d: set type UID to %s",
                             part.num, part.part_type)
                exec_native_cmd("sgdisk --typecode=%d:%s %s" % \
                                         (part.num, part.part_type,
                                          self.path), self.native_sysroot)

            if part.uuid and self.ptable_format == "gpt":
                logger.debug("partition %d: set UUID to %s",
                             part.num, part.uuid)
                exec_native_cmd("sgdisk --partition-guid=%d:%s %s" % \
                                (part.num, part.uuid, self.path),
                                self.native_sysroot)

            if part.label and self.ptable_format == "gpt":
                logger.debug("partition %d: set name to %s",
                             part.num, part.label)
                exec_native_cmd("parted -s %s name %d %s" % \
                                (self.path, part.num, part.label),
                                self.native_sysroot)

            if part.active:
                flag_name = "legacy_boot" if self.ptable_format == 'gpt' else "boot"
                logger.debug("Set '%s' flag for partition '%s' on disk '%s'",
                             flag_name, part.num, self.path)
                exec_native_cmd("parted -s %s set %d %s on" % \
                                (self.path, part.num, flag_name),
                                self.native_sysroot)
            if part.system_id:
                exec_native_cmd("sfdisk --part-type %s %s %s" % \
                                (self.path, part.num, part.system_id),
                                self.native_sysroot)

    def cleanup(self):
        # remove partition images
        for image in set(self.partimages):
            os.remove(image)

    def assemble(self):
        logger.debug("Installing partitions")

        for part in self.partitions:
            source = part.source_file
            if source:
                # install source_file contents into a partition
                sparse_copy(source, self.path, seek=part.start * self.sector_size)

                logger.debug("Installed %s in partition %d, sectors %d-%d, "
                             "size %d sectors", source, part.num, part.start,
                             part.start + part.size_sec - 1, part.size_sec)

                partimage = self.path + '.p%d' % part.num
                os.rename(source, partimage)
                self.partimages.append(partimage)