util/gem5img.py

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