Statistics
| Branch: | Tag: | Revision:

root / snf-image-helper / disklabel.py @ 501e19ec

History | View | Annotate | Download (22.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
from collections import OrderedDict
32

    
33
BLOCKSIZE = 512
34

    
35
LABELSECTOR = 1
36
LABELOFFSET = 0
37

    
38
BBSIZE = 8192  # size of boot area with label
39
SBSIZE = 8192  # max size of fs superblock
40

    
41
DISKMAGIC = 0x82564557
42

    
43

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

    
50
        def __init__(self, raw_part):
51
            """Create a Partition instance"""
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 Disk(object):
167
    """Represents an BSD Disk"""
168

    
169
    def __init__(self, device):
170
        """Create a Disk instance"""
171
        self.device = device
172
        self.part_num = None
173
        self.disklabel = None
174

    
175
        with open(device, "rb") as d:
176
            sector0 = d.read(BLOCKSIZE)
177
            self.mbr = MBR(sector0)
178

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

    
192
        assert self.disklabel is not None, "No *BSD partition found"
193

    
194
    def write(self):
195
        """Write the changes back to the media"""
196
        with open(self.device, 'rw+b') as d:
197
            d.write(self.mbr.pack())
198

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

    
202
    def __str__(self):
203
        return str(self.mbr) + str(self.disklabel)
204

    
205
    def enlarge(self, new_size):
206
        """Enlarge the disk and return the last usable sector"""
207

    
208
        # Fix the disklabel
209
        end = self.disklabel.enlarge(new_size)
210

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

    
215
        ntracks = self.disklabel.field['ntracks']
216
        nsectors = self.disklabel.field['nsectors']
217

    
218
        cylinder = end // (ntracks * nsectors)
219
        header = (end // nsectors) % ntracks
220
        sector = (end % nsectors) + 1
221
        chs = MBR.Partition.pack_chs(cylinder, header, sector)
222
        self.mbr.part[self.part_num].end = chs
223

    
224
    def enlarge_last_partition(self):
225
        """Enlarge the last partition to cover up all the free space"""
226
        self.disklabel.enlarge_last_partition()
227

    
228
    def get_last_partition_id(self):
229
        """Get the ID of the last partition"""
230
        return self.disklabel.get_last_partition_id()
231

    
232
    def get_duid(self):
233
        return self.disklabel.field['uid']
234

    
235

    
236
class DisklabelBase(object):
237
    """Disklabel base class"""
238

    
239
    def __init__(self, device):
240
        """Create a Disklabel instance"""
241
        raise NotImplementedError
242

    
243
    def pack(self, checksum=None):
244
        """Return a binary copy of the Disklabel block"""
245

    
246
        out = OrderedDict()
247
        for k, v in self.field.items():
248
            out[k] = v
249

    
250
        if checksum is not None:
251
            out['checksum'] = checksum
252

    
253
        return struct.pack(self.format, * out.values() + [self.ptable.pack()])
254

    
255
    def compute_checksum(self):
256
        """Compute the checksum of the disklabel"""
257

    
258
        raw = cStringIO.StringIO(self.pack(0))
259
        checksum = 0
260
        try:
261
            uint16 = raw.read(2)
262
            while uint16 != "":
263
                checksum ^= struct.unpack('<H', uint16)[0]
264
                uint16 = raw.read(2)
265
        finally:
266
            raw.close()
267

    
268
        return checksum
269

    
270
    def enlarge(self, new_size):
271
        """Enlarge the disk and return the last usable sector"""
272
        raise NotImplementedError
273

    
274
    def write_to(self, device):
275
        """Write the disklabel to a device"""
276

    
277
        # The disklabel starts at sector 1
278
        device.seek(BLOCKSIZE, os.SEEK_CUR)
279
        device.write(self.pack())
280

    
281
    def enlarge_last_partition(self):
282
        """Enlarge the last partition to consume all the usable space"""
283
        raise NotImplementedError
284

    
285
    def get_last_partition_id(self):
286
        """Get the ID of the last partition"""
287
        raise NotImplementedError
288

    
289
    def __str__(self):
290
        """Print the Disklabel"""
291
        raise NotImplementedError
292

    
293

    
294
class PartitionTableBase(object):
295
    """Base Class for disklabel partition tables"""
296

    
297
    @property
298
    def format(self):
299
        """Partition table format string"""
300
        raise NotImplementedError
301

    
302
    Partition = namedtuple('Partition', '')
303

    
304
    def __init__(self, ptable, pnumber):
305
        """Create a Partition Table instance"""
306
        self.part = []
307

    
308
        size = struct.calcsize(self.format)
309

    
310
        raw = cStringIO.StringIO(ptable)
311
        try:
312
            for i in range(pnumber):
313
                p = self.Partition(
314
                    *struct.unpack(self.format, raw.read(size)))
315
                self.part.append(p)
316
        finally:
317
            raw.close()
318

    
319
    def __str__(self):
320
        """Print the Partition table"""
321
        val = ""
322
        for i in range(len(self.part)):
323
            val += "%c: %s\n" % (chr(ord('a') + i), str(self.part[i]))
324
        return val
325

    
326
    def pack(self):
327
        """Packs the partition table into a binary string."""
328
        ret = ""
329
        for i in range(len(self.part)):
330
            ret += struct.pack(self.format, *self.part[i])
331
        return ret + ((364 - len(self.part) * 16) * '\x00')
332

    
333

    
334
class BSD_Disklabel(DisklabelBase):
335
    """Represents an BSD Disklabel"""
336

    
337
    class PartitionTable(PartitionTableBase):
338
        """Represents a BSD Partition Table"""
339
        format = "<IIIBBH"
340
        """
341
        Partition Entry:
342
        Offset  Length          Contents
343
        0       4               Number of sectors in partition
344
        4       4               Starting sector
345
        8       4               Filesystem basic fragment size
346
        12      1               Filesystem type
347
        13      1               Filesystem fragments per block
348
        14      2               Filesystem cylinders per group
349
        """
350

    
351
        Partition = namedtuple(
352
            'Partition', 'size, offset, fsize, fstype, frag, cpg')
353

    
354
    format = "<IHH16s16sIIIIIIHHIHHHHIII20s20sIHHII364s"
355
    """
356
    Offset  Length          Contents
357
    0       4               Magic
358
    4       2               Drive Type
359
    6       2               Subtype
360
    8       16              Type Name
361
    24      16              Pack Identifier
362
    32      4               Bytes per sector
363
    36      4               Data sectors per track
364
    40      4               Tracks per cylinder
365
    44      4               Data cylinders per unit
366
    48      4               Data sectors per cylinder
367
    52      4               Data sectors per unit
368
    56      2               Spare sectors per track
369
    58      2               Spare sectors per cylinder
370
    60      4               Alternative cylinders per unit
371
    64      2               Rotation Speed
372
    66      2               Hardware sector interleave
373
    68      2               Sector 0 skew, per track
374
    70      2               Sector 0 skew, per cylinder
375
    72      4               Head switch time
376
    76      4               Track-to-track seek
377
    80      4               Generic Flags
378
    84      5*4             Drive-type specific information
379
    104     5*4             Reserved for future use
380
    124     4               Magic Number
381
    128     2               Xor of data including partitions
382
    130     2               Number of partitions following
383
    132     4               size of boot area at sn0, bytes
384
    136     4               Max size of fs superblock, bytes
385
    140     16*16           Partition Table
386
    """
387

    
388

    
389
class OpenBSD_Disklabel(DisklabelBase):
390
    """Represents an OpenBSD Disklabel"""
391

    
392
    class PartitionTable(PartitionTableBase):
393
        """Reprepsents an OpenBSD Partition Table"""
394
        format = "<IIHHBBH"
395

    
396
        """
397
        Partition Entry:
398
        Offset  Length          Contents
399
        0       4               Number of sectors in the partition
400
        4       4               Starting sector
401
        8       2               Starting sector (high part)
402
        10      2               Number of sectors (high part)
403
        12      1               Filesystem type
404
        13      1               Filesystem Fragment per block
405
        14      2               FS cylinders per group
406
        """
407

    
408
        Partition = namedtuple(
409
            'Partition', 'size, offset, offseth, sizeh, fstype, frag, cpg')
410

    
411
        def setpsize(self, i, size):
412
            """Set size for partition i"""
413
            tmp = self.part[i]
414
            self.part[i] = self.Partition(
415
                size & 0xffffffff, tmp.offset, tmp.offseth, size >> 32,
416
                tmp.fstype, tmp.frag, tmp.cpg)
417

    
418
        def getpsize(self, i):
419
            """Get size for partition i"""
420
            return (self.part[i].sizeh << 32) + self.part[i].size
421

    
422
        def setpoffset(self, i, offset):
423
            """Set offset for partition i"""
424
            tmp = self.part[i]
425
            self.part[i] = self.Partition(
426
                tmp.size, offset & 0xffffffff, offset >> 32, tmp.sizeh,
427
                tmp.frag, tmp.cpg)
428

    
429
        def getpoffset(self, i):
430
            """Get offset for partition i"""
431
            return (self.part[i].offseth << 32) + self.part[i].offset
432

    
433
    format = "<IHH16s16sIIIIII8sIHHIII20sHH16sIHHII364s"
434

    
435
    def __init__(self, device):
436
        """Create a DiskLabel instance"""
437

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

    
443
        d_ = OrderedDict()   # Off  Len    Content
444
        (d_["magic"],        # 0    4      Magic
445
         d_["dtype"],        # 4    2      Drive Type
446
         d_["subtype"],      # 6    2      Subtype
447
         d_["typename"],     # 8    16     Type Name
448
         d_["packname"],     # 24   16     Pack Identifier
449
         d_["secsize"],      # 32   4      Bytes per sector
450
         d_["nsectors"],     # 36   4      Data sectors per track
451
         d_["ntracks"],      # 40   4      Tracks per cylinder
452
         d_["ncylinders"],   # 44   4      Data cylinders per unit
453
         d_["secpercyl"],    # 48   4      Data sectors per cylinder
454
         d_["secperunit"],   # 52   4      Data sectors per unit
455
         d_["uid"],          # 56   8      Unique label identifier
456
         d_["acylinders"],   # 64   4      Alt cylinders per unit
457
         d_["bstarth"],      # 68   2      Start of useable region (high part)
458
         d_["bendh"],        # 70   2      Size of usable region (high part)
459
         d_["bstart"],       # 72   4      Start of useable region
460
         d_["bend"],         # 76   4      End of usable region
461
         d_["flags"],        # 80   4      Generic Flags
462
         d_["drivedata"],    # 84   5*4    Drive-type specific information
463
         d_["secperunith"],  # 104  2      Number of data sectors (high part)
464
         d_["version"],      # 106  2      Version
465
         d_["spare"],        # 108  4*4    Reserved for future use
466
         d_["magic2"],       # 124  4      Magic number
467
         d_["checksum"],     # 128  2      Xor of data including partitions
468
         d_["npartitions"],  # 130  2      Number of partitions in following
469
         d_["bbsize"],       # 132  4      size of boot area at sn0, bytes
470
         d_["sbsize"],       # 136  4      Max size of fs superblock, bytes
471
         ptable_raw          # 140  16*16  Partition Table
472
         ) = struct.unpack(self.format, sector1)
473

    
474
        assert d_['magic'] == d_['magic2'] == DISKMAGIC, "Disklabel not valid"
475
        self.ptable = self.PartitionTable(ptable_raw, d_['npartitions'])
476
        self.field = d_
477

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

    
483
    def getdsize(self):
484
        """Get disk size"""
485
        return (self.field['secperunith'] << 32) + self.field['secperunit']
486

    
487
    dsize = property(getdsize, setdsize, None, "disk size")
488

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

    
494
    def getbstart(self):
495
        """Get start of usable region"""
496
        return (self.field['bstarth'] << 32) + self.field['bstart']
497

    
498
    bstart = property(getbstart, setbstart, None, "usable region start")
499

    
500
    def setbend(self, bend):
501
        """Set size of useable region"""
502
        self.field['bendh'] = bend >> 32
503
        self.field['bend'] = bend & 0xffffffff
504

    
505
    def getbend(self):
506
        """Get size of usable region"""
507
        return (self.field['bendh'] << 32) + self.field['bend']
508

    
509
    bend = property(getbend, setbend, None, "usable region size")
510

    
511
    def enlarge(self, new_size):
512
        """Enlarge the disk and return the last usable sector"""
513

    
514
        assert new_size >= self.dsize, \
515
            "New size cannot be smaller that %s" % self.dsize
516

    
517
        # Fix the disklabel
518
        self.dsize = new_size
519
        self.field['ncylinders'] = self.dsize // (self.field['nsectors'] *
520
                                                  self.field['ntracks'])
521
        self.bend = (self.field['ncylinders'] * self.field['nsectors'] *
522
                     self.field['ntracks'])
523

    
524
        # Partition 'c' descriptes the entire disk
525
        self.ptable.setpsize(2, new_size)
526

    
527
        # Update the checksum
528
        self.checksum = self.compute_checksum()
529

    
530
        # bend is the size and not the end of the usable region. I named it
531
        # like this because this is how it is named in OpenBSD. To get the last
532
        # usable sector you need to reduce this value by one.
533
        return self.bend - 1
534

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

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

    
548
        return part
549

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

    
553
        part_num = self.get_last_partition_id()
554

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

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

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

    
563
        if end > (self.bend - 1024):
564
            return
565

    
566
        self.ptable.setpsize(
567
            part_num, self.bend - self.ptable.getpoffset(part_num) - 1024)
568

    
569
        self.field['checksum'] = self.compute_checksum()
570

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

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

    
609

    
610
if __name__ == '__main__':
611

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

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

    
632
    options, args = parser.parse_args(sys.argv[1:])
633

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

    
637
    disk = Disk(args[0])
638

    
639
    if options.list:
640
        print disk
641
        sys.exit(0)
642

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

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

    
650
    if options.disk_size is not None:
651
        disk.enlarge(options.disk_size)
652

    
653
    if options.enlarge_partition:
654
        disk.enlarge_last_partition()
655

    
656
    disk.write()
657

    
658
sys.exit(0)
659

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