util/gem5img.py

   1 #!/usr/bin/python2.7
   2 #
   3 # gem5img.py
   4 # Script for managing a gem5 disk image.
   5 #
   6
   7 from optparse import OptionParser
   8 import os
   9 from os import environ as env
  10 import string
  11 from subprocess import CalledProcessError, Popen, PIPE, STDOUT
  12 from sys import exit, argv
  13
  14
  15 # Some constants.
  16 MaxLBACylinders = 16383
  17 MaxLBAHeads = 16
  18 MaxLBASectors = 63
  19 MaxLBABlocks = MaxLBACylinders * MaxLBAHeads * MaxLBASectors
  20
  21 BlockSize = 512
  22 MB = 1024 * 1024
  23
  24 # Setup PATH to look in the sbins.
  25 env['PATH'] += ':/sbin:/usr/sbin'
  26
  27 # Whether to print debug output.
  28 debug = False
  29
  30 # Figure out cylinders, heads and sectors from a size in blocks.
  31 def chsFromSize(sizeInBlocks):
  32     if sizeInBlocks >= MaxLBABlocks:
  33         sizeInMBs = (sizeInBlocks * BlockSize) / MB
  34         print '%d MB is too big for LBA, truncating file.' % sizeInMBs
  35         return (MaxLBACylinders, MaxLBAHeads, MaxLBASectors)
  36
  37     sectors = sizeInBlocks
  38     if sizeInBlocks > 63:
  39         sectors = 63
  40
  41     headSize = sizeInBlocks / sectors
  42     heads = 16
  43     if headSize < 16:
  44         heads = sizeInBlocks
  45
  46     cylinders = sizeInBlocks / (sectors * heads)
  47
  48     return (cylinders, heads, sectors)
  49
  50
  51 # Figure out if we should use sudo.
  52 def needSudo():
  53     if not hasattr(needSudo, 'notRoot'):
  54         needSudo.notRoot = (os.geteuid() != 0)
  55         if needSudo.notRoot:
  56             print 'You are not root. Using sudo.'
  57     return needSudo.notRoot
  58
  59 # Run an external command.
  60 def runCommand(command, inputVal=''):
  61     print "%>", ' '.join(command)
  62     proc = Popen(command, stdin=PIPE)
  63     proc.communicate(inputVal)
  64     return proc.returncode
  65
  66 # Run an external command and capture its output. This is intended to be
  67 # used with non-interactive commands where the output is for internal use.
  68 def getOutput(command, inputVal=''):
  69     global debug
  70     if debug:
  71         print "%>", ' '.join(command)
  72     proc = Popen(command, stderr=STDOUT,
  73                  stdin=PIPE, stdout=PIPE)
  74     (out, err) = proc.communicate(inputVal)
  75     return (out, proc.returncode)
  76
  77 # Run a command as root, using sudo if necessary.
  78 def runPriv(command, inputVal=''):
  79     realCommand = command
  80     if needSudo():
  81         realCommand = [findProg('sudo')] + command
  82     return runCommand(realCommand, inputVal)
  83
  84 def privOutput(command, inputVal=''):
  85     realCommand = command
  86     if needSudo():
  87         realCommand = [findProg('sudo')] + command
  88     return getOutput(realCommand, inputVal)
  89
  90 # Find the path to a program.
  91 def findProg(program, cleanupDev=None):
  92     (out, returncode) = getOutput(['which', program])
  93     if returncode != 0:
  94         if cleanupDev:
  95             cleanupDev.destroy()
  96         exit("Unable to find program %s, check your PATH variable." % program)
  97     return string.strip(out)
  98
  99 class LoopbackDevice(object):
 100     def __init__(self, devFile=None):
 101         self.devFile = devFile
 102     def __str__(self):
 103         return str(self.devFile)
 104
 105     def setup(self, fileName, offset=False):
 106         assert not self.devFile
 107         (out, returncode) = privOutput([findProg('losetup'), '-f'])
 108         if returncode != 0:
 109             print out
 110             return returncode
 111         self.devFile = string.strip(out)
 112         command = [findProg('losetup'), self.devFile, fileName]
 113         if offset:
 114             off = findPartOffset(self.devFile, fileName, 0)
 115             command = command[:1] + \
 116                       ["-o", "%d" % off] + \
 117                       command[1:]
 118         return runPriv(command)
 119
 120     def destroy(self):
 121         assert self.devFile
 122         returncode = runPriv([findProg('losetup'), '-d', self.devFile])
 123         self.devFile = None
 124         return returncode
 125
 126 def findPartOffset(devFile, fileName, partition):
 127     # Attach a loopback device to the file so we can use sfdisk on it.
 128     dev = LoopbackDevice()
 129     dev.setup(fileName)
 130     # Dump the partition information.
 131     command = [findProg('sfdisk'), '-d', dev.devFile]
 132     (out, returncode) = privOutput(command)
 133     if returncode != 0:
 134         print out
 135         exit(returncode)
 136     lines = out.splitlines()
 137     # Make sure the first few lines of the output look like what we expect.
 138     assert(lines[0][0] == '#' or lines[0].startswith('label:'))
 139     assert(lines[1] == 'unit: sectors' or lines[1].startswith('label-id:'))
 140     assert(lines[2] == '' or lines[2].startswith('device:'))
 141     if lines[0][0] == '#' :
 142         # Parsing an 'old style' dump oputput
 143         # Line 4 has information about the first partition.
 144         chunks = lines[3].split()
 145     else :
 146         # Parsing a 'new style' dump oputput
 147         # Line 6 has information about the first partition.
 148         chunks = lines[5].split()
 149     # The fourth chunk is the offset of the partition in sectors followed by
 150     # a comma. We drop the comma and convert that to an integer.
 151     sectors = string.atoi(chunks[3][:-1])
 152     # Free the loopback device and return an answer.
 153     dev.destroy()
 154     return sectors * BlockSize
 155
 156 def mountPointToDev(mountPoint):
 157     (mountTable, returncode) = getOutput([findProg('mount')])
 158     if returncode != 0:
 159         print mountTable
 160         exit(returncode)
 161     mountTable = mountTable.splitlines()
 162     for line in mountTable:
 163         chunks = line.split()
 164         if os.path.samefile(chunks[2], mountPoint):
 165             return LoopbackDevice(chunks[0])
 166     return None
 167
 168
 169 # Commands for the gem5img.py script
 170 commands = {}
 171 commandOrder = []
 172
 173 class Command(object):
 174     def addOption(self, *args, **kargs):
 175         self.parser.add_option(*args, **kargs)
 176
 177     def __init__(self, name, description, posArgs):
 178         self.name = name
 179         self.description = description
 180         self.func = None
 181         self.posArgs = posArgs
 182         commands[self.name] = self
 183         commandOrder.append(self.name)
 184         usage = 'usage: %prog [options]'
 185         posUsage = ''
 186         for posArg in posArgs:
 187             (argName, argDesc) = posArg
 188             usage += ' %s' % argName
 189             posUsage += '\n  %s: %s' % posArg
 190         usage += posUsage
 191         self.parser = OptionParser(usage=usage, description=description)
 192         self.addOption('-d', '--debug', dest='debug', action='store_true',
 193                        help='Verbose output.')
 194
 195     def parseArgs(self, argv):
 196         (self.options, self.args) = self.parser.parse_args(argv[2:])
 197         if len(self.args) != len(self.posArgs):
 198             self.parser.error('Incorrect number of arguments')
 199         global debug
 200         if self.options.debug:
 201             debug = True
 202
 203     def runCom(self):
 204         if not self.func:
 205             exit('Unimplemented command %s!' % self.name)
 206         self.func(self.options, self.args)
 207
 208
 209 # A command which prepares an image with an partition table and an empty file
 210 # system.
 211 initCom = Command('init', 'Create an image with an empty file system.',
 212                   [('file', 'Name of the image file.'),
 213                    ('mb', 'Size of the file in MB.')])
 214 initCom.addOption('-t', '--type', dest='fstype', action='store',
 215                   default='ext2',
 216                   help='Type of file system to use. Appended to mkfs.')
 217
 218 # A command to mount the first partition in the image.
 219 mountCom = Command('mount', 'Mount the first partition in the disk image.',
 220                    [('file', 'Name of the image file.'),
 221                     ('mount point', 'Where to mount the image.')])
 222
 223 def mountComFunc(options, args):
 224     (path, mountPoint) = args
 225     if not os.path.isdir(mountPoint):
 226         print "Mount point %s is not a directory." % mountPoint
 227
 228     dev = LoopbackDevice()
 229     if dev.setup(path, offset=True) != 0:
 230         exit(1)
 231
 232     if runPriv([findProg('mount'), str(dev), mountPoint]) != 0:
 233         dev.destroy()
 234         exit(1)
 235
 236 mountCom.func = mountComFunc
 237
 238 # A command to unmount the first partition in the image.
 239 umountCom = Command('umount', 'Unmount the first partition in the disk image.',
 240                     [('mount point', 'What mount point to unmount.')])
 241
 242 def umountComFunc(options, args):
 243     (mountPoint,) = args
 244     if not os.path.isdir(mountPoint):
 245         print "Mount point %s is not a directory." % mountPoint
 246         exit(1)
 247
 248     dev = mountPointToDev(mountPoint)
 249     if not dev:
 250         print "Unable to find mount information for %s." % mountPoint
 251
 252     # Unmount the loopback device.
 253     if runPriv([findProg('umount'), mountPoint]) != 0:
 254         exit(1)
 255
 256     # Destroy the loopback device.
 257     dev.destroy()
 258
 259 umountCom.func = umountComFunc
 260
 261
 262 # A command to create an empty file to hold the image.
 263 newCom = Command('new', 'File creation part of "init".',
 264                  [('file', 'Name of the image file.'),
 265                   ('mb', 'Size of the file in MB.')])
 266
 267 def newImage(file, mb):
 268     (cylinders, heads, sectors) = chsFromSize((mb * MB) / BlockSize)
 269     size = cylinders * heads * sectors * BlockSize
 270
 271     # We lseek to the end of the file and only write one byte there. This
 272     # leaves a "hole" which many file systems are smart enough not to actually
 273     # store to disk and which is defined to read as zero.
 274     fd = os.open(file, os.O_WRONLY | os.O_CREAT)
 275     os.lseek(fd, size - 1, os.SEEK_SET)
 276     os.write(fd, '\0')
 277
 278 def newComFunc(options, args):
 279     (file, mb) = args
 280     mb = string.atoi(mb)
 281     newImage(file, mb)
 282
 283
 284 newCom.func = newComFunc
 285
 286 # A command to partition the image file like a raw disk device.
 287 partitionCom = Command('partition', 'Partition part of "init".',
 288                        [('file', 'Name of the image file.')])
 289
 290 def partition(dev, cylinders, heads, sectors):
 291     # Use sfdisk to partition the device
 292     # The specified options are intended to work with both new and old
 293     # versions of sfdisk (see https://askubuntu.com/a/819614)
 294     comStr = ';'
 295     return runPriv([findProg('sfdisk'), '--no-reread', '-u', 'S', '-L', \
 296                    str(dev)], inputVal=comStr)
 297
 298 def partitionComFunc(options, args):
 299     (path,) = args
 300
 301     dev = LoopbackDevice()
 302     if dev.setup(path) != 0:
 303         exit(1)
 304
 305     # Figure out the dimensions of the file.
 306     size = os.path.getsize(path)
 307     if partition(dev, *chsFromSize(size / BlockSize)) != 0:
 308         dev.destroy()
 309         exit(1)
 310
 311     dev.destroy()
 312
 313 partitionCom.func = partitionComFunc
 314
 315 # A command to format the first partition in the image.
 316 formatCom = Command('format', 'Formatting part of "init".',
 317                     [('file', 'Name of the image file.')])
 318 formatCom.addOption('-t', '--type', dest='fstype', action='store',
 319                     default='ext2',
 320                     help='Type of file system to use. Appended to mkfs.')
 321
 322 def formatImage(dev, fsType):
 323     return runPriv([findProg('mkfs.%s' % fsType, dev), str(dev)])
 324
 325 def formatComFunc(options, args):
 326     (path,) = args
 327
 328     dev = LoopbackDevice()
 329     if dev.setup(path, offset=True) != 0:
 330         exit(1)
 331
 332     # Format the device.
 333     if formatImage(dev, options.fstype) != 0:
 334         dev.destroy()
 335         exit(1)
 336
 337     dev.destroy()
 338
 339 formatCom.func = formatComFunc
 340
 341 def initComFunc(options, args):
 342     (path, mb) = args
 343     mb = string.atoi(mb)
 344     newImage(path, mb)
 345     dev = LoopbackDevice()
 346     if dev.setup(path) != 0:
 347         exit(1)
 348     size = os.path.getsize(path)
 349     if partition(dev, *chsFromSize((mb * MB) / BlockSize)) != 0:
 350         dev.destroy()
 351         exit(1)
 352     dev.destroy()
 353     if dev.setup(path, offset=True) != 0:
 354         exit(1)
 355     if formatImage(dev, options.fstype) != 0:
 356         dev.destroy()
 357         exit(1)
 358     dev.destroy()
 359
 360 initCom.func = initComFunc
 361
 362
 363 # Figure out what command was requested and execute it.
 364 if len(argv) < 2 or argv[1] not in commands:
 365     print 'Usage: %s [command] <command arguments>'
 366     print 'where [command] is one of '
 367     for name in commandOrder:
 368         command = commands[name]
 369         print '    %s: %s' % (command.name, command.description)
 370     print 'Watch for orphaned loopback devices and delete them with'
 371     print 'losetup -d. Mounted images will belong to root, so you may need'
 372     print 'to use sudo to modify their contents.'
 373     exit(1)
 374
 375 command = commands[argv[1]]
 376 command.parseArgs(argv)
 377 command.runCom()