Statistics
| Branch: | Tag: | Revision:

root / snf-image-helper / disklabel.py @ be5f0160

History | View | Annotate | Download (22.7 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
DISKMAGIC = 0x82564557
41

    
42

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

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

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

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

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

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

    
85
            assert len(chs) == 3
86

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

    
92
            return (cylinder, head, sector)
93

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

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

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

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

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

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

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

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

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

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

    
165

    
166
class PartitionTableBase(object):
167
    """Base Class for disklabel partition tables"""
168
    format = ""
169

    
170
    Partition = namedtuple('Partition', '')
171

    
172
    def __init__(self, ptable, pnumber):
173
        """Create a Partition Table instance"""
174
        self.part = []
175

    
176
        size = struct.calcsize(self.format)
177

    
178
        raw = cStringIO.StringIO(ptable)
179
        try:
180
            for i in range(pnumber):
181
                p = self.Partition(
182
                    *struct.unpack(self.format, raw.read(size)))
183
                self.part.append(p)
184
        finally:
185
            raw.close()
186

    
187
    def __str__(self):
188
        """Print the Partition table"""
189
        val = ""
190
        for i in range(len(self.part)):
191
            val += "%c: %s\n" % (chr(ord('a') + i), str(self.part[i]))
192
        return val
193

    
194
    def pack(self):
195
        """Packs the partition table into a binary string."""
196
        ret = ""
197
        for i in range(len(self.part)):
198
            ret += struct.pack(self.format, *self.part[i])
199
        return ret
200

    
201

    
202
class Disk(object):
203
    """Represents an BSD Disk"""
204

    
205
    def __init__(self, device):
206
        """Create a Disk instance"""
207
        self.device = device
208
        self.part_num = None
209
        self.disklabel = None
210

    
211
        with open(device, "rb") as d:
212
            sector0 = d.read(BLOCKSIZE)
213
            self.mbr = MBR(sector0)
214

    
215
            for i in range(4):
216
                ptype = self.mbr.part[i].type
217
                if ptype in (0xa5, 0xa6, 0xa9):
218
                    d.seek(BLOCKSIZE * self.mbr.part[i].first_sector)
219
                    self.part_num = i
220
                    if ptype == 0xa5:  # FreeBSD
221
                        self.disklabel = BSD_Disklabel(d)
222
                    elif ptype == 0xa6:  # OpenBSD
223
                        self.disklabel = OpenBSD_Disklabel(d)
224
                    else:  # NetBSD
225
                        self.disklabel = BSD_Disklabel(d)
226
                    break
227

    
228
        assert self.disklabel is not None, "No *BSD partition found"
229

    
230
    def write(self):
231
        """Write the changes back to the media"""
232
        with open(self.device, 'rw+b') as d:
233
            d.write(self.mbr.pack())
234

    
235
            d.seek(self.mbr.part[self.part_num].first_sector * BLOCKSIZE)
236
            self.disklabel.write_to(d)
237

    
238
    def __str__(self):
239
        return str(self.mbr) + str(self.disklabel)
240

    
241
    def enlarge(self, new_size):
242
        """Enlarge the disk and return the last usable sector"""
243

    
244
        # Fix the disklabel
245
        end = self.disklabel.enlarge(new_size)
246

    
247
        # Fix the MBR
248
        start = self.mbr.part[self.part_num].first_sector
249
        self.mbr.part[self.part_num].sector_count = end - start + 1
250

    
251
        cylinder = end // (self.disklabel.ntracks * self.disklabel.nsectors)
252
        header = (end // self.disklabel.nsectors) % self.disklabel.ntracks
253
        sector = (end % self.disklabel.nsectors) + 1
254
        chs = MBR.Partition.pack_chs(cylinder, header, sector)
255
        self.mbr.part[self.part_num].end = chs
256

    
257
    def enlarge_last_partition(self):
258
        self.disklabel.enlarge_last_partition()
259

    
260

    
261
class BSD_Disklabel(object):
262
    """Represents an BSD Disklabel"""
263

    
264
    class PartitionTable(PartitionTableBase):
265
        """Represents a BSD Partition Table"""
266
        format = "<IIIBBH"
267
        """
268
        Partition Entry:
269
        Offset  Length          Contents
270
        0       4               Number of sectors in partition
271
        4       4               Starting sector
272
        8       4               Filesystem basic fragment size
273
        12      1               Filesystem type
274
        13      1               Filesystem fragments per block
275
        14      2               Filesystem cylinders per group
276
        """
277

    
278
        Partition = namedtuple(
279
            'Partition', 'size, offset, fsize, fstype, frag, cpg')
280

    
281
    format = "<IHH16s16sIIIIIIHHIHHHHIII20s20sIHHII64s"
282
    """
283
    Offset  Length          Contents
284
    0       4               Magic
285
    4       2               Drive Type
286
    6       2               Subtype
287
    8       16              Type Name
288
    24      16              Pack Identifier
289
    32      4               Bytes per sector
290
    36      4               Data sectors per track
291
    40      4               Tracks per cylinder
292
    44      4               Data cylinders per unit
293
    48      4               Data sectors per cylinder
294
    52      4               Data sectors per unit
295
    56      2               Spare sectors per track
296
    58      2               Spare sectors per cylinder
297
    60      4               Alternative cylinders per unit
298
    64      2               Rotation Speed
299
    66      2               Hardware sector interleave
300
    68      2               Sector 0 skew, per track
301
    70      2               Sector 0 skew, per cylinder
302
    72      4               Head switch time
303
    76      4               Track-to-track seek
304
    80      4               Generic Flags
305
    84      5*4             Drive-type specific information
306
    104     5*4             Reserved for future use
307
    124     4               Magic Number
308
    128     2               Xor of data including partitions
309
    130     2               Number of partitions following
310
    132     4               size of boot area at sn0, bytes
311
    136     4               Max size of fs superblock, bytes
312
    140     16*16           Partition Table
313
    """
314

    
315

    
316
class OpenBSD_Disklabel(object):
317
    """Represents an OpenBSD Disklabel"""
318

    
319
    class PartitionTable(PartitionTableBase):
320
        """Reprepsents an OpenBSD Partition Table"""
321
        format = "<IIHHBBH"
322
        """
323
        Partition Entry:
324
        Offset  Length          Contents
325
        0       4               Number of sectors in the partition
326
        4       4               Starting sector
327
        8       2               Starting sector (high part)
328
        10      2               Number of sectors (high part)
329
        12      1               Filesystem type
330
        13      1               Filesystem Fragment per block
331
        14      2               FS cylinders per group
332
        """
333

    
334
        Partition = namedtuple(
335
            'Partition', 'size, offset, offseth, sizeh, fstype, frag, cpg')
336

    
337
        def setpsize(self, i, size):
338
            """Set size for partition i"""
339
            tmp = self.part[i]
340
            self.part[i] = self.Partition(size & 0xffffffff, tmp.offset,
341
                                          tmp.offseth, size >> 32, tmp.fstype,
342
                                          tmp.frag, tmp.cpg)
343

    
344
        def getpsize(self, i):
345
            """Get size for partition i"""
346
            return (self.part[i].sizeh << 32) + self.part[i].size
347

    
348
        def setpoffset(self, i, offset):
349
            """Set offset for partition i"""
350
            tmp = self.part[i]
351
            self.part[i] = self.Partition(tmp.size, offset & 0xffffffff,
352
                                          offset >> 32, tmp.sizeh, tmp.frag,
353
                                          tmp.cpg)
354

    
355
        def getpoffset(self, i):
356
            """Get offset for partition i"""
357
            return (self.part[i].offseth << 32) + self.part[i].offset
358

    
359
    format = "<IHH16s16sIIIIII8sIHHIII20sHH16sIHHII364s"
360
    """
361
    Offset  Length          Contents
362
    0       4               Magic
363
    4       2               Drive Type
364
    6       2               Subtype
365
    8       16              Type Name
366
    24      16              Pack Identifier
367
    32      4               Bytes per sector
368
    36      4               Data sectors per track
369
    40      4               Tracks per cylinder
370
    44      4               Data cylinders per unit
371
    48      4               Data sectors per cylinder
372
    52      4               Data sectors per unit
373
    56      8               Unique label identifier
374
    64      4               Alt cylinders per unit
375
    68      2               Start of useable region (high part)
376
    70      2               Size of usable region (high part)
377
    72      4               Start of useable region
378
    76      4               End of usable region
379
    80      4               Generic Flags
380
    84      5*4             Drive-type specific information
381
    104     2               Number of data sectors (high part)
382
    106     2               Version
383
    108     4*4             Reserved for future use
384
    124     4               Magic number
385
    128     2               Xor of data including partitions
386
    130     2               Number of partitions in following
387
    132     4               size of boot area at sn0, bytes
388
    136     4               Max size of fs superblock, bytes
389
    140     16*16           Partition Table
390
    """
391
    def __init__(self, device):
392
        """Create a DiskLabel instance"""
393

    
394
        device.seek(BLOCKSIZE, os.SEEK_CUR)
395
        # The offset of the disklabel from the beginning of the partition is
396
        # one sector
397
        sector1 = device.read(BLOCKSIZE)
398

    
399
        (self.magic,
400
         self.dtype,
401
         self.subtype,
402
         self.typename,
403
         self.packname,
404
         self.secsize,
405
         self.nsectors,
406
         self.ntracks,
407
         self.ncylinders,
408
         self.secpercyl,
409
         self.secperunit,
410
         self.uid,
411
         self.acylinders,
412
         self.bstarth,
413
         self.bendh,
414
         self.bstart,
415
         self.bend,
416
         self.flags,
417
         self.drivedata,
418
         self.secperunith,
419
         self.version,
420
         self.spare,
421
         self.magic2,
422
         self.checksum,
423
         self.npartitions,
424
         self.bbsize,
425
         self.sbsize,
426
         ptable_raw) = struct.unpack(self.format, sector1)
427

    
428
        assert self.magic == DISKMAGIC, "Disklabel is not valid"
429

    
430
        self.ptable = self.PartitionTable(ptable_raw, self.npartitions)
431

    
432
    def pack(self, checksum=None):
433
        return struct.pack(self.format,
434
                           self.magic,
435
                           self.dtype,
436
                           self.subtype,
437
                           self.typename,
438
                           self.packname,
439
                           self.secsize,
440
                           self.nsectors,
441
                           self.ntracks,
442
                           self.ncylinders,
443
                           self.secpercyl,
444
                           self.secperunit,
445
                           self.uid,
446
                           self.acylinders,
447
                           self.bstarth,
448
                           self.bendh,
449
                           self.bstart,
450
                           self.bend,
451
                           self.flags,
452
                           self.drivedata,
453
                           self.secperunith,
454
                           self.version,
455
                           self.spare,
456
                           self.magic2,
457
                           self.checksum if checksum is None else checksum,
458
                           self.npartitions,
459
                           self.bbsize,
460
                           self.sbsize,
461
                           self.ptable.pack() +
462
                           ((364 - self.npartitions * 16) * '\x00'))
463

    
464
    def compute_checksum(self):
465
        """Compute the checksum of the disklabel"""
466

    
467
        raw = cStringIO.StringIO(self.pack(0))
468
        checksum = 0
469
        try:
470
            uint16 = raw.read(2)
471
            while uint16 != "":
472
                checksum ^= struct.unpack('<H', uint16)[0]
473
                uint16 = raw.read(2)
474
        finally:
475
            raw.close()
476

    
477
        return checksum
478

    
479
    def setdsize(self, dsize):
480
        """Set disk size"""
481
        self.secperunith = dsize >> 32
482
        self.secperunit = dsize & 0xffffffff
483

    
484
    def getdsize(self):
485
        """Get disk size"""
486
        return (self.secperunith << 32) + self.secperunit
487

    
488
    def setbstart(self, bstart):
489
        """Set start of useable region"""
490
        self.bstarth = bstart >> 32
491
        self.bstart = bstart & 0xffffffff
492

    
493
    def getbstart(self):
494
        """Get start of usable region"""
495
        return (self.bstarth << 32) + self.bstart
496

    
497
    def setbend(self, bend):
498
        """Set size of useable region"""
499
        self.bendh = bend >> 32
500
        self.bend = bend & 0xffffffff
501

    
502
    def getbend(self):
503
        """Get size of usable region"""
504
        return (self.bendh << 32) + self.bend
505

    
506
    def enlarge(self, new_size):
507
        """Enlarge the size of the disk and return the last usable sector"""
508

    
509
        assert new_size >= self.getdsize(), \
510
            "New size cannot be smaller that %s" % self.getdsize()
511

    
512
        # Fix the disklabel
513
        self.setdsize(new_size)
514
        self.ncylinders = self.getdsize() // (self.nsectors * self.ntracks)
515
        self.setbend(self.ncylinders * self.nsectors * self.ntracks)
516

    
517
        # Partition 'c' descriptes the entire disk
518
        self.ptable.setpsize(2, new_size)
519

    
520
        # Update the checksum
521
        self.checksum = self.compute_checksum()
522

    
523
        # getbend() gives back the size of the usable region and not the end of
524
        # the usable region. I named it like this because this is how it is
525
        # named in OpenBSD. To get the last usable sector you need to reduce
526
        # this value by one.
527
        return self.getbend() - 1
528

    
529
    def write_to(self, device):
530
        """Write the disklabel to a device"""
531

    
532
        # The disklabel starts at sector 1
533
        device.seek(BLOCKSIZE, os.SEEK_CUR)
534
        device.write(self.pack())
535

    
536
    def get_last_partition_id(self):
537
        """Returns the id of the last partition"""
538
        part = 0
539
        end = 0
540
        # Don't check partition 'c' which is the whole disk
541
        for i in filter(lambda x: x != 2, range(self.npartitions)):
542
            curr_end = self.ptable.getpsize(i) + self.ptable.getpoffset(i)
543
            if end < curr_end:
544
                end = curr_end
545
                part = i
546

    
547
        assert end > 0, "No partition found"
548

    
549
        return part
550

    
551
    def enlarge_last_partition(self):
552
        """Enlarge the last partition to cover up all the free space"""
553

    
554
        part_num = self.get_last_partition_id()
555

    
556
        end = self.ptable.getpsize(part_num) + self.ptable.getpoffset(part_num)
557

    
558
        assert end > 0, "No partition found"
559

    
560
        if self.ptable.part[part_num].fstype == 1:  # Swap partition.
561
            #TODO: Maybe create a warning?
562
            return
563

    
564
        if end > (self.getbend() - 1024):
565
            return
566

    
567
        self.ptable.setpsize(
568
            part_num, self.getbend() - self.ptable.getpoffset(part_num) - 1024)
569

    
570
        self.checksum = self.compute_checksum()
571

    
572
    def __str__(self):
573
        """Print the Disklabel"""
574

    
575
        title = "Disklabel"
576
        return \
577
            "%s\n%s\n" % (title, len(title) * "=") + \
578
            "Magic Number: 0x%x\n" % self.magic + \
579
            "Drive type: %d\n" % self.dtype + \
580
            "Subtype: %d\n" % self.subtype + \
581
            "Typename: %s\n" % self.typename.strip('\x00').strip() + \
582
            "Pack Identifier: %s\n" % self.packname.strip('\x00').strip() + \
583
            "Number of bytes per sector: %d\n" % self.secsize + \
584
            "Number of data sectors per track: %d\n" % self.nsectors + \
585
            "Number of tracks per cylinder: %d\n" % self.ntracks + \
586
            "Number of data cylinders per unit: %d\n" % self.ncylinders + \
587
            "Number of data sectors per cylinder: %d\n" % self.secpercyl + \
588
            "Number of data sectors per unit: %d\n" % self.secperunit + \
589
            "DUID: %s\n" % "".join(x.encode('hex') for x in self.uid) + \
590
            "Alt. cylinders per unit: %d\n" % self.acylinders + \
591
            "Start of useable region (high part): %d\n" % self.bstarth + \
592
            "Size of useable region (high part): %d\n" % self.bendh + \
593
            "Start of useable region: %d\n" % self.bstart + \
594
            "End of usable region: %d\n" % self.bend + \
595
            "Generic Flags: %r\n" % self.flags + \
596
            "Drive data: %r\n" % self.drivedata + \
597
            "Number of data sectors (high part): %d\n" % self.secperunith + \
598
            "Version: %d\n" % self.version + \
599
            "Reserved for future use: %r\n" % self.spare + \
600
            "The magic number again: 0x%x\n" % self.magic2 + \
601
            "Checksum: %d\n" % self.checksum + \
602
            "Number of partitions: %d\n" % self.npartitions + \
603
            "Size of boot aread at sn0: %d\n" % self.bbsize + \
604
            "Max size of fs superblock: %d\n" % self.sbsize + \
605
            "%s" % self.ptable
606

    
607

    
608
if __name__ == '__main__':
609

    
610
    usage = "Usage: %prog [options] <input_media>"
611
    parser = optparse.OptionParser(usage=usage)
612

    
613
    parser.add_option("-l", "--list", action="store_true", dest="list",
614
                      default=False,
615
                      help="list the disklabel on the specified media")
616
    parser.add_option("--get-last-partition", action="store_true",
617
                      dest="last_part", default=False,
618
                      help="print the label of the last partition")
619
    parser.add_option("--get-duid", action="store_true", dest="duid",
620
                      default=False,
621
                      help="print the disklabel unique identifier")
622
    parser.add_option("-d", "--enlarge-disk", type="int", dest="disk_size",
623
                      default=None, metavar="SIZE",
624
                      help="Enlarge the disk to this SIZE (in sectors)")
625
    parser.add_option(
626
        "-p", "--enlarge-partition", action="store_true",
627
        dest="enlarge_partition", default=False,
628
        help="Enlarge the last partition to cover up the free space")
629

    
630
    options, args = parser.parse_args(sys.argv[1:])
631

    
632
    if len(args) != 1:
633
        parser.error("Wrong number of arguments")
634

    
635
    disk = Disk(args[0])
636

    
637
    if options.list:
638
        print disk
639
        sys.exit(0)
640

    
641
    if options.duid:
642
        print "%s" % "".join(x.encode('hex') for x in disk.uid)
643
        sys.exit(0)
644

    
645
    if options.last_part:
646
        print "%c" % chr(ord('a') + disk.get_last_partition_id())
647

    
648
    if options.disk_size is not None:
649
        disk.enlarge(options.disk_size)
650

    
651
    if options.enlarge_partition:
652
        disk.enlarge_last_partition()
653

    
654
    disk.write()
655

    
656
sys.exit(0)
657

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