mic/imager/liveusb.py

   1 #!/usr/bin/python -tt
   2 #
   3 # Copyright (c) 2011 Intel, Inc.
   4 #
   5 # This program is free software; you can redistribute it and/or modify it
   6 # under the terms of the GNU General Public License as published by the Free
   7 # Software Foundation; version 2 of the License
   8 #
   9 # This program is distributed in the hope that it will be useful, but
  10 # WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
  11 # or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
  12 # for more details.
  13 #
  14 # You should have received a copy of the GNU General Public License along
  15 # with this program; if not, write to the Free Software Foundation, Inc., 59
  16 # Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  17
  18 import os
  19 import shutil
  20 import re
  21
  22 from mic import msger
  23 from mic.utils import misc, fs_related, runner
  24 from mic.utils.errors import CreatorError, MountError
  25 from mic.utils.partitionedfs import PartitionedMount
  26 from mic.imager.livecd import LiveCDImageCreator
  27 from mic.archive import packing
  28
  29 class LiveUSBImageCreator(LiveCDImageCreator):
  30     img_format = 'liveusb'
  31
  32     def __init__(self, *args):
  33         LiveCDImageCreator.__init__(self, *args)
  34
  35         self._dep_checks.extend(["kpartx", "parted"])
  36
  37         # remove dependency of genisoimage in parent class
  38         if "genisoimage" in self._dep_checks:
  39             self._dep_checks.remove("genisoimage")
  40
  41     def _create_usbimg(self, isodir):
  42         overlaysizemb = 64 #default
  43         #skipcompress = self.skip_compression?
  44         fstype = "vfat"
  45         homesizemb=0
  46         swapsizemb=0
  47         homefile="home.img"
  48         plussize=128
  49         kernelargs=None
  50
  51         if fstype == 'vfat':
  52             if overlaysizemb > 2047:
  53                 raise CreatorError("Can't have an overlay of 2048MB or "
  54                                    "greater on VFAT")
  55
  56             if homesizemb > 2047:
  57                 raise CreatorError("Can't have an home overlay of 2048MB or "
  58                                    "greater on VFAT")
  59
  60             if swapsizemb > 2047:
  61                 raise CreatorError("Can't have an swap overlay of 2048MB or "
  62                                    "greater on VFAT")
  63
  64         livesize = misc.get_file_size(isodir + "/LiveOS")
  65
  66         usbimgsize = (overlaysizemb + \
  67                       homesizemb + \
  68                       swapsizemb + \
  69                       livesize + \
  70                       plussize) * 1024L * 1024L
  71
  72         disk = fs_related.SparseLoopbackDisk("%s/%s.usbimg" \
  73                                                  % (self._outdir, self.name),
  74                                              usbimgsize)
  75         usbmnt = self._mkdtemp("usb-mnt")
  76         usbloop = PartitionedMount(usbmnt)
  77         usbloop.add_disk('/dev/sdb', disk)
  78
  79         usbloop.add_partition(usbimgsize/1024/1024,
  80                               "/dev/sdb",
  81                               "/",
  82                               fstype,
  83                               boot=True)
  84
  85         usbloop.mount()
  86
  87         try:
  88             fs_related.makedirs(usbmnt + "/LiveOS")
  89
  90             if os.path.exists(isodir + "/LiveOS/squashfs.img"):
  91                 shutil.copyfile(isodir + "/LiveOS/squashfs.img",
  92                                 usbmnt + "/LiveOS/squashfs.img")
  93             else:
  94                 fs_related.mksquashfs(os.path.dirname(self._image),
  95                                       usbmnt + "/LiveOS/squashfs.img")
  96
  97             if os.path.exists(isodir + "/LiveOS/osmin.img"):
  98                 shutil.copyfile(isodir + "/LiveOS/osmin.img",
  99                                 usbmnt + "/LiveOS/osmin.img")
 100
 101             if fstype == "vfat" or fstype == "msdos":
 102                 uuid = usbloop.partitions[0]['mount'].uuid
 103                 label = usbloop.partitions[0]['mount'].fslabel
 104                 usblabel = "UUID=%s" % (uuid)
 105                 overlaysuffix = "-%s-%s" % (label, uuid)
 106             else:
 107                 diskmount = usbloop.partitions[0]['mount']
 108                 usblabel = "UUID=%s" % diskmount.uuid
 109                 overlaysuffix = "-%s-%s" % (diskmount.fslabel, diskmount.uuid)
 110
 111             args = ['cp', "-Rf", isodir + "/isolinux", usbmnt + "/syslinux"]
 112             rc = runner.show(args)
 113             if rc:
 114                 raise CreatorError("Can't copy isolinux directory %s" \
 115                                    % (isodir + "/isolinux/*"))
 116
 117             if os.path.isfile("/usr/share/syslinux/isolinux.bin"):
 118                 syslinux_path = "/usr/share/syslinux"
 119             elif  os.path.isfile("/usr/lib/syslinux/isolinux.bin"):
 120                 syslinux_path = "/usr/lib/syslinux"
 121             elif  os.path.isfile("/usr/lib/syslinux/bios/isolinux.bin"):
 122                 syslinux_path = "/usr/lib/syslinux/bios"
 123             else:
 124                 raise CreatorError("syslinux not installed : "
 125                                    "cannot find syslinux installation path")
 126
 127             for f in ("isolinux.bin", "vesamenu.c32"):
 128                 path = os.path.join(syslinux_path, f)
 129                 if os.path.isfile(path):
 130                     args = ['cp', path, usbmnt + "/syslinux/"]
 131                     rc = runner.show(args)
 132                     if rc:
 133                         raise CreatorError("Can't copy syslinux file " + path)
 134                 else:
 135                     raise CreatorError("syslinux not installed: "
 136                                        "syslinux file %s not found" % path)
 137
 138             fd = open(isodir + "/isolinux/isolinux.cfg", "r")
 139             text = fd.read()
 140             fd.close()
 141             pattern = re.compile('CDLABEL=[^ ]*')
 142             text = pattern.sub(usblabel, text)
 143             pattern = re.compile('rootfstype=[^ ]*')
 144             text = pattern.sub("rootfstype=" + fstype, text)
 145             if kernelargs:
 146                 text = text.replace("rd.live.image", "rd.live.image " + kernelargs)
 147
 148             if overlaysizemb > 0:
 149                 msger.info("Initializing persistent overlay file")
 150                 overfile = "overlay" + overlaysuffix
 151                 if fstype == "vfat":
 152                     args = ['dd',
 153                             "if=/dev/zero",
 154                             "of=" + usbmnt + "/LiveOS/" + overfile,
 155                             "count=%d" % overlaysizemb,
 156                             "bs=1M"]
 157                 else:
 158                     args = ['dd',
 159                             "if=/dev/null",
 160                             "of=" + usbmnt + "/LiveOS/" + overfile,
 161                             "count=1",
 162                             "bs=1M",
 163                             "seek=%d" % overlaysizemb]
 164                 rc = runner.show(args)
 165                 if rc:
 166                     raise CreatorError("Can't create overlay file")
 167                 text = text.replace("rd.live.image", "rd.live.image rd.live.overlay=" + usblabel)
 168                 text = text.replace(" ro ", " rw ")
 169
 170             if swapsizemb > 0:
 171                 msger.info("Initializing swap file")
 172                 swapfile = usbmnt + "/LiveOS/" + "swap.img"
 173                 args = ['dd',
 174                         "if=/dev/zero",
 175                         "of=" + swapfile,
 176                         "count=%d" % swapsizemb,
 177                         "bs=1M"]
 178                 rc = runner.show(args)
 179                 if rc:
 180                     raise CreatorError("Can't create swap file")
 181                 args = ["mkswap", "-f", swapfile]
 182                 rc = runner.show(args)
 183                 if rc:
 184                     raise CreatorError("Can't mkswap on swap file")
 185
 186             if homesizemb > 0:
 187                 msger.info("Initializing persistent /home")
 188                 homefile = usbmnt + "/LiveOS/" + homefile
 189                 if fstype == "vfat":
 190                     args = ['dd',
 191                             "if=/dev/zero",
 192                             "of=" + homefile,
 193                             "count=%d" % homesizemb,
 194                             "bs=1M"]
 195                 else:
 196                     args = ['dd',
 197                             "if=/dev/null",
 198                             "of=" + homefile,
 199                             "count=1",
 200                             "bs=1M",
 201                             "seek=%d" % homesizemb]
 202                 rc = runner.show(args)
 203                 if rc:
 204                     raise CreatorError("Can't create home file")
 205
 206                 mkfscmd = fs_related.find_binary_path("/sbin/mkfs." + fstype)
 207                 if fstype == "ext2" or fstype == "ext3":
 208                     args = [mkfscmd, "-F", "-j", homefile]
 209                 else:
 210                     args = [mkfscmd, homefile]
 211                 rc = runner.show(args)
 212                 if rc:
 213                     raise CreatorError("Can't mke2fs home file")
 214                 if fstype == "ext2" or fstype == "ext3":
 215                     tune2fs = fs_related.find_binary_path("tune2fs")
 216                     args = [tune2fs,
 217                             "-c0",
 218                             "-i0",
 219                             "-ouser_xattr,acl",
 220                             homefile]
 221                     rc = runner.show(args)
 222                     if rc:
 223                         raise CreatorError("Can't tune2fs home file")
 224
 225             if fstype == "vfat" or fstype == "msdos":
 226                 syslinuxcmd = fs_related.find_binary_path("syslinux")
 227                 syslinuxcfg = usbmnt + "/syslinux/syslinux.cfg"
 228                 args = [syslinuxcmd,
 229                         "-d",
 230                         "syslinux",
 231                         usbloop.partitions[0]["device"]]
 232
 233             elif fstype == "ext2" or fstype == "ext3":
 234                 extlinuxcmd = fs_related.find_binary_path("extlinux")
 235                 syslinuxcfg = usbmnt + "/syslinux/extlinux.conf"
 236                 args = [extlinuxcmd,
 237                         "-i",
 238                         usbmnt + "/syslinux"]
 239
 240             else:
 241                 raise CreatorError("Invalid file system type: %s" % (fstype))
 242
 243             os.unlink(usbmnt + "/syslinux/isolinux.cfg")
 244             fd = open(syslinuxcfg, "w")
 245             fd.write(text)
 246             fd.close()
 247             rc = runner.show(args)
 248             if rc:
 249                 raise CreatorError("Can't install boot loader.")
 250
 251         finally:
 252             usbloop.unmount()
 253             usbloop.cleanup()
 254
 255         # Need to do this after image is unmounted and device mapper is closed
 256         msger.info("set MBR")
 257         found = False
 258         for mbrfile in ["/usr/lib/syslinux/mbr.bin", "/usr/lib/syslinux/bios/mbr.bin", "/usr/share/syslinux/mbr.bin"]:
 259             if os.path.exists(mbrfile):
 260                 found = True
 261                 break
 262         if not found:
 263             raise CreatorError("mbr.bin file didn't exist.")
 264         mbrsize = os.path.getsize(mbrfile)
 265         outimg = "%s/%s.usbimg" % (self._outdir, self.name)
 266
 267         args = ['dd',
 268                 "if=" + mbrfile,
 269                 "of=" + outimg,
 270                 "seek=0",
 271                 "conv=notrunc",
 272                 "bs=1",
 273                 "count=%d" % (mbrsize)]
 274         rc = runner.show(args)
 275         if rc:
 276             raise CreatorError("Can't set MBR.")
 277
 278     def _stage_final_image(self):
 279         try:
 280             isodir = self._get_isodir()
 281             fs_related.makedirs(isodir + "/LiveOS")
 282
 283             minimal_size = self._resparse()
 284
 285             if not self.skip_minimize:
 286                 fs_related.create_image_minimizer(isodir + "/LiveOS/osmin.img",
 287                                                   self._image,
 288                                                   minimal_size)
 289
 290             if self.skip_compression:
 291                 shutil.move(self._image,
 292                             isodir + "/LiveOS/ext3fs.img")
 293             else:
 294                 fs_related.makedirs(os.path.join(
 295                                         os.path.dirname(self._image),
 296                                         "LiveOS"))
 297                 shutil.move(self._image,
 298                             os.path.join(os.path.dirname(self._image),
 299                                          "LiveOS", "ext3fs.img"))
 300                 fs_related.mksquashfs(os.path.dirname(self._image),
 301                            isodir + "/LiveOS/squashfs.img")
 302
 303                 self._create_usbimg(isodir)
 304
 305                 if self.pack_to:
 306                     self.image_files.update({'image_files': self.pack_to})
 307                     usbimg = os.path.join(self._outdir, self.name + ".usbimg")
 308                     packimg = os.path.join(self._outdir, self.pack_to)
 309                     packing(packimg, usbimg)
 310                     os.unlink(usbimg)
 311                 else:
 312                     self.image_files.update({'image_files': self.name + ".usbimg"})
 313
 314         finally:
 315             shutil.rmtree(isodir, ignore_errors = True)
 316             self._set_isodir(None)
 317