Statistics
| Branch: | Tag: | Revision:

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

History | View | Annotate | Download (23.5 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 DisklabelBase(object):
262
    """Disklabel base class"""
263

    
264
    def __init__(self, device):
265
        """Create a Disklabel instance"""
266
        raise NotImplementedError
267

    
268
    def pack(self, checksum=None):
269
        """Return a binary copy of the Disklabel block"""
270
        raise NotImplementedError
271

    
272
    def compute_checksum(self):
273
        """Compute the checksum of the disklabel"""
274

    
275
        raw = cStringIO.StringIO(self.pack(0))
276
        checksum = 0
277
        try:
278
            uint16 = raw.read(2)
279
            while uint16 != "":
280
                checksum ^= struct.unpack('<H', uint16)[0]
281
                uint16 = raw.read(2)
282
        finally:
283
            raw.close()
284

    
285
        return checksum
286

    
287
    def enlarge(self, new_size):
288
        """Enlarge the disk and return the last usable sector"""
289
        raise NotImplementedError
290

    
291
    def write_to(self, device):
292
        """Write the disklabel to a device"""
293
        raise NotImplementedError
294

    
295
    def enlarge_last_partition(self):
296
        """Enlarge the last partition to consume all the usable space"""
297
        raise NotImplementedError
298

    
299
    def __str__(self):
300
        """Print the Disklabel"""
301
        raise NotImplementedError
302

    
303

    
304
class BSD_Disklabel(DisklabelBase):
305
    """Represents an BSD Disklabel"""
306

    
307
    class PartitionTable(PartitionTableBase):
308
        """Represents a BSD Partition Table"""
309
        format = "<IIIBBH"
310
        """
311
        Partition Entry:
312
        Offset  Length          Contents
313
        0       4               Number of sectors in partition
314
        4       4               Starting sector
315
        8       4               Filesystem basic fragment size
316
        12      1               Filesystem type
317
        13      1               Filesystem fragments per block
318
        14      2               Filesystem cylinders per group
319
        """
320

    
321
        Partition = namedtuple(
322
            'Partition', 'size, offset, fsize, fstype, frag, cpg')
323

    
324
    format = "<IHH16s16sIIIIIIHHIHHHHIII20s20sIHHII64s"
325
    """
326
    Offset  Length          Contents
327
    0       4               Magic
328
    4       2               Drive Type
329
    6       2               Subtype
330
    8       16              Type Name
331
    24      16              Pack Identifier
332
    32      4               Bytes per sector
333
    36      4               Data sectors per track
334
    40      4               Tracks per cylinder
335
    44      4               Data cylinders per unit
336
    48      4               Data sectors per cylinder
337
    52      4               Data sectors per unit
338
    56      2               Spare sectors per track
339
    58      2               Spare sectors per cylinder
340
    60      4               Alternative cylinders per unit
341
    64      2               Rotation Speed
342
    66      2               Hardware sector interleave
343
    68      2               Sector 0 skew, per track
344
    70      2               Sector 0 skew, per cylinder
345
    72      4               Head switch time
346
    76      4               Track-to-track seek
347
    80      4               Generic Flags
348
    84      5*4             Drive-type specific information
349
    104     5*4             Reserved for future use
350
    124     4               Magic Number
351
    128     2               Xor of data including partitions
352
    130     2               Number of partitions following
353
    132     4               size of boot area at sn0, bytes
354
    136     4               Max size of fs superblock, bytes
355
    140     16*16           Partition Table
356
    """
357

    
358

    
359
class OpenBSD_Disklabel(DisklabelBase):
360
    """Represents an OpenBSD Disklabel"""
361

    
362
    class PartitionTable(PartitionTableBase):
363
        """Reprepsents an OpenBSD Partition Table"""
364
        format = "<IIHHBBH"
365
        """
366
        Partition Entry:
367
        Offset  Length          Contents
368
        0       4               Number of sectors in the partition
369
        4       4               Starting sector
370
        8       2               Starting sector (high part)
371
        10      2               Number of sectors (high part)
372
        12      1               Filesystem type
373
        13      1               Filesystem Fragment per block
374
        14      2               FS cylinders per group
375
        """
376

    
377
        Partition = namedtuple(
378
            'Partition', 'size, offset, offseth, sizeh, fstype, frag, cpg')
379

    
380
        def setpsize(self, i, size):
381
            """Set size for partition i"""
382
            tmp = self.part[i]
383
            self.part[i] = self.Partition(size & 0xffffffff, tmp.offset,
384
                                          tmp.offseth, size >> 32, tmp.fstype,
385
                                          tmp.frag, tmp.cpg)
386

    
387
        def getpsize(self, i):
388
            """Get size for partition i"""
389
            return (self.part[i].sizeh << 32) + self.part[i].size
390

    
391
        def setpoffset(self, i, offset):
392
            """Set offset for partition i"""
393
            tmp = self.part[i]
394
            self.part[i] = self.Partition(tmp.size, offset & 0xffffffff,
395
                                          offset >> 32, tmp.sizeh, tmp.frag,
396
                                          tmp.cpg)
397

    
398
        def getpoffset(self, i):
399
            """Get offset for partition i"""
400
            return (self.part[i].offseth << 32) + self.part[i].offset
401

    
402
    format = "<IHH16s16sIIIIII8sIHHIII20sHH16sIHHII364s"
403
    """
404
    Offset  Length          Contents
405
    0       4               Magic
406
    4       2               Drive Type
407
    6       2               Subtype
408
    8       16              Type Name
409
    24      16              Pack Identifier
410
    32      4               Bytes per sector
411
    36      4               Data sectors per track
412
    40      4               Tracks per cylinder
413
    44      4               Data cylinders per unit
414
    48      4               Data sectors per cylinder
415
    52      4               Data sectors per unit
416
    56      8               Unique label identifier
417
    64      4               Alt cylinders per unit
418
    68      2               Start of useable region (high part)
419
    70      2               Size of usable region (high part)
420
    72      4               Start of useable region
421
    76      4               End of usable region
422
    80      4               Generic Flags
423
    84      5*4             Drive-type specific information
424
    104     2               Number of data sectors (high part)
425
    106     2               Version
426
    108     4*4             Reserved for future use
427
    124     4               Magic number
428
    128     2               Xor of data including partitions
429
    130     2               Number of partitions in following
430
    132     4               size of boot area at sn0, bytes
431
    136     4               Max size of fs superblock, bytes
432
    140     16*16           Partition Table
433
    """
434
    def __init__(self, device):
435
        """Create a DiskLabel instance"""
436

    
437
        device.seek(BLOCKSIZE, os.SEEK_CUR)
438
        # The offset of the disklabel from the beginning of the partition is
439
        # one sector
440
        sector1 = device.read(BLOCKSIZE)
441

    
442
        (self.magic,
443
         self.dtype,
444
         self.subtype,
445
         self.typename,
446
         self.packname,
447
         self.secsize,
448
         self.nsectors,
449
         self.ntracks,
450
         self.ncylinders,
451
         self.secpercyl,
452
         self.secperunit,
453
         self.uid,
454
         self.acylinders,
455
         self.bstarth,
456
         self.bendh,
457
         self.bstart,
458
         self.bend,
459
         self.flags,
460
         self.drivedata,
461
         self.secperunith,
462
         self.version,
463
         self.spare,
464
         self.magic2,
465
         self.checksum,
466
         self.npartitions,
467
         self.bbsize,
468
         self.sbsize,
469
         ptable_raw) = struct.unpack(self.format, sector1)
470

    
471
        assert self.magic == DISKMAGIC, "Disklabel is not valid"
472

    
473
        self.ptable = self.PartitionTable(ptable_raw, self.npartitions)
474

    
475
    def pack(self, checksum=None):
476
        return struct.pack(self.format,
477
                           self.magic,
478
                           self.dtype,
479
                           self.subtype,
480
                           self.typename,
481
                           self.packname,
482
                           self.secsize,
483
                           self.nsectors,
484
                           self.ntracks,
485
                           self.ncylinders,
486
                           self.secpercyl,
487
                           self.secperunit,
488
                           self.uid,
489
                           self.acylinders,
490
                           self.bstarth,
491
                           self.bendh,
492
                           self.bstart,
493
                           self.bend,
494
                           self.flags,
495
                           self.drivedata,
496
                           self.secperunith,
497
                           self.version,
498
                           self.spare,
499
                           self.magic2,
500
                           self.checksum if checksum is None else checksum,
501
                           self.npartitions,
502
                           self.bbsize,
503
                           self.sbsize,
504
                           self.ptable.pack() +
505
                           ((364 - self.npartitions * 16) * '\x00'))
506

    
507
    def setdsize(self, dsize):
508
        """Set disk size"""
509
        self.secperunith = dsize >> 32
510
        self.secperunit = dsize & 0xffffffff
511

    
512
    def getdsize(self):
513
        """Get disk size"""
514
        return (self.secperunith << 32) + self.secperunit
515

    
516
    def setbstart(self, bstart):
517
        """Set start of useable region"""
518
        self.bstarth = bstart >> 32
519
        self.bstart = bstart & 0xffffffff
520

    
521
    def getbstart(self):
522
        """Get start of usable region"""
523
        return (self.bstarth << 32) + self.bstart
524

    
525
    def setbend(self, bend):
526
        """Set size of useable region"""
527
        self.bendh = bend >> 32
528
        self.bend = bend & 0xffffffff
529

    
530
    def getbend(self):
531
        """Get size of usable region"""
532
        return (self.bendh << 32) + self.bend
533

    
534
    def enlarge(self, new_size):
535
        """Enlarge the disk and return the last usable sector"""
536

    
537
        assert new_size >= self.getdsize(), \
538
            "New size cannot be smaller that %s" % self.getdsize()
539

    
540
        # Fix the disklabel
541
        self.setdsize(new_size)
542
        self.ncylinders = self.getdsize() // (self.nsectors * self.ntracks)
543
        self.setbend(self.ncylinders * self.nsectors * self.ntracks)
544

    
545
        # Partition 'c' descriptes the entire disk
546
        self.ptable.setpsize(2, new_size)
547

    
548
        # Update the checksum
549
        self.checksum = self.compute_checksum()
550

    
551
        # getbend() gives back the size of the usable region and not the end of
552
        # the usable region. I named it like this because this is how it is
553
        # named in OpenBSD. To get the last usable sector you need to reduce
554
        # this value by one.
555
        return self.getbend() - 1
556

    
557
    def write_to(self, device):
558
        """Write the disklabel to a device"""
559

    
560
        # The disklabel starts at sector 1
561
        device.seek(BLOCKSIZE, os.SEEK_CUR)
562
        device.write(self.pack())
563

    
564
    def get_last_partition_id(self):
565
        """Returns the id of the last partition"""
566
        part = 0
567
        end = 0
568
        # Don't check partition 'c' which is the whole disk
569
        for i in filter(lambda x: x != 2, range(self.npartitions)):
570
            curr_end = self.ptable.getpsize(i) + self.ptable.getpoffset(i)
571
            if end < curr_end:
572
                end = curr_end
573
                part = i
574

    
575
        assert end > 0, "No partition found"
576

    
577
        return part
578

    
579
    def enlarge_last_partition(self):
580
        """Enlarge the last partition to cover up all the free space"""
581

    
582
        part_num = self.get_last_partition_id()
583

    
584
        end = self.ptable.getpsize(part_num) + self.ptable.getpoffset(part_num)
585

    
586
        assert end > 0, "No partition found"
587

    
588
        if self.ptable.part[part_num].fstype == 1:  # Swap partition.
589
            #TODO: Maybe create a warning?
590
            return
591

    
592
        if end > (self.getbend() - 1024):
593
            return
594

    
595
        self.ptable.setpsize(
596
            part_num, self.getbend() - self.ptable.getpoffset(part_num) - 1024)
597

    
598
        self.checksum = self.compute_checksum()
599

    
600
    def __str__(self):
601
        """Print the Disklabel"""
602

    
603
        title = "Disklabel"
604
        return \
605
            "%s\n%s\n" % (title, len(title) * "=") + \
606
            "Magic Number: 0x%x\n" % self.magic + \
607
            "Drive type: %d\n" % self.dtype + \
608
            "Subtype: %d\n" % self.subtype + \
609
            "Typename: %s\n" % self.typename.strip('\x00').strip() + \
610
            "Pack Identifier: %s\n" % self.packname.strip('\x00').strip() + \
611
            "Number of bytes per sector: %d\n" % self.secsize + \
612
            "Number of data sectors per track: %d\n" % self.nsectors + \
613
            "Number of tracks per cylinder: %d\n" % self.ntracks + \
614
            "Number of data cylinders per unit: %d\n" % self.ncylinders + \
615
            "Number of data sectors per cylinder: %d\n" % self.secpercyl + \
616
            "Number of data sectors per unit: %d\n" % self.secperunit + \
617
            "DUID: %s\n" % "".join(x.encode('hex') for x in self.uid) + \
618
            "Alt. cylinders per unit: %d\n" % self.acylinders + \
619
            "Start of useable region (high part): %d\n" % self.bstarth + \
620
            "Size of useable region (high part): %d\n" % self.bendh + \
621
            "Start of useable region: %d\n" % self.bstart + \
622
            "End of usable region: %d\n" % self.bend + \
623
            "Generic Flags: %r\n" % self.flags + \
624
            "Drive data: %r\n" % self.drivedata + \
625
            "Number of data sectors (high part): %d\n" % self.secperunith + \
626
            "Version: %d\n" % self.version + \
627
            "Reserved for future use: %r\n" % self.spare + \
628
            "The magic number again: 0x%x\n" % self.magic2 + \
629
            "Checksum: %d\n" % self.checksum + \
630
            "Number of partitions: %d\n" % self.npartitions + \
631
            "Size of boot aread at sn0: %d\n" % self.bbsize + \
632
            "Max size of fs superblock: %d\n" % self.sbsize + \
633
            "%s" % self.ptable
634

    
635

    
636
if __name__ == '__main__':
637

    
638
    usage = "Usage: %prog [options] <input_media>"
639
    parser = optparse.OptionParser(usage=usage)
640

    
641
    parser.add_option("-l", "--list", action="store_true", dest="list",
642
                      default=False,
643
                      help="list the disklabel on the specified media")
644
    parser.add_option("--get-last-partition", action="store_true",
645
                      dest="last_part", default=False,
646
                      help="print the label of the last partition")
647
    parser.add_option("--get-duid", action="store_true", dest="duid",
648
                      default=False,
649
                      help="print the disklabel unique identifier")
650
    parser.add_option("-d", "--enlarge-disk", type="int", dest="disk_size",
651
                      default=None, metavar="SIZE",
652
                      help="Enlarge the disk to this SIZE (in sectors)")
653
    parser.add_option(
654
        "-p", "--enlarge-partition", action="store_true",
655
        dest="enlarge_partition", default=False,
656
        help="Enlarge the last partition to cover up the free space")
657

    
658
    options, args = parser.parse_args(sys.argv[1:])
659

    
660
    if len(args) != 1:
661
        parser.error("Wrong number of arguments")
662

    
663
    disk = Disk(args[0])
664

    
665
    if options.list:
666
        print disk
667
        sys.exit(0)
668

    
669
    if options.duid:
670
        print "%s" % "".join(x.encode('hex') for x in disk.uid)
671
        sys.exit(0)
672

    
673
    if options.last_part:
674
        print "%c" % chr(ord('a') + disk.get_last_partition_id())
675

    
676
    if options.disk_size is not None:
677
        disk.enlarge(options.disk_size)
678

    
679
    if options.enlarge_partition:
680
        disk.enlarge_last_partition()
681

    
682
    disk.write()
683

    
684
sys.exit(0)
685

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