Statistics
| Branch: | Tag: | Revision:

root / snf-image-helper / disklabel.py @ 0d413fc6

History | View | Annotate | Download (19.7 kB)

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

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

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

    
28
from collections import namedtuple
29

    
30
BLOCKSIZE = 512
31

    
32
LABELSECTOR = 1
33
LABELOFFSET = 0
34

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

    
38
DISKMAGIC = 0x82564557
39

    
40

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

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

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

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

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

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

    
83
            assert len(chs) == 3
84

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

    
90
            return (cylinder, head, sector)
91

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

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

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

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

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

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

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

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

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

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

    
162

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    
372
        return checksum
373

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

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

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

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

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

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

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

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

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

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

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

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

    
425
        self.checksum = self.compute_checksum()
426

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

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

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

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

    
448
        return part
449

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

    
453
        part_num = self.get_last_partition_id()
454

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

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

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

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

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

    
469
        self.checksum = self.compute_checksum()
470

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

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

    
508

    
509
if __name__ == '__main__':
510

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

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

    
534
    options, args = parser.parse_args(sys.argv[1:])
535

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

    
539
    disklabel = Disklabel(args[0])
540

    
541
    if options.list:
542
        print disklabel
543
        sys.exit(0)
544

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

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

    
552
    if options.last_linux:
553
        part_id = disklabel.get_last_partition_id()
554
        # The linux kernel does not assign a partition for label 'c' that
555
        # describes the whole disk
556
        print part_id + (4 if part_id > 2 else 5)
557

    
558
    if options.disk_size is not None:
559
        disklabel.enlarge_disk(options.disk_size)
560

    
561
    if options.enlarge_partition:
562
        disklabel.enlarge_last_partition()
563

    
564
    disklabel.write()
565

    
566
sys.exit(0)
567

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