Statistics
| Branch: | Tag: | Revision:

root / snf-image-helper / disklabel.py @ 1de1eff5

History | View | Annotate | Download (19.4 kB)

1
#!/usr/bin/env python
2
#
3
# -*- coding: utf-8 -*-
4
#
5
# Copyright (C) 2013 GRNET S.A.
6
#
7
# This program is free software; you can redistribute it and/or modify
8
# it under the terms of the GNU General Public License as published by
9
# the Free Software Foundation; either version 2 of the License, or
10
# (at your option) any later version.
11
#
12
# This program is distributed in the hope that it will be useful, but
13
# WITHOUT ANY WARRANTY; without even the implied warranty of
14
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
15
# General Public License for more details.
16
#
17
# You should have received a copy of the GNU General Public License
18
# along with this program; if not, write to the Free Software
19
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
20
# 02110-1301, USA.
21

    
22
"""This module provides the code for handling OpenBSD disklabels"""
23

    
24
import struct
25
import sys
26
import cStringIO
27
import optparse
28

    
29
from collections import namedtuple
30

    
31
BLOCKSIZE = 512
32

    
33
LABELSECTOR = 1
34
LABELOFFSET = 0
35

    
36
BBSIZE = 8192  # size of boot area with label
37
SBSIZE = 8192  # max size of fs superblock
38

    
39
DISKMAGIC = 0x82564557
40

    
41

    
42
class MBR(object):
43
    """Represents a Master Boot Record."""
44
    class Partition(object):
45
        """Represents a partition entry in MBR"""
46
        format = "<B3sB3sLL"
47

    
48
        def __init__(self, raw_part):
49
            """Create a Partition instance"""
50
            (
51
                self.status,
52
                self.start,
53
                self.type,
54
                self.end,
55
                self.first_sector,
56
                self.sector_count
57
            ) = struct.unpack(self.format, raw_part)
58

    
59
        def pack(self):
60
            """Pack the partition values into a binary string"""
61
            return struct.pack(self.format,
62
                               self.status,
63
                               self.start,
64
                               self.type,
65
                               self.end,
66
                               self.first_sector,
67
                               self.sector_count)
68

    
69
        @staticmethod
70
        def size():
71
            """Returns the size of an MBR partition entry"""
72
            return struct.calcsize(MBR.Partition.format)
73

    
74
        def __str__(self):
75
            start = self.unpack_chs(self.start)
76
            end = self.unpack_chs(self.end)
77
            return "%d %s %d %s %d %d" % (self.status, start, self.type, end,
78
                                          self.first_sector, self.sector_count)
79

    
80
        @staticmethod
81
        def unpack_chs(chs):
82
            """Unpacks a CHS address string to a tuple."""
83

    
84
            assert len(chs) == 3
85

    
86
            head = struct.unpack('<B', chs[0])[0]
87
            sector = struct.unpack('<B', chs[1])[0] & 0x3f
88
            cylinder = (struct.unpack('<B', chs[1])[0] & 0xC0) << 2 | \
89
                struct.unpack('<B', chs[2])[0]
90

    
91
            return (cylinder, head, sector)
92

    
93
        @staticmethod
94
        def pack_chs(cylinder, head, sector):
95
            """Packs a CHS tuple to an address string."""
96

    
97
            assert 1 <= sector <= 63
98
            assert 0 <= head <= 255
99
            assert 0 <= cylinder
100

    
101
            # If the cylinders overflow then put the value (1023, 254, 63) to
102
            # the tuple. At least this is what OpenBSD does.
103
            if cylinder > 1023:
104
                cylinder = 1023
105
                head = 254
106
                sector = 63
107

    
108
            byte0 = head
109
            byte1 = (cylinder >> 2) & 0xC0 | sector
110
            byte2 = cylinder & 0xff
111

    
112
            return struct.pack('<BBB', byte0, byte1, byte2)
113

    
114
    format = "<444s2x16s16s16s16s2s"
115
    """
116
    Offset  Length          Contents
117
    0       440(max. 446)   code area
118
    440     2(optional)     disk signature
119
    444     2               Usually nulls
120
    446     16              Partition 0
121
    462     16              Partition 1
122
    478     16              Partition 2
123
    494     16              Partition 3
124
    510     2               MBR signature
125
    """
126
    def __init__(self, block):
127
        """Create an MBR instance"""
128
        raw_part = {}
129
        (self.code_area,
130
         raw_part[0],
131
         raw_part[1],
132
         raw_part[2],
133
         raw_part[3],
134
         self.signature) = struct.unpack(self.format, block)
135

    
136
        self.part = {}
137
        for i in range(4):
138
            self.part[i] = self.Partition(raw_part[i])
139

    
140
    @staticmethod
141
    def size():
142
        """Return the size of a Master Boot Record."""
143
        return struct.calcsize(MBR.format)
144

    
145
    def pack(self):
146
        """Pack an MBR to a binary string."""
147
        return struct.pack(self.format,
148
                           self.code_area,
149
                           self.part[0].pack(),
150
                           self.part[1].pack(),
151
                           self.part[2].pack(),
152
                           self.part[3].pack(),
153
                           self.signature)
154

    
155
    def __str__(self):
156
        ret = ""
157
        for i in range(4):
158
            ret += "Partition %d: %s\n" % (i, self.part[i])
159
        ret += "Signature: %s %s\n" % (hex(ord(self.signature[0])),
160
                                       hex(ord(self.signature[1])))
161
        return ret
162

    
163

    
164
class Disklabel:
165
    """Represents an OpenBSD Disklabel"""
166
    format = "<IHH16s16sIIIIII8sIHHIII20sHH16sIHHII364s"
167
    """
168
    Offset  Length          Contents
169
    0       4               Magic
170
    4       2               Drive Type
171
    6       2               Subtype
172
    8       16              Type Name
173
    24      16              Pack Identifier
174
    32      4               Bytes per sector
175
    36      4               Data sectors per track
176
    40      4               Tracks per cilinder
177
    44      4               Data cylinders per unit
178
    48      4               Data sectors per cylynder
179
    52      4               Data sectors per unit
180
    56      8               Unique label identifier
181
    64      4               Alt cylinders per unit
182
    68      2               Start of useable region (high part)
183
    70      2               Size of usable region (high part)
184
    72      4               Start of useable region
185
    76      4               End of usable region
186
    80      4               Generic Flags
187
    84      5*4             Drive-type specific information
188
    104     2               Number of data sectors (high part)
189
    106     2               Version
190
    108     4*4             Reserved for future use
191
    124     4               Magic number
192
    128     2               Xor of data Inclu. partitions
193
    130     2               Number of partitions in following
194
    132     4               size of boot area at sn0, bytes
195
    136     4               Max size of fs superblock, bytes
196
    140     16*16           Partition Table
197
    """
198

    
199
    class PartitionTable:
200
        """Reprepsents an OpenBSD Partition Table"""
201
        format = "<IIHHBBH"
202
        """
203
        Partition Entry:
204
        Offset  Length          Contents
205
        0       4               Number of sectors in the partition
206
        4       4               Starting sector
207
        8       2               Starting sector (high part)
208
        10      2               Number of sectors (high part)
209
        12      1               Filesystem type
210
        13      1               Filesystem Fragment per block
211
        14      2               FS cylinders per group
212
        """
213

    
214
        Partition = namedtuple(
215
            'Partition', 'size, offset, offseth, sizeh, fstype, frag, cpg')
216

    
217
        def __init__(self, ptable, pnumber):
218
            """Create a Partition Table instance"""
219
            self.part = []
220

    
221
            size = struct.calcsize(self.format)
222

    
223
            raw = cStringIO.StringIO(ptable)
224
            try:
225
                for i in range(pnumber):
226
                    p = self.Partition(
227
                        *struct.unpack(self.format, raw.read(size)))
228
                    self.part.append(p)
229
            finally:
230
                raw.close()
231

    
232
        def __str__(self):
233
            """Print the Partition table"""
234
            val = ""
235
            for i in range(len(self.part)):
236
                val += "%c: %s\n" % (chr(ord('a') + i), str(self.part[i]))
237
            return val
238

    
239
        def pack(self):
240
            """Packs the partition table into a binary string."""
241
            ret = ""
242
            for i in range(len(self.part)):
243
                ret += struct.pack(self.format,
244
                                   self.part[i].size,
245
                                   self.part[i].offset,
246
                                   self.part[i].offseth,
247
                                   self.part[i].sizeh,
248
                                   self.part[i].fstype,
249
                                   self.part[i].frag,
250
                                   self.part[i].cpg)
251
            return ret
252

    
253
        def setpsize(self, i, size):
254
            """Set size for partition i"""
255
            tmp = self.part[i]
256
            self.part[i] = self.Partition(size & 0xffffffff, tmp.offset,
257
                                          tmp.offseth, size >> 32, tmp.fstype,
258
                                          tmp.frag, tmp.cpg)
259

    
260
        def getpsize(self, i):
261
            return (self.part[i].sizeh << 32) + self.part[i].size
262

    
263
        def setpoffset(self, i, offset):
264
            """Set  offset for partition i"""
265
            tmp = self.part[i]
266
            self.part[i] = self.Partition(tmp.size, offset & 0xffffffff,
267
                                          offset >> 32, tmp.sizeh, tmp.frag,
268
                                          tmp.cpg)
269

    
270
        def getpoffset(self, i):
271
            return (self.part[i].offseth << 32) + self.part[i].offset
272

    
273
    def __init__(self, disk):
274
        """Create a DiskLabel instance"""
275
        self.disk = disk
276
        self.part_num = None
277

    
278
        with open(disk, "rb") as d:
279
            sector0 = d.read(BLOCKSIZE)
280
            self.mbr = MBR(sector0)
281

    
282
            for i in range(4):
283
                if self.mbr.part[i].type == 0xa6:  # OpenBSD type
284
                    self.part_num = i
285
                    break
286

    
287
            assert self.part_num is not None, "No OpenBSD partition found"
288

    
289
            d.seek(BLOCKSIZE * self.mbr.part[self.part_num].first_sector)
290
            part_sector0 = d.read(BLOCKSIZE)
291
            # The offset of the disklabel from the begining of the
292
            # partition is one sector
293
            part_sector1 = d.read(BLOCKSIZE)
294

    
295
        (self.magic,
296
         self.dtype,
297
         self.subtype,
298
         self.typename,
299
         self.packname,
300
         self.secsize,
301
         self.nsectors,
302
         self.ntracks,
303
         self.ncylinders,
304
         self.secpercyl,
305
         self.secperunit,
306
         self.uid,
307
         self.acylinders,
308
         self.bstarth,
309
         self.bendh,
310
         self.bstart,
311
         self.bend,
312
         self.flags,
313
         self.drivedata,
314
         self.secperunith,
315
         self.version,
316
         self.spare,
317
         self.magic2,
318
         self.checksum,
319
         self.npartitions,
320
         self.bbsize,
321
         self.sbsize,
322
         ptable_raw) = struct.unpack(self.format, part_sector1)
323

    
324
        assert self.magic == DISKMAGIC, "Disklabel is not valid"
325

    
326
        self.ptable = self.PartitionTable(ptable_raw, self.npartitions)
327

    
328
    def pack(self, checksum=None):
329
        return struct.pack(self.format,
330
                           self.magic,
331
                           self.dtype,
332
                           self.subtype,
333
                           self.typename,
334
                           self.packname,
335
                           self.secsize,
336
                           self.nsectors,
337
                           self.ntracks,
338
                           self.ncylinders,
339
                           self.secpercyl,
340
                           self.secperunit,
341
                           self.uid,
342
                           self.acylinders,
343
                           self.bstarth,
344
                           self.bendh,
345
                           self.bstart,
346
                           self.bend,
347
                           self.flags,
348
                           self.drivedata,
349
                           self.secperunith,
350
                           self.version,
351
                           self.spare,
352
                           self.magic2,
353
                           self.checksum if checksum is None else checksum,
354
                           self.npartitions,
355
                           self.bbsize,
356
                           self.sbsize,
357
                           self.ptable.pack() +
358
                           ((364 - self.npartitions * 16) * '\x00'))
359

    
360
    def compute_checksum(self):
361
        """Compute the checksum of the disklabel"""
362

    
363
        raw = cStringIO.StringIO(self.pack(0))
364
        checksum = 0
365
        try:
366
            uint16 = raw.read(2)
367
            while uint16 != "":
368
                checksum ^= struct.unpack('<H', uint16)[0]
369
                uint16 = raw.read(2)
370
        finally:
371
            raw.close()
372

    
373
        return checksum
374

    
375
    def setdsize(self, dsize):
376
        """Set disk size"""
377
        self.secperunith = dsize >> 32
378
        self.secperunit = dsize & 0xffffffff
379

    
380
    def getdsize(self):
381
        """Get disk size"""
382
        return (self.secperunith << 32) + self.secperunit
383

    
384
    def setbstart(self, bstart):
385
        """Set start of useable region"""
386
        self.bstarth = bstart >> 32
387
        self.bstart = bstart & 0xffffffff
388

    
389
    def getbstart(self):
390
        """Get start of usable region"""
391
        return (self.bstarth << 32) + self.bstart
392

    
393
    def setbend(self, bend):
394
        """Set end of useable region"""
395
        self.bendh = bend >> 32
396
        self.bend = bend & 0xffffffff
397

    
398
    def getbend(self):
399
        return (self.bendh << 32) + self.bend
400

    
401
    def enlarge_disk(self, new_size):
402
        """Enlarge the size of the disk"""
403

    
404
        assert new_size >= self.secperunit, \
405
            "New size cannot be smaller that %s" % self.secperunit
406

    
407
        # Fix the disklabel
408
        self.setdsize(new_size)
409
        self.ncylinders = self.getdsize() // (self.nsectors * self.ntracks)
410
        self.setbend(self.ncylinders * self.nsectors * self.ntracks)
411

    
412
        # Partition 'c' descriptes the entire disk
413
        self.ptable.setpsize(2, new_size)
414

    
415
        # Fix the MBR table
416
        start = self.mbr.part[self.part_num].first_sector
417
        self.mbr.part[self.part_num].sector_count = self.getbend() - start
418

    
419
        lba = self.getbend() - 1
420
        cylinder = lba // (self.ntracks * self.nsectors)
421
        header = (lba // self.nsectors) % self.ntracks
422
        sector = (lba % self.nsectors) + 1
423
        chs = MBR.Partition.pack_chs(cylinder, header, sector)
424
        self.mbr.part[self.part_num].end = chs
425

    
426
        self.checksum = self.compute_checksum()
427

    
428
    def write(self):
429
        """Write the disklabel back to the media"""
430
        with open(self.disk, 'rw+b') as d:
431
            d.write(self.mbr.pack())
432

    
433
            d.seek((self.mbr.part[self.part_num].first_sector + 1) * BLOCKSIZE)
434
            d.write(self.pack())
435

    
436
    def get_last_partition_id(self):
437
        """Returns the id of the last partition"""
438
        part = 0
439
        end = 0
440
        # Don't check partition 'c' which is the whole disk
441
        for i in filter(lambda x: x != 2, range(self.npartitions)):
442
            curr_end = self.ptable.getpsize(i) + self.ptable.getpoffset(i)
443
            if end < curr_end:
444
                end = curr_end
445
                part = i
446

    
447
        assert end > 0, "No partition found"
448

    
449
        return part
450

    
451
    def enlarge_last_partition(self):
452
        """Enlarge the last partition to cover up all the free space"""
453

    
454
        part_num = self.get_last_partition_id()
455

    
456
        end = self.ptable.getpsize(part_num) + self.ptable.getpoffset(part_num)
457

    
458
        assert end > 0, "No partition found"
459

    
460
        if self.ptable.part[part_num].fstype == 1:  # Swap partition.
461
            #TODO: Maybe create a warning?
462
            return
463

    
464
        if end > (self.getbend() - 1024):
465
            return
466

    
467
        self.ptable.setpsize(
468
            part_num, self.getbend() - self.ptable.getpoffset(part_num) - 1024)
469

    
470
        self.checksum = self.compute_checksum()
471

    
472
    def __str__(self):
473
        """Print the Disklabel"""
474
        title1 = "Master Boot Record"
475
        title2 = "Disklabel"
476

    
477
        return \
478
            "%s\n%s\n%s\n" % (title1, len(title1) * "=", str(self.mbr)) + \
479
            "%s\n%s\n" % (title2, len(title2) * "=") + \
480
            "Magic Number: 0x%x\n" % self.magic + \
481
            "Drive type: %d\n" % self.dtype + \
482
            "Subtype: %d\n" % self.subtype + \
483
            "Typename: %s\n" % self.typename.strip('\x00').strip() + \
484
            "Pack Identifier: %s\n" % self.packname.strip('\x00').strip() + \
485
            "Number of bytes per sector: %d\n" % self.secsize + \
486
            "Number of data sectors per track: %d\n" % self.nsectors + \
487
            "Number of tracks per cylinder: %d\n" % self.ntracks + \
488
            "Number of data cylinders per unit: %d\n" % self.ncylinders + \
489
            "Number of data sectors per cylinder: %d\n" % self.secpercyl + \
490
            "Number of data sectors per unit: %d\n" % self.secperunit + \
491
            "DUID: %s\n" % "".join(x.encode('hex') for x in self.uid) + \
492
            "Alt. cylinders per unit: %d\n" % self.acylinders + \
493
            "Start of useable region (high part): %d\n" % self.bstarth + \
494
            "Size of useable region (high part): %d\n" % self.bendh + \
495
            "Start of useable region: %d\n" % self.bstart + \
496
            "End of usable region: %d\n" % self.bend + \
497
            "Generic Flags: %r\n" % self.flags + \
498
            "Drive data: %r\n" % self.drivedata + \
499
            "Number of data sectors (high part): %d\n" % self.secperunith + \
500
            "Version: %d\n" % self.version + \
501
            "Reserved for future use: %r\n" % self.spare + \
502
            "The magic number again: 0x%x\n" % self.magic2 + \
503
            "Checksum: %d\n" % self.checksum + \
504
            "Number of partitions: %d\n" % self.npartitions + \
505
            "Size of boot aread at sn0: %d\n" % self.bbsize + \
506
            "Max size of fs superblock: %d\n" % self.sbsize + \
507
            "%s" % self.ptable
508

    
509

    
510
if __name__ == '__main__':
511

    
512
    usage = "Usage: %prog [options] <input_media>"
513
    parser = optparse.OptionParser(usage=usage)
514

    
515
    parser.add_option("-l", "--list", action="store_true", dest="list",
516
                      default=False,
517
                      help="list the disklabel on the specified media")
518
    parser.add_option("--get-last-partition", action="store_true",
519
                      dest="last_part", default=False,
520
                      help="print the label of the last partition")
521
    parser.add_option("--get-duid", action="store_true", dest="duid",
522
                      default=False,
523
                      help="print the disklabel unique identifier")
524
    parser.add_option("-d", "--enlarge-disk", type="int", dest="disk_size",
525
                      default=None, metavar="SIZE",
526
                      help="Enlarge the disk to this SIZE (in sectors)")
527
    parser.add_option(
528
        "-p", "--enlarge-partition", action="store_true",
529
        dest="enlarge_partition", default=False,
530
        help="Enlarge the last partition to cover up the free space")
531

    
532
    options, args = parser.parse_args(sys.argv[1:])
533

    
534
    if len(args) != 1:
535
        parser.error("Wrong number of arguments")
536

    
537
    disklabel = Disklabel(args[0])
538

    
539
    if options.list:
540
        print disklabel
541
        sys.exit(0)
542

    
543
    if options.duid:
544
        print "%s" % "".join(x.encode('hex') for x in disklabel.uid)
545
        sys.exit(0)
546

    
547
    if options.last_part:
548
        print "%c" % chr(ord('a') + disklabel.get_last_partition_id())
549

    
550
    if options.disk_size is not None:
551
        disklabel.enlarge_disk(options.disk_size)
552

    
553
    if options.enlarge_partition:
554
        disklabel.enlarge_last_partition()
555

    
556
    disklabel.write()
557

    
558
sys.exit(0)
559

    
560
# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :