1#!/usr/bin/python2.7
2#
3# gem5img.py
4# Script for managing a gem5 disk image.
5#
6
7from optparse import OptionParser
8import os
9from os import environ as env
10import string
11from subprocess import CalledProcessError, Popen, PIPE, STDOUT
12from sys import exit, argv
13
14
15# Some constants.
16MaxLBACylinders = 16383
17MaxLBAHeads = 16
18MaxLBASectors = 63
19MaxLBABlocks = MaxLBACylinders * MaxLBAHeads * MaxLBASectors
20
21BlockSize = 512
22MB = 1024 * 1024
23
24# Setup PATH to look in the sbins.
25env['PATH'] += ':/sbin:/usr/sbin'
26
27# Whether to print debug output.
28debug = False
29
30# Figure out cylinders, heads and sectors from a size in blocks.
31def 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.
52def 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.
60def 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.
68def 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.
78def runPriv(command, inputVal=''):
79    realCommand = command
80    if needSudo():
81        realCommand = [findProg('sudo')] + command
82    return runCommand(realCommand, inputVal)
83
84def 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.
91def 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
99class 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
126def 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
150def 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
164commands = {}
165commandOrder = []
166
167class 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.
205initCom = 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.')])
208initCom.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.
213mountCom = 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
217def 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
230mountCom.func = mountComFunc
231
232# A command to unmount the first partition in the image.
233umountCom = Command('umount', 'Unmount the first partition in the disk image.',
234                    [('mount point', 'What mount point to unmount.')])
235
236def 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
253umountCom.func = umountComFunc
254
255
256# A command to create an empty file to hold the image.
257newCom = Command('new', 'File creation part of "init".',
258                 [('file', 'Name of the image file.'),
259                  ('mb', 'Size of the file in MB.')])
260
261def 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
272def newComFunc(options, args):
273    (file, mb) = args
274    mb = string.atoi(mb)
275    newImage(file, mb)
276
277
278newCom.func = newComFunc
279
280# A command to partition the image file like a raw disk device.
281partitionCom = Command('partition', 'Partition part of "init".',
282                       [('file', 'Name of the image file.')])
283
284def 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
293def 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
308partitionCom.func = partitionComFunc
309
310# A command to format the first partition in the image.
311formatCom = Command('format', 'Formatting part of "init".',
312                    [('file', 'Name of the image file.')])
313formatCom.addOption('-t', '--type', dest='fstype', action='store',
314                    default='ext2',
315                    help='Type of file system to use. Appended to mkfs.')
316
317def formatImage(dev, fsType):
318    return runPriv([findProg('mkfs.%s' % fsType, dev), str(dev)])
319
320def 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
334formatCom.func = formatComFunc
335
336def 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
355initCom.func = initComFunc
356
357
358# Figure out what command was requested and execute it.
359if 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
370command = commands[argv[1]]
371command.parseArgs(argv)
372command.runCom()
373