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] == '#')
 139     assert(lines[1] == 'unit: sectors')
 140     assert(lines[2] == '')
 141     # This line has information about the first partition.
 142     chunks = lines[3].split()
 143     # The fourth chunk is the offset of the partition in sectors followed by
 144     # a comma. We drop the comma and convert that to an integer.
 145     sectors = string.atoi(chunks[3][:-1])
 146     # Free the loopback device and return an answer.
 147     dev.destroy()
 148     return sectors * BlockSize
 149
 150 def mountPointToDev(mountPoint):
 151     (mountTable, returncode) = getOutput([findProg('mount')])
 152     if returncode != 0:
 153         print mountTable
 154         exit(returncode)
 155     mountTable = mountTable.splitlines()
 156     for line in mountTable:
 157         chunks = line.split()
 158         if os.path.samefile(chunks[2], mountPoint):
 159             return LoopbackDevice(chunks[0])
 160     return None
 161
 162
 163 # Commands for the gem5img.py script
 164 commands = {}
 165 commandOrder = []
 166
 167 class Command(object):
 168     def addOption(self, *args, **kargs):
 169         self.parser.add_option(*args, **kargs)
 170
 171     def __init__(self, name, description, posArgs):
 172         self.name = name
 173         self.description = description
 174         self.func = None
 175         self.posArgs = posArgs
 176         commands[self.name] = self
 177         commandOrder.append(self.name)
 178         usage = 'usage: %prog [options]'
 179         posUsage = ''
 180         for posArg in posArgs:
 181             (argName, argDesc) = posArg
 182             usage += ' %s' % argName
 183             posUsage += '\n  %s: %s' % posArg
 184         usage += posUsage
 185         self.parser = OptionParser(usage=usage, description=description)
 186         self.addOption('-d', '--debug', dest='debug', action='store_true',
 187                        help='Verbose output.')
 188
 189     def parseArgs(self, argv):
 190         (self.options, self.args) = self.parser.parse_args(argv[2:])
 191         if len(self.args) != len(self.posArgs):
 192             self.parser.error('Incorrect number of arguments')
 193         global debug
 194         if self.options.debug:
 195             debug = True
 196
 197     def runCom(self):
 198         if not self.func:
 199             exit('Unimplemented command %s!' % self.name)
 200         self.func(self.options, self.args)
 201
 202
 203 # A command which prepares an image with an partition table and an empty file
 204 # system.
 205 initCom = Command('init', 'Create an image with an empty file system.',
 206                   [('file', 'Name of the image file.'),
 207                    ('mb', 'Size of the file in MB.')])
 208 initCom.addOption('-t', '--type', dest='fstype', action='store',
 209                   default='ext2',
 210                   help='Type of file system to use. Appended to mkfs.')
 211
 212 # A command to mount the first partition in the image.
 213 mountCom = Command('mount', 'Mount the first partition in the disk image.',
 214                    [('file', 'Name of the image file.'),
 215                     ('mount point', 'Where to mount the image.')])
 216
 217 def mountComFunc(options, args):
 218     (path, mountPoint) = args
 219     if not os.path.isdir(mountPoint):
 220         print "Mount point %s is not a directory." % mountPoint
 221
 222     dev = LoopbackDevice()
 223     if dev.setup(path, offset=True) != 0:
 224         exit(1)
 225
 226     if runPriv([findProg('mount'), str(dev), mountPoint]) != 0:
 227         dev.destroy()
 228         exit(1)
 229
 230 mountCom.func = mountComFunc
 231
 232 # A command to unmount the first partition in the image.
 233 umountCom = Command('umount', 'Unmount the first partition in the disk image.',
 234                     [('mount point', 'What mount point to unmount.')])
 235
 236 def umountComFunc(options, args):
 237     (mountPoint,) = args
 238     if not os.path.isdir(mountPoint):
 239         print "Mount point %s is not a directory." % mountPoint
 240         exit(1)
 241
 242     dev = mountPointToDev(mountPoint)
 243     if not dev:
 244         print "Unable to find mount information for %s." % mountPoint
 245
 246     # Unmount the loopback device.
 247     if runPriv([findProg('umount'), mountPoint]) != 0:
 248         exit(1)
 249
 250     # Destroy the loopback device.
 251     dev.destroy()
 252
 253 umountCom.func = umountComFunc
 254
 255
 256 # A command to create an empty file to hold the image.
 257 newCom = Command('new', 'File creation part of "init".',
 258                  [('file', 'Name of the image file.'),
 259                   ('mb', 'Size of the file in MB.')])
 260
 261 def newImage(file, mb):
 262     (cylinders, heads, sectors) = chsFromSize((mb * MB) / BlockSize)
 263     size = cylinders * heads * sectors * BlockSize
 264
 265     # We lseek to the end of the file and only write one byte there. This
 266     # leaves a "hole" which many file systems are smart enough not to actually
 267     # store to disk and which is defined to read as zero.
 268     fd = os.open(file, os.O_WRONLY | os.O_CREAT)
 269     os.lseek(fd, size - 1, os.SEEK_SET)
 270     os.write(fd, '\0')
 271
 272 def newComFunc(options, args):
 273     (file, mb) = args
 274     mb = string.atoi(mb)
 275     newImage(file, mb)
 276
 277
 278 newCom.func = newComFunc
 279
 280 # A command to partition the image file like a raw disk device.
 281 partitionCom = Command('partition', 'Partition part of "init".',
 282                        [('file', 'Name of the image file.')])
 283
 284 def partition(dev, cylinders, heads, sectors):
 285     # Use fdisk to partition the device
 286     comStr = '0,\n;\n;\n;\n'
 287     return runPriv([findProg('sfdisk'), '--no-reread', '-D', \
 288                    '-C', "%d" % cylinders, \
 289                    '-H', "%d" % heads, \
 290                    '-S', "%d" % sectors, \
 291                    str(dev)], inputVal=comStr)
 292
 293 def partitionComFunc(options, args):
 294     (path,) = args
 295
 296     dev = LoopbackDevice()
 297     if dev.setup(path) != 0:
 298         exit(1)
 299
 300     # Figure out the dimensions of the file.
 301     size = os.path.getsize(path)
 302     if partition(dev, *chsFromSize(size / BlockSize)) != 0:
 303         dev.destroy()
 304         exit(1)
 305
 306     dev.destroy()
 307
 308 partitionCom.func = partitionComFunc
 309
 310 # A command to format the first partition in the image.
 311 formatCom = Command('format', 'Formatting part of "init".',
 312                     [('file', 'Name of the image file.')])
 313 formatCom.addOption('-t', '--type', dest='fstype', action='store',
 314                     default='ext2',
 315                     help='Type of file system to use. Appended to mkfs.')
 316
 317 def formatImage(dev, fsType):
 318     return runPriv([findProg('mkfs.%s' % fsType, dev), str(dev)])
 319
 320 def formatComFunc(options, args):
 321     (path,) = args
 322
 323     dev = LoopbackDevice()
 324     if dev.setup(path, offset=True) != 0:
 325         exit(1)
 326
 327     # Format the device.
 328     if formatImage(dev, options.fstype) != 0:
 329         dev.destroy()
 330         exit(1)
 331
 332     dev.destroy()
 333
 334 formatCom.func = formatComFunc
 335
 336 def initComFunc(options, args):
 337     (path, mb) = args
 338     mb = string.atoi(mb)
 339     newImage(path, mb)
 340     dev = LoopbackDevice()
 341     if dev.setup(path) != 0:
 342         exit(1)
 343     size = os.path.getsize(path)
 344     if partition(dev, *chsFromSize((mb * MB) / BlockSize)) != 0:
 345         dev.destroy()
 346         exit(1)
 347     dev.destroy()
 348     if dev.setup(path, offset=True) != 0:
 349         exit(1)
 350     if formatImage(dev, options.fstype) != 0:
 351         dev.destroy()
 352         exit(1)
 353     dev.destroy()
 354
 355 initCom.func = initComFunc
 356
 357
 358 # Figure out what command was requested and execute it.
 359 if len(argv) < 2 or argv[1] not in commands:
 360     print 'Usage: %s [command] <command arguments>'
 361     print 'where [command] is one of '
 362     for name in commandOrder:
 363         command = commands[name]
 364         print '    %s: %s' % (command.name, command.description)
 365     print 'Watch for orphaned loopback devices and delete them with'
 366     print 'losetup -d. Mounted images will belong to root, so you may need'
 367     print 'to use sudo to modify their contents.'
 368     exit(1)
 369
 370 command = commands[argv[1]]
 371 command.parseArgs(argv)
 372 command.runCom()