xref: /openbmc/openbmc/poky/scripts/lib/wic/engine.py (revision eb8dc403)
1# ex:ts=4:sw=4:sts=4:et
2# -*- tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*-
3#
4# Copyright (c) 2013, Intel Corporation.
5# All rights reserved.
6#
7# This program is free software; you can redistribute it and/or modify
8# it under the terms of the GNU General Public License version 2 as
9# published by the Free Software Foundation.
10#
11# This program is distributed in the hope that it will be useful,
12# but WITHOUT ANY WARRANTY; without even the implied warranty of
13# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
14# GNU General Public License for more details.
15#
16# You should have received a copy of the GNU General Public License along
17# with this program; if not, write to the Free Software Foundation, Inc.,
18# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
19#
20# DESCRIPTION
21
22# This module implements the image creation engine used by 'wic' to
23# create images.  The engine parses through the OpenEmbedded kickstart
24# (wks) file specified and generates images that can then be directly
25# written onto media.
26#
27# AUTHORS
28# Tom Zanussi <tom.zanussi (at] linux.intel.com>
29#
30
31import logging
32import os
33import tempfile
34import json
35import subprocess
36
37from collections import namedtuple, OrderedDict
38from distutils.spawn import find_executable
39
40from wic import WicError
41from wic.filemap import sparse_copy
42from wic.pluginbase import PluginMgr
43from wic.misc import get_bitbake_var, exec_cmd
44
45logger = logging.getLogger('wic')
46
47def verify_build_env():
48    """
49    Verify that the build environment is sane.
50
51    Returns True if it is, false otherwise
52    """
53    if not os.environ.get("BUILDDIR"):
54        raise WicError("BUILDDIR not found, exiting. (Did you forget to source oe-init-build-env?)")
55
56    return True
57
58
59CANNED_IMAGE_DIR = "lib/wic/canned-wks" # relative to scripts
60SCRIPTS_CANNED_IMAGE_DIR = "scripts/" + CANNED_IMAGE_DIR
61WIC_DIR = "wic"
62
63def build_canned_image_list(path):
64    layers_path = get_bitbake_var("BBLAYERS")
65    canned_wks_layer_dirs = []
66
67    if layers_path is not None:
68        for layer_path in layers_path.split():
69            for wks_path in (WIC_DIR, SCRIPTS_CANNED_IMAGE_DIR):
70                cpath = os.path.join(layer_path, wks_path)
71                if os.path.isdir(cpath):
72                    canned_wks_layer_dirs.append(cpath)
73
74    cpath = os.path.join(path, CANNED_IMAGE_DIR)
75    canned_wks_layer_dirs.append(cpath)
76
77    return canned_wks_layer_dirs
78
79def find_canned_image(scripts_path, wks_file):
80    """
81    Find a .wks file with the given name in the canned files dir.
82
83    Return False if not found
84    """
85    layers_canned_wks_dir = build_canned_image_list(scripts_path)
86
87    for canned_wks_dir in layers_canned_wks_dir:
88        for root, dirs, files in os.walk(canned_wks_dir):
89            for fname in files:
90                if fname.endswith("~") or fname.endswith("#"):
91                    continue
92                if fname.endswith(".wks") and wks_file + ".wks" == fname:
93                    fullpath = os.path.join(canned_wks_dir, fname)
94                    return fullpath
95    return None
96
97
98def list_canned_images(scripts_path):
99    """
100    List the .wks files in the canned image dir, minus the extension.
101    """
102    layers_canned_wks_dir = build_canned_image_list(scripts_path)
103
104    for canned_wks_dir in layers_canned_wks_dir:
105        for root, dirs, files in os.walk(canned_wks_dir):
106            for fname in files:
107                if fname.endswith("~") or fname.endswith("#"):
108                    continue
109                if fname.endswith(".wks"):
110                    fullpath = os.path.join(canned_wks_dir, fname)
111                    with open(fullpath) as wks:
112                        for line in wks:
113                            desc = ""
114                            idx = line.find("short-description:")
115                            if idx != -1:
116                                desc = line[idx + len("short-description:"):].strip()
117                                break
118                    basename = os.path.splitext(fname)[0]
119                    print("  %s\t\t%s" % (basename.ljust(30), desc))
120
121
122def list_canned_image_help(scripts_path, fullpath):
123    """
124    List the help and params in the specified canned image.
125    """
126    found = False
127    with open(fullpath) as wks:
128        for line in wks:
129            if not found:
130                idx = line.find("long-description:")
131                if idx != -1:
132                    print()
133                    print(line[idx + len("long-description:"):].strip())
134                    found = True
135                continue
136            if not line.strip():
137                break
138            idx = line.find("#")
139            if idx != -1:
140                print(line[idx + len("#:"):].rstrip())
141            else:
142                break
143
144
145def list_source_plugins():
146    """
147    List the available source plugins i.e. plugins available for --source.
148    """
149    plugins = PluginMgr.get_plugins('source')
150
151    for plugin in plugins:
152        print("  %s" % plugin)
153
154
155def wic_create(wks_file, rootfs_dir, bootimg_dir, kernel_dir,
156               native_sysroot, options):
157    """
158    Create image
159
160    wks_file - user-defined OE kickstart file
161    rootfs_dir - absolute path to the build's /rootfs dir
162    bootimg_dir - absolute path to the build's boot artifacts directory
163    kernel_dir - absolute path to the build's kernel directory
164    native_sysroot - absolute path to the build's native sysroots dir
165    image_output_dir - dirname to create for image
166    options - wic command line options (debug, bmap, etc)
167
168    Normally, the values for the build artifacts values are determined
169    by 'wic -e' from the output of the 'bitbake -e' command given an
170    image name e.g. 'core-image-minimal' and a given machine set in
171    local.conf.  If that's the case, the variables get the following
172    values from the output of 'bitbake -e':
173
174    rootfs_dir:        IMAGE_ROOTFS
175    kernel_dir:        DEPLOY_DIR_IMAGE
176    native_sysroot:    STAGING_DIR_NATIVE
177
178    In the above case, bootimg_dir remains unset and the
179    plugin-specific image creation code is responsible for finding the
180    bootimg artifacts.
181
182    In the case where the values are passed in explicitly i.e 'wic -e'
183    is not used but rather the individual 'wic' options are used to
184    explicitly specify these values.
185    """
186    try:
187        oe_builddir = os.environ["BUILDDIR"]
188    except KeyError:
189        raise WicError("BUILDDIR not found, exiting. (Did you forget to source oe-init-build-env?)")
190
191    if not os.path.exists(options.outdir):
192        os.makedirs(options.outdir)
193
194    pname = 'direct'
195    plugin_class = PluginMgr.get_plugins('imager').get(pname)
196    if not plugin_class:
197        raise WicError('Unknown plugin: %s' % pname)
198
199    plugin = plugin_class(wks_file, rootfs_dir, bootimg_dir, kernel_dir,
200                          native_sysroot, oe_builddir, options)
201
202    plugin.do_create()
203
204    logger.info("The image(s) were created using OE kickstart file:\n  %s", wks_file)
205
206
207def wic_list(args, scripts_path):
208    """
209    Print the list of images or source plugins.
210    """
211    if args.list_type is None:
212        return False
213
214    if args.list_type == "images":
215
216        list_canned_images(scripts_path)
217        return True
218    elif args.list_type == "source-plugins":
219        list_source_plugins()
220        return True
221    elif len(args.help_for) == 1 and args.help_for[0] == 'help':
222        wks_file = args.list_type
223        fullpath = find_canned_image(scripts_path, wks_file)
224        if not fullpath:
225            raise WicError("No image named %s found, exiting. "
226                           "(Use 'wic list images' to list available images, "
227                           "or specify a fully-qualified OE kickstart (.wks) "
228                           "filename)" % wks_file)
229
230        list_canned_image_help(scripts_path, fullpath)
231        return True
232
233    return False
234
235
236class Disk:
237    def __init__(self, imagepath, native_sysroot, fstypes=('fat', 'ext')):
238        self.imagepath = imagepath
239        self.native_sysroot = native_sysroot
240        self.fstypes = fstypes
241        self._partitions = None
242        self._partimages = {}
243        self._lsector_size = None
244        self._psector_size = None
245        self._ptable_format = None
246
247        # find parted
248        self.paths = "/bin:/usr/bin:/usr/sbin:/sbin/"
249        if native_sysroot:
250            for path in self.paths.split(':'):
251                self.paths = "%s%s:%s" % (native_sysroot, path, self.paths)
252
253        self.parted = find_executable("parted", self.paths)
254        if not self.parted:
255            raise WicError("Can't find executable parted")
256
257        self.partitions = self.get_partitions()
258
259    def __del__(self):
260        for path in self._partimages.values():
261            os.unlink(path)
262
263    def get_partitions(self):
264        if self._partitions is None:
265            self._partitions = OrderedDict()
266            out = exec_cmd("%s -sm %s unit B print" % (self.parted, self.imagepath))
267            parttype = namedtuple("Part", "pnum start end size fstype")
268            splitted = out.splitlines()
269            lsector_size, psector_size, self._ptable_format = splitted[1].split(":")[3:6]
270            self._lsector_size = int(lsector_size)
271            self._psector_size = int(psector_size)
272            for line in splitted[2:]:
273                pnum, start, end, size, fstype = line.split(':')[:5]
274                partition = parttype(int(pnum), int(start[:-1]), int(end[:-1]),
275                                     int(size[:-1]), fstype)
276                self._partitions[pnum] = partition
277
278        return self._partitions
279
280    def __getattr__(self, name):
281        """Get path to the executable in a lazy way."""
282        if name in ("mdir", "mcopy", "mdel", "mdeltree", "sfdisk", "e2fsck",
283                    "resize2fs", "mkswap", "mkdosfs", "debugfs"):
284            aname = "_%s" % name
285            if aname not in self.__dict__:
286                setattr(self, aname, find_executable(name, self.paths))
287                if aname not in self.__dict__:
288                    raise WicError("Can't find executable {}".format(name))
289            return self.__dict__[aname]
290        return self.__dict__[name]
291
292    def _get_part_image(self, pnum):
293        if pnum not in self.partitions:
294            raise WicError("Partition %s is not in the image")
295        part = self.partitions[pnum]
296        # check if fstype is supported
297        for fstype in self.fstypes:
298            if part.fstype.startswith(fstype):
299                break
300        else:
301            raise WicError("Not supported fstype: {}".format(part.fstype))
302        if pnum not in self._partimages:
303            tmpf = tempfile.NamedTemporaryFile(prefix="wic-part")
304            dst_fname = tmpf.name
305            tmpf.close()
306            sparse_copy(self.imagepath, dst_fname, skip=part.start, length=part.size)
307            self._partimages[pnum] = dst_fname
308
309        return self._partimages[pnum]
310
311    def _put_part_image(self, pnum):
312        """Put partition image into partitioned image."""
313        sparse_copy(self._partimages[pnum], self.imagepath,
314                    seek=self.partitions[pnum].start)
315
316    def dir(self, pnum, path):
317        if self.partitions[pnum].fstype.startswith('ext'):
318            return exec_cmd("{} {} -R 'ls -l {}'".format(self.debugfs,
319                                                         self._get_part_image(pnum),
320                                                         path), as_shell=True)
321        else: # fat
322            return exec_cmd("{} -i {} ::{}".format(self.mdir,
323                                                   self._get_part_image(pnum),
324                                                   path))
325
326    def copy(self, src, pnum, path):
327        """Copy partition image into wic image."""
328        if self.partitions[pnum].fstype.startswith('ext'):
329            cmd = "echo -e 'cd {}\nwrite {} {}' | {} -w {}".\
330                      format(path, src, os.path.basename(src),
331                             self.debugfs, self._get_part_image(pnum))
332        else: # fat
333            cmd = "{} -i {} -snop {} ::{}".format(self.mcopy,
334                                                  self._get_part_image(pnum),
335                                                  src, path)
336        exec_cmd(cmd, as_shell=True)
337        self._put_part_image(pnum)
338
339    def remove(self, pnum, path):
340        """Remove files/dirs from the partition."""
341        partimg = self._get_part_image(pnum)
342        if self.partitions[pnum].fstype.startswith('ext'):
343            exec_cmd("{} {} -wR 'rm {}'".format(self.debugfs,
344                                                self._get_part_image(pnum),
345                                                path), as_shell=True)
346        else: # fat
347            cmd = "{} -i {} ::{}".format(self.mdel, partimg, path)
348            try:
349                exec_cmd(cmd)
350            except WicError as err:
351                if "not found" in str(err) or "non empty" in str(err):
352                    # mdel outputs 'File ... not found' or 'directory .. non empty"
353                    # try to use mdeltree as path could be a directory
354                    cmd = "{} -i {} ::{}".format(self.mdeltree,
355                                                 partimg, path)
356                    exec_cmd(cmd)
357                else:
358                    raise err
359        self._put_part_image(pnum)
360
361    def write(self, target, expand):
362        """Write disk image to the media or file."""
363        def write_sfdisk_script(outf, parts):
364            for key, val in parts['partitiontable'].items():
365                if key in ("partitions", "device", "firstlba", "lastlba"):
366                    continue
367                if key == "id":
368                    key = "label-id"
369                outf.write("{}: {}\n".format(key, val))
370            outf.write("\n")
371            for part in parts['partitiontable']['partitions']:
372                line = ''
373                for name in ('attrs', 'name', 'size', 'type', 'uuid'):
374                    if name == 'size' and part['type'] == 'f':
375                        # don't write size for extended partition
376                        continue
377                    val = part.get(name)
378                    if val:
379                        line += '{}={}, '.format(name, val)
380                if line:
381                    line = line[:-2] # strip ', '
382                if part.get('bootable'):
383                    line += ' ,bootable'
384                outf.write("{}\n".format(line))
385            outf.flush()
386
387        def read_ptable(path):
388            out = exec_cmd("{} -dJ {}".format(self.sfdisk, path))
389            return json.loads(out)
390
391        def write_ptable(parts, target):
392            with tempfile.NamedTemporaryFile(prefix="wic-sfdisk-", mode='w') as outf:
393                write_sfdisk_script(outf, parts)
394                cmd = "{} --no-reread {} < {} 2>/dev/null".format(self.sfdisk, target, outf.name)
395                try:
396                    subprocess.check_output(cmd, shell=True)
397                except subprocess.CalledProcessError as err:
398                    raise WicError("Can't run '{}' command: {}".format(cmd, err))
399
400        if expand is None:
401            sparse_copy(self.imagepath, target)
402        else:
403            # copy first sectors that may contain bootloader
404            sparse_copy(self.imagepath, target, length=2048 * self._lsector_size)
405
406            # copy source partition table to the target
407            parts = read_ptable(self.imagepath)
408            write_ptable(parts, target)
409
410            # get size of unpartitioned space
411            free = None
412            for line in exec_cmd("{} -F {}".format(self.sfdisk, target)).splitlines():
413                if line.startswith("Unpartitioned space ") and line.endswith("sectors"):
414                    free = int(line.split()[-2])
415            if free is None:
416                raise WicError("Can't get size of unpartitioned space")
417
418            # calculate expanded partitions sizes
419            sizes = {}
420            for num, part in enumerate(parts['partitiontable']['partitions'], 1):
421                if num in expand:
422                    if expand[num] != 0: # don't resize partition if size is set to 0
423                        sectors = expand[num] // self._lsector_size
424                        free -= sectors - part['size']
425                        part['size'] = sectors
426                        sizes[num] = sectors
427                elif part['type'] != 'f':
428                    sizes[num] = -1
429
430            for num, part in enumerate(parts['partitiontable']['partitions'], 1):
431                if sizes.get(num) == -1:
432                    part['size'] += free // len(sizes)
433
434            # write resized partition table to the target
435            write_ptable(parts, target)
436
437            # read resized partition table
438            parts = read_ptable(target)
439
440            # copy partitions content
441            for num, part in enumerate(parts['partitiontable']['partitions'], 1):
442                pnum = str(num)
443                fstype = self.partitions[pnum].fstype
444
445                # copy unchanged partition
446                if part['size'] == self.partitions[pnum].size // self._lsector_size:
447                    logger.info("copying unchanged partition {}".format(pnum))
448                    sparse_copy(self._get_part_image(pnum), target, seek=part['start'] * self._lsector_size)
449                    continue
450
451                # resize or re-create partitions
452                if fstype.startswith('ext') or fstype.startswith('fat') or \
453                   fstype.startswith('linux-swap'):
454
455                    partfname = None
456                    with tempfile.NamedTemporaryFile(prefix="wic-part{}-".format(pnum)) as partf:
457                        partfname = partf.name
458
459                    if fstype.startswith('ext'):
460                        logger.info("resizing ext partition {}".format(pnum))
461                        partimg = self._get_part_image(pnum)
462                        sparse_copy(partimg, partfname)
463                        exec_cmd("{} -pf {}".format(self.e2fsck, partfname))
464                        exec_cmd("{} {} {}s".format(\
465                                 self.resize2fs, partfname, part['size']))
466                    elif fstype.startswith('fat'):
467                        logger.info("copying content of the fat partition {}".format(pnum))
468                        with tempfile.TemporaryDirectory(prefix='wic-fatdir-') as tmpdir:
469                            # copy content to the temporary directory
470                            cmd = "{} -snompi {} :: {}".format(self.mcopy,
471                                                               self._get_part_image(pnum),
472                                                               tmpdir)
473                            exec_cmd(cmd)
474                            # create new msdos partition
475                            label = part.get("name")
476                            label_str = "-n {}".format(label) if label else ''
477
478                            cmd = "{} {} -C {} {}".format(self.mkdosfs, label_str, partfname,
479                                                          part['size'])
480                            exec_cmd(cmd)
481                            # copy content from the temporary directory to the new partition
482                            cmd = "{} -snompi {} {}/* ::".format(self.mcopy, partfname, tmpdir)
483                            exec_cmd(cmd, as_shell=True)
484                    elif fstype.startswith('linux-swap'):
485                        logger.info("creating swap partition {}".format(pnum))
486                        label = part.get("name")
487                        label_str = "-L {}".format(label) if label else ''
488                        uuid = part.get("uuid")
489                        uuid_str = "-U {}".format(uuid) if uuid else ''
490                        with open(partfname, 'w') as sparse:
491                            os.ftruncate(sparse.fileno(), part['size'] * self._lsector_size)
492                        exec_cmd("{} {} {} {}".format(self.mkswap, label_str, uuid_str, partfname))
493                    sparse_copy(partfname, target, seek=part['start'] * self._lsector_size)
494                    os.unlink(partfname)
495                elif part['type'] != 'f':
496                    logger.warn("skipping partition {}: unsupported fstype {}".format(pnum, fstype))
497
498def wic_ls(args, native_sysroot):
499    """List contents of partitioned image or vfat partition."""
500    disk = Disk(args.path.image, native_sysroot)
501    if not args.path.part:
502        if disk.partitions:
503            print('Num     Start        End          Size      Fstype')
504            for part in disk.partitions.values():
505                print("{:2d}  {:12d} {:12d} {:12d}  {}".format(\
506                          part.pnum, part.start, part.end,
507                          part.size, part.fstype))
508    else:
509        path = args.path.path or '/'
510        print(disk.dir(args.path.part, path))
511
512def wic_cp(args, native_sysroot):
513    """
514    Copy local file or directory to the vfat partition of
515    partitioned image.
516    """
517    disk = Disk(args.dest.image, native_sysroot)
518    disk.copy(args.src, args.dest.part, args.dest.path)
519
520def wic_rm(args, native_sysroot):
521    """
522    Remove files or directories from the vfat partition of
523    partitioned image.
524    """
525    disk = Disk(args.path.image, native_sysroot)
526    disk.remove(args.path.part, args.path.path)
527
528def wic_write(args, native_sysroot):
529    """
530    Write image to a target device.
531    """
532    disk = Disk(args.image, native_sysroot, ('fat', 'ext', 'swap'))
533    disk.write(args.target, args.expand)
534
535def find_canned(scripts_path, file_name):
536    """
537    Find a file either by its path or by name in the canned files dir.
538
539    Return None if not found
540    """
541    if os.path.exists(file_name):
542        return file_name
543
544    layers_canned_wks_dir = build_canned_image_list(scripts_path)
545    for canned_wks_dir in layers_canned_wks_dir:
546        for root, dirs, files in os.walk(canned_wks_dir):
547            for fname in files:
548                if fname == file_name:
549                    fullpath = os.path.join(canned_wks_dir, fname)
550                    return fullpath
551
552def get_custom_config(boot_file):
553    """
554    Get the custom configuration to be used for the bootloader.
555
556    Return None if the file can't be found.
557    """
558    # Get the scripts path of poky
559    scripts_path = os.path.abspath("%s/../.." % os.path.dirname(__file__))
560
561    cfg_file = find_canned(scripts_path, boot_file)
562    if cfg_file:
563        with open(cfg_file, "r") as f:
564            config = f.read()
565        return config
566