Statistics
| Branch: | Tag: | Revision:

root / snf-image-helper / disklabel.py @ 34fd52a7

History | View | Annotate | Download (20.9 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 BSD disklabels"""
23

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

    
30
from collections import namedtuple
31

    
32
BLOCKSIZE = 512
33

    
34
LABELSECTOR = 1
35
LABELOFFSET = 0
36

    
37
BBSIZE = 8192  # size of boot area with label
38
SBSIZE = 8192  # max size of fs superblock
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
        title = "Master Boot Record"
161
        return "%s\n%s\n%s\n" % (title, len(title) * "=", ret)
162

    
163

    
164
class Disklabel:
165
    """Represents an BSD Disklabel"""
166

    
167
    def __init__(self, disk):
168
        """Create a DiskLabel instance"""
169
        self.disk = disk
170
        self.part_num = None
171
        self.disklabel = None
172

    
173
        with open(disk, "rb") as d:
174
            sector0 = d.read(BLOCKSIZE)
175
            self.mbr = MBR(sector0)
176

    
177
            for i in range(4):
178
                ptype = self.mbr.part[i].type
179
                if ptype in (0xa5, 0xa6, 0xa9):
180
                    d.seek(BLOCKSIZE * self.mbr.part[i].first_sector)
181
                    self.part_num = i
182
                    if ptype == 0xa5:  # FreeBSD
183
                        self.disklabel = BSD_Disklabel(d)
184
                    elif ptype == 0xa6:  # OpenBSD
185
                        self.disklabel = OpenBSD_Disklabel(d)
186
                    else:  # NetBSD
187
                        self.disklabel = BSD_Disklabel(d)
188
                    break
189

    
190
        assert self.disklabel is not None, "No *BSD partition found"
191

    
192
    def write(self):
193
        """Write the disklabel back to the media"""
194
        with open(self.disk, 'rw+b') as d:
195
            d.write(self.mbr.pack())
196

    
197
            d.seek(self.mbr.part[self.part_num].first_sector * BLOCKSIZE)
198
            self.disklabel.write_to(d)
199

    
200
    def __str__(self):
201
        return str(self.mbr) + str(self.disklabel)
202

    
203
    def enlarge_disk(self, new_size):
204
        """Enlarge the size of the disk and return the last usable sector"""
205

    
206
        # Fix the disklabel
207
        end = self.disklabel.enlarge_disk(new_size)
208

    
209
        # Fix the MBR
210
        start = self.mbr.part[self.part_num].first_sector
211
        self.mbr.part[self.part_num].sector_count = end - start + 1
212

    
213
        cylinder = end // (self.disklabel.ntracks * self.disklabel.nsectors)
214
        header = (end // self.disklabel.nsectors) % self.disklabel.ntracks
215
        sector = (end % self.disklabel.nsectors) + 1
216
        chs = MBR.Partition.pack_chs(cylinder, header, sector)
217
        self.mbr.part[self.part_num].end = chs
218

    
219
    def enlarge_last_partition(self):
220
        self.disklabel.enlarge_last_partition()
221

    
222

    
223
class BSD_Disklabel:
224
    pass
225

    
226

    
227
class OpenBSD_Disklabel:
228
    """Represents an OpenBSD Disklabel"""
229
    format = "<IHH16s16sIIIIII8sIHHIII20sHH16sIHHII364s"
230
    """
231
    Offset  Length          Contents
232
    0       4               Magic
233
    4       2               Drive Type
234
    6       2               Subtype
235
    8       16              Type Name
236
    24      16              Pack Identifier
237
    32      4               Bytes per sector
238
    36      4               Data sectors per track
239
    40      4               Tracks per cilinder
240
    44      4               Data cylinders per unit
241
    48      4               Data sectors per cylynder
242
    52      4               Data sectors per unit
243
    56      8               Unique label identifier
244
    64      4               Alt cylinders per unit
245
    68      2               Start of useable region (high part)
246
    70      2               Size of usable region (high part)
247
    72      4               Start of useable region
248
    76      4               End of usable region
249
    80      4               Generic Flags
250
    84      5*4             Drive-type specific information
251
    104     2               Number of data sectors (high part)
252
    106     2               Version
253
    108     4*4             Reserved for future use
254
    124     4               Magic number
255
    128     2               Xor of data Inclu. partitions
256
    130     2               Number of partitions in following
257
    132     4               size of boot area at sn0, bytes
258
    136     4               Max size of fs superblock, bytes
259
    140     16*16           Partition Table
260
    """
261

    
262
    class PartitionTable:
263
        """Reprepsents an OpenBSD Partition Table"""
264
        format = "<IIHHBBH"
265
        """
266
        Partition Entry:
267
        Offset  Length          Contents
268
        0       4               Number of sectors in the partition
269
        4       4               Starting sector
270
        8       2               Starting sector (high part)
271
        10      2               Number of sectors (high part)
272
        12      1               Filesystem type
273
        13      1               Filesystem Fragment per block
274
        14      2               FS cylinders per group
275
        """
276

    
277
        Partition = namedtuple(
278
            'Partition', 'size, offset, offseth, sizeh, fstype, frag, cpg')
279

    
280
        def __init__(self, ptable, pnumber):
281
            """Create a Partition Table instance"""
282
            self.part = []
283

    
284
            size = struct.calcsize(self.format)
285

    
286
            raw = cStringIO.StringIO(ptable)
287
            try:
288
                for i in range(pnumber):
289
                    p = self.Partition(
290
                        *struct.unpack(self.format, raw.read(size)))
291
                    self.part.append(p)
292
            finally:
293
                raw.close()
294

    
295
        def __str__(self):
296
            """Print the Partition table"""
297
            val = ""
298
            for i in range(len(self.part)):
299
                val += "%c: %s\n" % (chr(ord('a') + i), str(self.part[i]))
300
            return val
301

    
302
        def pack(self):
303
            """Packs the partition table into a binary string."""
304
            ret = ""
305
            for i in range(len(self.part)):
306
                ret += struct.pack(self.format,
307
                                   self.part[i].size,
308
                                   self.part[i].offset,
309
                                   self.part[i].offseth,
310
                                   self.part[i].sizeh,
311
                                   self.part[i].fstype,
312
                                   self.part[i].frag,
313
                                   self.part[i].cpg)
314
            return ret
315

    
316
        def setpsize(self, i, size):
317
            """Set size for partition i"""
318
            tmp = self.part[i]
319
            self.part[i] = self.Partition(size & 0xffffffff, tmp.offset,
320
                                          tmp.offseth, size >> 32, tmp.fstype,
321
                                          tmp.frag, tmp.cpg)
322

    
323
        def getpsize(self, i):
324
            """Get size for partition i"""
325
            return (self.part[i].sizeh << 32) + self.part[i].size
326

    
327
        def setpoffset(self, i, offset):
328
            """Set offset for partition i"""
329
            tmp = self.part[i]
330
            self.part[i] = self.Partition(tmp.size, offset & 0xffffffff,
331
                                          offset >> 32, tmp.sizeh, tmp.frag,
332
                                          tmp.cpg)
333

    
334
        def getpoffset(self, i):
335
            """Get offset for partition i"""
336
            return (self.part[i].offseth << 32) + self.part[i].offset
337

    
338
    DISKMAGIC = 0x82564557
339

    
340
    def __init__(self, device):
341
        """Create a DiskLabel instance"""
342

    
343
        device.seek(BLOCKSIZE, os.SEEK_CUR)
344
        # The offset of the disklabel from the beginning of the partition is
345
        # one sector
346
        sector1 = device.read(BLOCKSIZE)
347

    
348
        (self.magic,
349
         self.dtype,
350
         self.subtype,
351
         self.typename,
352
         self.packname,
353
         self.secsize,
354
         self.nsectors,
355
         self.ntracks,
356
         self.ncylinders,
357
         self.secpercyl,
358
         self.secperunit,
359
         self.uid,
360
         self.acylinders,
361
         self.bstarth,
362
         self.bendh,
363
         self.bstart,
364
         self.bend,
365
         self.flags,
366
         self.drivedata,
367
         self.secperunith,
368
         self.version,
369
         self.spare,
370
         self.magic2,
371
         self.checksum,
372
         self.npartitions,
373
         self.bbsize,
374
         self.sbsize,
375
         ptable_raw) = struct.unpack(self.format, sector1)
376

    
377
        assert self.magic == self.DISKMAGIC, "Disklabel is not valid"
378

    
379
        self.ptable = self.PartitionTable(ptable_raw, self.npartitions)
380

    
381
    def pack(self, checksum=None):
382
        return struct.pack(self.format,
383
                           self.magic,
384
                           self.dtype,
385
                           self.subtype,
386
                           self.typename,
387
                           self.packname,
388
                           self.secsize,
389
                           self.nsectors,
390
                           self.ntracks,
391
                           self.ncylinders,
392
                           self.secpercyl,
393
                           self.secperunit,
394
                           self.uid,
395
                           self.acylinders,
396
                           self.bstarth,
397
                           self.bendh,
398
                           self.bstart,
399
                           self.bend,
400
                           self.flags,
401
                           self.drivedata,
402
                           self.secperunith,
403
                           self.version,
404
                           self.spare,
405
                           self.magic2,
406
                           self.checksum if checksum is None else checksum,
407
                           self.npartitions,
408
                           self.bbsize,
409
                           self.sbsize,
410
                           self.ptable.pack() +
411
                           ((364 - self.npartitions * 16) * '\x00'))
412

    
413
    def compute_checksum(self):
414
        """Compute the checksum of the disklabel"""
415

    
416
        raw = cStringIO.StringIO(self.pack(0))
417
        checksum = 0
418
        try:
419
            uint16 = raw.read(2)
420
            while uint16 != "":
421
                checksum ^= struct.unpack('<H', uint16)[0]
422
                uint16 = raw.read(2)
423
        finally:
424
            raw.close()
425

    
426
        return checksum
427

    
428
    def setdsize(self, dsize):
429
        """Set disk size"""
430
        self.secperunith = dsize >> 32
431
        self.secperunit = dsize & 0xffffffff
432

    
433
    def getdsize(self):
434
        """Get disk size"""
435
        return (self.secperunith << 32) + self.secperunit
436

    
437
    def setbstart(self, bstart):
438
        """Set start of useable region"""
439
        self.bstarth = bstart >> 32
440
        self.bstart = bstart & 0xffffffff
441

    
442
    def getbstart(self):
443
        """Get start of usable region"""
444
        return (self.bstarth << 32) + self.bstart
445

    
446
    def setbend(self, bend):
447
        """Set size of useable region"""
448
        self.bendh = bend >> 32
449
        self.bend = bend & 0xffffffff
450

    
451
    def getbend(self):
452
        """Get size of usable region"""
453
        return (self.bendh << 32) + self.bend
454

    
455
    def enlarge_disk(self, new_size):
456
        """Enlarge the size of the disk and return the last usable sector"""
457

    
458
        assert new_size >= self.getdsize(), \
459
            "New size cannot be smaller that %s" % self.getdsize()
460

    
461
        # Fix the disklabel
462
        self.setdsize(new_size)
463
        self.ncylinders = self.getdsize() // (self.nsectors * self.ntracks)
464
        self.setbend(self.ncylinders * self.nsectors * self.ntracks)
465

    
466
        # Partition 'c' descriptes the entire disk
467
        self.ptable.setpsize(2, new_size)
468

    
469
        # Update the checksum
470
        self.checksum = self.compute_checksum()
471

    
472
        # getbend() gives back the size of the usable region and not the end of
473
        # the usable region. I named it like this because this is how it is
474
        # named in OpenBSD. To get the last usable sector you need to reduce
475
        # this value by one.
476
        return self.getbend() - 1
477

    
478
    def write_to(self, device):
479
        """Write the disklabel to a device"""
480

    
481
        # The disklabel starts at sector 1
482
        device.seek(BLOCKSIZE, os.SEEK_CUR)
483
        device.write(self.pack())
484

    
485
    def get_last_partition_id(self):
486
        """Returns the id of the last partition"""
487
        part = 0
488
        end = 0
489
        # Don't check partition 'c' which is the whole disk
490
        for i in filter(lambda x: x != 2, range(self.npartitions)):
491
            curr_end = self.ptable.getpsize(i) + self.ptable.getpoffset(i)
492
            if end < curr_end:
493
                end = curr_end
494
                part = i
495

    
496
        assert end > 0, "No partition found"
497

    
498
        return part
499

    
500
    def enlarge_last_partition(self):
501
        """Enlarge the last partition to cover up all the free space"""
502

    
503
        part_num = self.get_last_partition_id()
504

    
505
        end = self.ptable.getpsize(part_num) + self.ptable.getpoffset(part_num)
506

    
507
        assert end > 0, "No partition found"
508

    
509
        if self.ptable.part[part_num].fstype == 1:  # Swap partition.
510
            #TODO: Maybe create a warning?
511
            return
512

    
513
        if end > (self.getbend() - 1024):
514
            return
515

    
516
        self.ptable.setpsize(
517
            part_num, self.getbend() - self.ptable.getpoffset(part_num) - 1024)
518

    
519
        self.checksum = self.compute_checksum()
520

    
521
    def __str__(self):
522
        """Print the Disklabel"""
523

    
524
        title = "Disklabel"
525
        return \
526
            "%s\n%s\n" % (title, len(title) * "=") + \
527
            "Magic Number: 0x%x\n" % self.magic + \
528
            "Drive type: %d\n" % self.dtype + \
529
            "Subtype: %d\n" % self.subtype + \
530
            "Typename: %s\n" % self.typename.strip('\x00').strip() + \
531
            "Pack Identifier: %s\n" % self.packname.strip('\x00').strip() + \
532
            "Number of bytes per sector: %d\n" % self.secsize + \
533
            "Number of data sectors per track: %d\n" % self.nsectors + \
534
            "Number of tracks per cylinder: %d\n" % self.ntracks + \
535
            "Number of data cylinders per unit: %d\n" % self.ncylinders + \
536
            "Number of data sectors per cylinder: %d\n" % self.secpercyl + \
537
            "Number of data sectors per unit: %d\n" % self.secperunit + \
538
            "DUID: %s\n" % "".join(x.encode('hex') for x in self.uid) + \
539
            "Alt. cylinders per unit: %d\n" % self.acylinders + \
540
            "Start of useable region (high part): %d\n" % self.bstarth + \
541
            "Size of useable region (high part): %d\n" % self.bendh + \
542
            "Start of useable region: %d\n" % self.bstart + \
543
            "End of usable region: %d\n" % self.bend + \
544
            "Generic Flags: %r\n" % self.flags + \
545
            "Drive data: %r\n" % self.drivedata + \
546
            "Number of data sectors (high part): %d\n" % self.secperunith + \
547
            "Version: %d\n" % self.version + \
548
            "Reserved for future use: %r\n" % self.spare + \
549
            "The magic number again: 0x%x\n" % self.magic2 + \
550
            "Checksum: %d\n" % self.checksum + \
551
            "Number of partitions: %d\n" % self.npartitions + \
552
            "Size of boot aread at sn0: %d\n" % self.bbsize + \
553
            "Max size of fs superblock: %d\n" % self.sbsize + \
554
            "%s" % self.ptable
555

    
556

    
557
if __name__ == '__main__':
558

    
559
    usage = "Usage: %prog [options] <input_media>"
560
    parser = optparse.OptionParser(usage=usage)
561

    
562
    parser.add_option("-l", "--list", action="store_true", dest="list",
563
                      default=False,
564
                      help="list the disklabel on the specified media")
565
    parser.add_option("--get-last-partition", action="store_true",
566
                      dest="last_part", default=False,
567
                      help="print the label of the last partition")
568
    parser.add_option("--get-duid", action="store_true", dest="duid",
569
                      default=False,
570
                      help="print the disklabel unique identifier")
571
    parser.add_option("-d", "--enlarge-disk", type="int", dest="disk_size",
572
                      default=None, metavar="SIZE",
573
                      help="Enlarge the disk to this SIZE (in sectors)")
574
    parser.add_option(
575
        "-p", "--enlarge-partition", action="store_true",
576
        dest="enlarge_partition", default=False,
577
        help="Enlarge the last partition to cover up the free space")
578

    
579
    options, args = parser.parse_args(sys.argv[1:])
580

    
581
    if len(args) != 1:
582
        parser.error("Wrong number of arguments")
583

    
584
    disklabel = Disklabel(args[0])
585

    
586
    if options.list:
587
        print disklabel
588
        sys.exit(0)
589

    
590
    if options.duid:
591
        print "%s" % "".join(x.encode('hex') for x in disklabel.uid)
592
        sys.exit(0)
593

    
594
    if options.last_part:
595
        print "%c" % chr(ord('a') + disklabel.get_last_partition_id())
596

    
597
    if options.disk_size is not None:
598
        disklabel.enlarge_disk(options.disk_size)
599

    
600
    if options.enlarge_partition:
601
        disklabel.enlarge_last_partition()
602

    
603
    disklabel.write()
604

    
605
sys.exit(0)
606

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