util/gem5img.py

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