Statistics
| Branch: | Tag: | Revision:

root / snf-image-helper / disklabel.py @ 94d0a699

History | View | Annotate | Download (28.8 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
import abc
30

    
31
from collections import namedtuple
32
from collections import OrderedDict
33

    
34
BLOCKSIZE = 512
35

    
36
LABELSECTOR = 1
37
LABELOFFSET = 0
38

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

    
42
DISKMAGIC = 0x82564557
43

    
44

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

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

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

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

    
76
        def __str__(self):
77
            return "%02Xh %s %02Xh %s %d %d" % (self.status,
78
                                                self.unpack_chs(self.start),
79
                                                self.type,
80
                                                self.unpack_chs(self.end),
81
                                                self.first_sector,
82
                                                self.sector_count)
83

    
84
        @staticmethod
85
        def unpack_chs(chs):
86
            """Unpacks a CHS address string to a tuple."""
87

    
88
            assert len(chs) == 3
89

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

    
95
            return (cylinder, head, sector)
96

    
97
        @staticmethod
98
        def pack_chs(cylinder, head, sector):
99
            """Packs a CHS tuple to an address string."""
100

    
101
            assert 1 <= sector < 2**6, "Invalid sector value"
102
            assert 0 <= head < 2**8, "Invalid head value"
103
            assert 0 <= cylinder < 2**10, "Invalid cylinder value"
104

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

    
109
            return struct.pack('<BBB', byte0, byte1, byte2)
110

    
111
    def __init__(self, block):
112
        """Create an MBR instance"""
113

    
114
        self.fmt = "<444s2x16s16s16s16s2s"
115
        raw_part = {}     # Offset  Length          Contents
116
        (self.code_area,  # 0       440(max. 446)   code area
117
                          # 440     2(optional)     disk signature
118
                          # 444     2               Usually nulls
119
         raw_part[0],     # 446     16              Partition 0
120
         raw_part[1],     # 462     16              Partition 1
121
         raw_part[2],     # 478     16              Partition 2
122
         raw_part[3],     # 494     16              Partition 3
123
         self.signature   # 510     2               MBR signature
124
         ) = struct.unpack(self.fmt, block)
125

    
126
        self.part = {}
127
        for i in range(4):
128
            self.part[i] = self.Partition(raw_part[i])
129

    
130
    def size(self):
131
        """Return the size of a Master Boot Record."""
132
        return struct.calcsize(self.fmt)
133

    
134
    def pack(self):
135
        """Pack an MBR to a binary string."""
136
        return struct.pack(self.fmt,
137
                           self.code_area,
138
                           self.part[0].pack(),
139
                           self.part[1].pack(),
140
                           self.part[2].pack(),
141
                           self.part[3].pack(),
142
                           self.signature)
143

    
144
    def __str__(self):
145
        """Print the MBR"""
146
        ret = ""
147
        for i in range(4):
148
            ret += "Partition %d: %s\n" % (i, self.part[i])
149
        ret += "Signature: %s %s\n" % (hex(ord(self.signature[0])),
150
                                       hex(ord(self.signature[1])))
151
        title = "Master Boot Record"
152
        return "%s\n%s\n%s\n" % (title, len(title) * "=", ret)
153

    
154

    
155
class Disk(object):
156
    """Represents a BSD Disk"""
157

    
158
    def __init__(self, device):
159
        """Create a Disk instance"""
160
        self.device = device
161
        self.part_num = None
162
        self.disklabel = None
163

    
164
        with open(device, "rb") as d:
165
            sector0 = d.read(BLOCKSIZE)
166
            self.mbr = MBR(sector0)
167

    
168
            for i in range(4):
169
                ptype = self.mbr.part[i].type
170
                if ptype in (0xa5, 0xa6, 0xa9):
171
                    d.seek(BLOCKSIZE * self.mbr.part[i].first_sector)
172
                    self.part_num = i
173
                    if ptype == 0xa5:  # FreeBSD
174
                        self.disklabel = BSDDisklabel(d)
175
                    elif ptype == 0xa6:  # OpenBSD
176
                        self.disklabel = OpenBSDDisklabel(d)
177
                    else:  # NetBSD
178
                        self.disklabel = BSDDisklabel(d)
179
                    break
180

    
181
        assert self.disklabel is not None, "No *BSD partition found"
182

    
183
    def write(self):
184
        """Write the changes back to the media"""
185
        with open(self.device, 'rw+b') as d:
186
            d.write(self.mbr.pack())
187

    
188
            d.seek(self.mbr.part[self.part_num].first_sector * BLOCKSIZE)
189
            self.disklabel.write_to(d)
190

    
191
    def __str__(self):
192
        """Print the partitioning info of the Disk"""
193
        return str(self.mbr) + str(self.disklabel)
194

    
195
    def enlarge(self, new_size):
196
        """Enlarge the disk and return the last usable sector"""
197

    
198
        # Fix the disklabel
199
        end = self.disklabel.enlarge(new_size)
200

    
201
        # Fix the MBR
202
        start = self.mbr.part[self.part_num].first_sector
203
        self.mbr.part[self.part_num].sector_count = end - start + 1
204

    
205
        # There are cases where the CHS address changes although the LBA
206
        # address remains the same. For example when the bios is configured
207
        # to use a disk in LBA-Assisted translation mode, the CHS
208
        # representation depends on the size of the disk. When we enlarge the
209
        # size of the disk we may have to update the starting sector's CHS
210
        # address too.
211
        start_chs = MBR.Partition.pack_chs(*self.disklabel.lba2chs(start))
212
        self.mbr.part[self.part_num].start = start_chs
213

    
214
        end_chs = MBR.Partition.pack_chs(*self.disklabel.lba2chs(end))
215
        self.mbr.part[self.part_num].end = end_chs
216

    
217
    def enlarge_last_partition(self):
218
        """Enlarge the last partition to cover up all the free space"""
219
        self.disklabel.enlarge_last_partition()
220

    
221
    def get_last_partition_id(self):
222
        """Get the ID of the last partition"""
223
        return self.disklabel.get_last_partition_id()
224

    
225
    def get_duid(self):
226
        """Get the Disklabel Unique Identifier (works only for OpenBSD)"""
227
        if 'uid' in self.disklabel.field:
228
            return self.disklabel.field['uid']
229

    
230
        return ""
231

    
232

    
233
class DisklabelBase(object):
234
    """Disklabel base class"""
235
    __metaclass__ = abc.ABCMeta
236

    
237
    def __init__(self, device):
238
        """Create a Disklabel instance"""
239

    
240
        # Subclasses need to overwrite this
241
        self.field = None
242
        self.ptable = None
243

    
244
    @abc.abstractproperty
245
    def fmt(self):
246
        """Fields format string for the disklabel fields"""
247
        pass
248

    
249
    def pack(self, checksum=None):
250
        """Return a binary copy of the Disklabel block"""
251

    
252
        out = OrderedDict()
253
        for k, v in self.field.iteritems():
254
            out[k] = v
255

    
256
        if checksum is not None:
257
            out['checksum'] = checksum
258

    
259
        return struct.pack(self.fmt, * out.values() + [self.ptable.pack()])
260

    
261
    def compute_checksum(self):
262
        """Compute the checksum of the disklabel"""
263

    
264
        raw = cStringIO.StringIO(self.pack(0))
265
        checksum = 0
266
        try:
267
            uint16 = raw.read(2)
268
            while uint16 != "":
269
                checksum ^= struct.unpack('<H', uint16)[0]
270
                uint16 = raw.read(2)
271
        finally:
272
            raw.close()
273

    
274
        return checksum
275

    
276
    def write_to(self, device):
277
        """Write the disklabel to a device"""
278

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

    
283
    def lba2chs(self, lba, hpc=None, spt=None):
284
        """Returns the CHS address for a given LBA address"""
285

    
286
        assert hpc is None or hpc > 0, "Invalid heads per cylinder value"
287
        assert spt is None or spt > 0, "Invalid sectors per track value"
288

    
289
        if hpc is None:
290
            hpc = self.field['ntracks']
291

    
292
        if spt is None:
293
            spt = self.field['nsectors']
294

    
295
        # See:
296
        # http://en.wikipedia.org/wiki/Logical_Block_Addressing#CHS_conversion
297
        #
298

    
299
        cylinder = lba // (spt * hpc)
300
        header = (lba // spt) % hpc
301
        sector = (lba % spt) + 1
302

    
303
        return (cylinder, header, sector)
304

    
305
    @abc.abstractmethod
306
    def enlarge(self, new_size):
307
        """Enlarge the disk and return the last usable sector"""
308
        pass
309

    
310
    @abc.abstractmethod
311
    def enlarge_last_partition(self):
312
        """Enlarge the last partition to consume all the usable space"""
313
        pass
314

    
315
    @abc.abstractmethod
316
    def get_last_partition_id(self):
317
        """Get the ID of the last partition"""
318
        pass
319

    
320
    @abc.abstractmethod
321
    def __str__(self):
322
        """Print the Disklabel"""
323
        pass
324

    
325

    
326
class PartitionTableBase(object):
327
    """Base Class for disklabel partition tables"""
328
    __metaclass__ = abc.ABCMeta
329

    
330
    @abc.abstractproperty
331
    def fmt(self):
332
        """Partition fields format string"""
333
        pass
334

    
335
    @abc.abstractproperty
336
    def fields(self):
337
        """The partition fields"""
338
        pass
339

    
340
    def __init__(self, ptable, pnumber):
341
        """Create a Partition Table instance"""
342

    
343
        self.Partition = namedtuple('Partition', self.fields)
344
        self.part = []
345

    
346
        size = struct.calcsize(self.fmt)
347
        raw = cStringIO.StringIO(ptable)
348
        try:
349
            for _ in xrange(pnumber):
350
                self.part.append(
351
                    self.Partition(*struct.unpack(self.fmt, raw.read(size)))
352
                    )
353
        finally:
354
            raw.close()
355

    
356
    def __str__(self):
357
        """Print the Partition table"""
358
        val = ""
359
        for i in xrange(len(self.part)):
360
            val += "%c: %s\n" % (chr(ord('a') + i), str(self.part[i]))
361
        return val
362

    
363
    def pack(self):
364
        """Packs the partition table into a binary string."""
365
        ret = ""
366
        for i in xrange(len(self.part)):
367
            ret += struct.pack(self.fmt, *self.part[i])
368
        return ret + ((364 - len(self.part) * 16) * '\x00')
369

    
370

    
371
class BSDDisklabel(DisklabelBase):
372
    """Represents an BSD Disklabel"""
373

    
374
    class PartitionTable(PartitionTableBase):
375
        """Represents a BSD Partition Table"""
376

    
377
        @property
378
        def fmt(self):
379
            """Partition fields format string"""
380
            return "<IIIBBH"
381

    
382
        @property
383
        def fields(self):
384
            """The partition fields"""
385
            return [    # Offset  Length Contents
386
                'size',     # 0       4      Number of sectors in partition
387
                'offset',   # 4       4      Starting sector of the partition
388
                'fsize',    # 8       4      File system basic fragment size
389
                'fstype',   # 12      1      File system type
390
                'frag',     # 13      1      File system fragments per block
391
                'cpg'       # 14      2      File system cylinders per group
392
                ]
393

    
394
    @property
395
    def fmt(self):
396
        return "<IHH16s16sIIIIIIHHIHHHHIII20s20sIHHII364s"
397

    
398
    def __init__(self, device):
399
        """Create a BSD DiskLabel instance"""
400
        super(BSDDisklabel, self).__init__(device)
401

    
402
        # Disklabel starts at offset one
403
        device.seek(BLOCKSIZE, os.SEEK_CUR)
404
        sector1 = device.read(BLOCKSIZE)
405

    
406
        d_ = OrderedDict()      # Off  Len    Content
407
        (d_["magic"],           # 0    4      Magic
408
         d_["dtype"],           # 4    2      Drive Type
409
         d_["subtype"],         # 6    2      Subtype
410
         d_["typename"],        # 8    16     Type Name
411
         d_["packname"],        # 24   16     Pack Identifier
412
         d_["secsize"],         # 32   4      Bytes per sector
413
         d_["nsectors"],        # 36   4      Data sectors per track
414
         d_["ntracks"],         # 40   4      Tracks per cylinder
415
         d_["ncylinders"],      # 44   4      Data cylinders per unit
416
         d_["secpercyl"],       # 48   4      Data sectors per cylinder
417
         d_["secperunit"],      # 52   4      Data sectors per unit
418
         d_["sparespertrack"],  # 56   2      Spare sectors per track
419
         d_["sparespercyl"],    # 58   2      Spare sectors per cylinder
420
         d_["acylinders"],      # 60   4      Alternative cylinders per unit
421
         d_["rpm"],             # 64   2      Rotation Speed
422
         d_["interleave"],      # 66   2      Hardware sector interleave
423
         d_["trackskew"],       # 68   2      Sector 0 skew, per track
424
         d_["cylskew"],         # 70   2      Sector 0 skew, per cylinder
425
         d_["headswitch"],      # 72   4      Head switch time
426
         d_["trkseek"],         # 76   4      Track-to-track seek
427
         d_["flags"],           # 80   4      Generic Flags
428
         d_["drivedata"],       # 84   5*4    Drive-type specific information
429
         d_["spare"],           # 104  5*4    Reserved for future use
430
         d_["magic2"],          # 124  4      Magic Number
431
         d_["checksum"],        # 128  2      Xor of data including partitions
432
         d_["npartitions"],     # 130  2      Number of partitions following
433
         d_["bbsize"],          # 132  4      size of boot area at sn0, bytes
434
         d_["sbsize"],          # 136  4      Max size of fs superblock, bytes
435
         ptable_raw             # 140  16*16  Partition Table
436
         ) = struct.unpack(self.fmt, sector1)
437

    
438
        assert d_['magic'] == d_['magic2'] == DISKMAGIC, "Disklabel not valid"
439
        self.ptable = self.PartitionTable(ptable_raw, d_['npartitions'])
440
        self.field = d_
441

    
442
    def enlarge(self, new_size):
443
        raise NotImplementedError
444

    
445
    def enlarge_last_partition(self):
446
        raise NotImplementedError
447

    
448
    def get_last_partition_id(self):
449
        raise NotImplementedError
450

    
451
    def lba2chs(self, lba, hpc=None, spt=None):
452
        """Return the CHS address for a given LBA address"""
453

    
454
        # NetBSD uses LBA-Assisted translation. The sectors per track (spt)
455
        # value is always 63 and the heads per cylinder (hpc) value depends on
456
        # the size of the disk
457

    
458
        disk_size = self.field['secperunit']
459

    
460
        # Maximum disk size that can be addressed is 1024 * HPC* SPT
461
        spt = 63
462
        for hpc in 16, 32, 64, 128, 255:
463
            if disk_size <= 1024 * hpc * spt:
464
                break
465

    
466
        chs = super(BSDDisklabel, self).lba2chs(lba, hpc, spt)
467

    
468
        if chs[0] > 1023:  # Cylinders overflowed
469
            assert hpc == 255  # Cylinders may overflow only for large disks
470
            return (1023, chs[1], chs[2])
471

    
472
        return chs
473

    
474
    def __str__(self):
475
        """Print the Disklabel"""
476

    
477
        title = "Disklabel"
478

    
479
        # Those fields may contain null bytes
480
        typename = self.field['typename'].strip('\x00').strip()
481
        packname = self.field['packname'].strip('\x00').strip()
482

    
483
        return \
484
            "%s\n%s\n" % (title, len(title) * "=") + \
485
            "Magic Number: 0x%(magic)x\n" \
486
            "Drive type: %(dtype)d\n" \
487
            "Subtype: %(subtype)d\n" % self.field + \
488
            "Typename: %s\n" % typename + \
489
            "Pack Identifier: %s\n" % packname + \
490
            "# of bytes per sector: %(secsize)d\n" \
491
            "# of data sectors per track: %(nsectors)d\n" \
492
            "# of tracks per cylinder: %(ntracks)d\n" \
493
            "# of data cylinders per unit: %(ncylinders)d\n" \
494
            "# of data sectors per cylinder: %(secpercyl)d\n" \
495
            "# of data sectors per unit: %(secperunit)d\n" \
496
            "# of spare sectors per track: %(sparespertrack)d\n" \
497
            "# of spare sectors per cylinder: %(sparespercyl)d\n" \
498
            "Alt. cylinders per unit: %(acylinders)d\n" \
499
            "Rotational speed: %(rpm)d\n" \
500
            "Hardware sector interleave: %(interleave)d\n" \
501
            "Sector 0 skew, per track: %(trackskew)d\n" \
502
            "Sector 0 skew, per cylinder: %(cylskew)d\n" \
503
            "Head switch time, usec: %(headswitch)d\n" \
504
            "Track-to-track seek, usec: %(trkseek)d\n" \
505
            "Generic Flags: %(flags)r\n" \
506
            "Drive data: %(drivedata)r\n" \
507
            "Reserved for future use: %(spare)r\n" \
508
            "The magic number again: 0x%(magic2)x\n" \
509
            "Checksum: %(checksum)d\n" \
510
            "Number of partitions: %(npartitions)d\n"  \
511
            "Size of boot aread at sn0: %(bbsize)d\n"  \
512
            "Max size of fs superblock: %(sbsize)d\n" % self.field + \
513
            "%s" % self.ptable
514

    
515

    
516
class OpenBSDDisklabel(DisklabelBase):
517
    """Represents an OpenBSD Disklabel"""
518

    
519
    class PartitionTable(PartitionTableBase):
520
        """Reprepsents an OpenBSD Partition Table"""
521

    
522
        @property
523
        def fmt(self):
524
            return "<IIHHBBH"
525

    
526
        @property
527
        def fields(self):
528
            return [        # Offset  Length Contents
529
                'size',     # 0       4      Number of sectors in the partition
530
                'offset',   # 4       4      Starting sector of the partition
531
                'offseth',  # 8       2      Starting sector (high part)
532
                'sizeh',    # 10      2      Number of sectors (high part)
533
                'fstype',   # 12      1      Filesystem type
534
                'frag',     # 13      1      Filesystem Fragments per block
535
                'cpg'       # 14      2      FS cylinders per group
536
                ]
537

    
538
        def setpsize(self, i, size):
539
            """Set size for partition i"""
540
            tmp = self.part[i]
541
            self.part[i] = self.Partition(
542
                size & 0xffffffff, tmp.offset, tmp.offseth, size >> 32,
543
                tmp.fstype, tmp.frag, tmp.cpg)
544

    
545
        def getpsize(self, i):
546
            """Get size for partition i"""
547
            return (self.part[i].sizeh << 32) + self.part[i].size
548

    
549
        def setpoffset(self, i, offset):
550
            """Set offset for partition i"""
551
            tmp = self.part[i]
552
            self.part[i] = self.Partition(
553
                tmp.size, offset & 0xffffffff, offset >> 32, tmp.sizeh,
554
                tmp.frag, tmp.cpg)
555

    
556
        def getpoffset(self, i):
557
            """Get offset for partition i"""
558
            return (self.part[i].offseth << 32) + self.part[i].offset
559

    
560
    @property
561
    def fmt(self):
562
        return "<IHH16s16sIIIIII8sIHHIII20sHH16sIHHII364s"
563

    
564
    def __init__(self, device):
565
        """Create a DiskLabel instance"""
566

    
567
        super(OpenBSDDisklabel, self).__init__(device)
568

    
569
        # Disklabel starts at offset one
570
        device.seek(BLOCKSIZE, os.SEEK_CUR)
571
        sector1 = device.read(BLOCKSIZE)
572

    
573
        d_ = OrderedDict()   # Off  Len    Content
574
        (d_["magic"],        # 0    4      Magic
575
         d_["dtype"],        # 4    2      Drive Type
576
         d_["subtype"],      # 6    2      Subtype
577
         d_["typename"],     # 8    16     Type Name
578
         d_["packname"],     # 24   16     Pack Identifier
579
         d_["secsize"],      # 32   4      Bytes per sector
580
         d_["nsectors"],     # 36   4      Data sectors per track
581
         d_["ntracks"],      # 40   4      Tracks per cylinder
582
         d_["ncylinders"],   # 44   4      Data cylinders per unit
583
         d_["secpercyl"],    # 48   4      Data sectors per cylinder
584
         d_["secperunit"],   # 52   4      Data sectors per unit
585
         d_["uid"],          # 56   8      Unique label identifier
586
         d_["acylinders"],   # 64   4      Alt cylinders per unit
587
         d_["bstarth"],      # 68   2      Start of useable region (high part)
588
         d_["bendh"],        # 70   2      Size of usable region (high part)
589
         d_["bstart"],       # 72   4      Start of useable region
590
         d_["bend"],         # 76   4      End of usable region
591
         d_["flags"],        # 80   4      Generic Flags
592
         d_["drivedata"],    # 84   5*4    Drive-type specific information
593
         d_["secperunith"],  # 104  2      Number of data sectors (high part)
594
         d_["version"],      # 106  2      Version
595
         d_["spare"],        # 108  4*4    Reserved for future use
596
         d_["magic2"],       # 124  4      Magic number
597
         d_["checksum"],     # 128  2      Xor of data including partitions
598
         d_["npartitions"],  # 130  2      Number of partitions in following
599
         d_["bbsize"],       # 132  4      size of boot area at sn0, bytes
600
         d_["sbsize"],       # 136  4      Max size of fs superblock, bytes
601
         ptable_raw          # 140  16*16  Partition Table
602
         ) = struct.unpack(self.fmt, sector1)
603

    
604
        assert d_['magic'] == d_['magic2'] == DISKMAGIC, "Disklabel not valid"
605
        self.ptable = self.PartitionTable(ptable_raw, d_['npartitions'])
606
        self.field = d_
607

    
608
    def setdsize(self, dsize):
609
        """Set disk size"""
610
        self.field['secperunith'] = dsize >> 32
611
        self.field['secperunit'] = dsize & 0xffffffff
612

    
613
    def getdsize(self):
614
        """Get disk size"""
615
        return (self.field['secperunith'] << 32) + self.field['secperunit']
616

    
617
    dsize = property(getdsize, setdsize, None, "disk size")
618

    
619
    def setbstart(self, bstart):
620
        """Set start of useable region"""
621
        self.field['bstarth'] = bstart >> 32
622
        self.field['bstart'] = bstart & 0xffffffff
623

    
624
    def getbstart(self):
625
        """Get start of usable region"""
626
        return (self.field['bstarth'] << 32) + self.field['bstart']
627

    
628
    bstart = property(getbstart, setbstart, None, "start of usable region")
629

    
630
    def setbend(self, bend):
631
        """Set end of useable region"""
632
        self.field['bendh'] = bend >> 32
633
        self.field['bend'] = bend & 0xffffffff
634

    
635
    def getbend(self):
636
        """Get end of usable region"""
637
        return (self.field['bendh'] << 32) + self.field['bend']
638

    
639
    bend = property(getbend, setbend, None, "end of usable region")
640

    
641
    def enlarge(self, new_size):
642
        """Enlarge the disk and return the last usable sector"""
643

    
644
        assert new_size >= self.dsize, \
645
            "New size cannot be smaller that %s" % self.dsize
646

    
647
        # Fix the disklabel
648
        self.dsize = new_size
649
        self.field['ncylinders'] = self.dsize // (self.field['nsectors'] *
650
                                                  self.field['ntracks'])
651
        self.bend = (self.field['ncylinders'] * self.field['nsectors'] *
652
                     self.field['ntracks'])
653

    
654
        # Partition 'c' descriptes the entire disk
655
        self.ptable.setpsize(2, new_size)
656

    
657
        # Update the checksum
658
        self.field['checksum'] = self.compute_checksum()
659

    
660
        # The last usable sector is the end of the usable region minus one
661
        return self.bend - 1
662

    
663
    def get_last_partition_id(self):
664
        """Returns the id of the last partition"""
665
        part = 0
666
        end = 0
667
        # Don't check partition 'c' which is the whole disk
668
        for i in filter(lambda x: x != 2, range(len(self.ptable.part))):
669
            curr_end = self.ptable.getpsize(i) + self.ptable.getpoffset(i)
670
            if end < curr_end:
671
                end = curr_end
672
                part = i
673

    
674
        assert end > 0, "No partition found"
675

    
676
        return part
677

    
678
    def enlarge_last_partition(self):
679
        """Enlarge the last partition to cover up all the free space"""
680

    
681
        part_num = self.get_last_partition_id()
682

    
683
        end = self.ptable.getpsize(part_num) + self.ptable.getpoffset(part_num)
684

    
685
        assert end > 0, "No partition found"
686

    
687
        if self.ptable.part[part_num].fstype == 1:  # Swap partition.
688
            #TODO: Maybe create a warning?
689
            return
690

    
691
        if end > (self.bend - 1024):
692
            return
693

    
694
        self.ptable.setpsize(
695
            part_num, self.bend - self.ptable.getpoffset(part_num) - 1024)
696

    
697
        self.field['checksum'] = self.compute_checksum()
698

    
699
    def lba2chs(self, lba, hpc=None, spt=None):
700
        """Returns the CHS address for a given LBA address"""
701

    
702
        chs = super(OpenBSDDisklabel, self).lba2chs(lba, hpc, spt)
703

    
704
        # If the cylinders overflow, then OpenBSD will put the value
705
        #(1023, 254, 63) to the tuple.
706
        if chs[0] > 1023:
707
            return (1023, 254, 63)
708

    
709
        return chs
710

    
711
    def __str__(self):
712
        """Print the Disklabel"""
713

    
714
        # Those values may contain null bytes
715
        typename = self.field['typename'].strip('\x00').strip()
716
        packname = self.field['packname'].strip('\x00').strip()
717

    
718
        duid = "".join(x.encode('hex') for x in self.field['uid'])
719

    
720
        title = "Disklabel"
721
        return \
722
            "%s\n%s\n" % (title, len(title) * "=") + \
723
            "Magic Number: 0x%(magic)x\n" \
724
            "Drive type: %(dtype)d\n" \
725
            "Subtype: %(subtype)d\n" % self.field + \
726
            "Typename: %s\n" % typename + \
727
            "Pack Identifier: %s\n" % packname + \
728
            "# of bytes per sector: %(secsize)d\n" \
729
            "# of data sectors per track: %(nsectors)d\n" \
730
            "# of tracks per cylinder: %(ntracks)d\n" \
731
            "# of data cylinders per unit: %(ncylinders)d\n" \
732
            "# of data sectors per cylinder: %(secpercyl)d\n" \
733
            "# of data sectors per unit: %(secperunit)d\n" % self.field + \
734
            "DUID: %s\n" % duid + \
735
            "Alt. cylinders per unit: %(acylinders)d\n" \
736
            "Start of useable region (high part): %(bstarth)d\n" \
737
            "Size of useable region (high part): %(bendh)d\n" \
738
            "Start of useable region: %(bstart)d\n" \
739
            "End of usable region: %(bend)d\n" \
740
            "Generic Flags: %(flags)r\n" \
741
            "Drive data: %(drivedata)r\n" \
742
            "Number of data sectors (high part): %(secperunith)d\n" \
743
            "Version: %(version)d\n" \
744
            "Reserved for future use: %(spare)r\n" \
745
            "The magic number again: 0x%(magic2)x\n" \
746
            "Checksum: %(checksum)d\n" \
747
            "Number of partitions: %(npartitions)d\n"  \
748
            "Size of boot aread at sn0: %(bbsize)d\n"  \
749
            "Max size of fs superblock: %(sbsize)d\n" % self.field + \
750
            "%s" % self.ptable
751

    
752

    
753
def main():
754
    """Main entry point"""
755
    usage = "Usage: %prog [options] <input_media>"
756
    parser = optparse.OptionParser(usage=usage)
757

    
758
    parser.add_option("-l", "--list", action="store_true", dest="list",
759
                      default=False,
760
                      help="list the disklabel on the specified media")
761
    parser.add_option("--get-last-partition", action="store_true",
762
                      dest="last_part", default=False,
763
                      help="print the label of the last partition")
764
    parser.add_option(
765
        "--get-duid", action="store_true", dest="duid", default=False,
766
        help="print the Disklabel Unique Identifier (OpenBSD only)")
767
    parser.add_option("-d", "--enlarge-disk", type="int", dest="disk_size",
768
                      default=None, metavar="SIZE",
769
                      help="Enlarge the disk to this SIZE (in sectors)")
770
    parser.add_option(
771
        "-p", "--enlarge-partition", action="store_true",
772
        dest="enlarge_partition", default=False,
773
        help="Enlarge the last partition to cover up the free space")
774

    
775
    options, args = parser.parse_args(sys.argv[1:])
776

    
777
    if len(args) != 1:
778
        parser.error("Wrong number of arguments")
779

    
780
    disk = Disk(args[0])
781

    
782
    if options.list:
783
        print disk
784
        return 0
785

    
786
    if options.duid:
787
        print "%s" % "".join(x.encode('hex') for x in disk.get_duid())
788
        return 0
789

    
790
    if options.last_part:
791
        print "%c" % chr(ord('a') + disk.get_last_partition_id())
792

    
793
    if options.disk_size is not None:
794
        disk.enlarge(options.disk_size)
795

    
796
    if options.enlarge_partition:
797
        disk.enlarge_last_partition()
798

    
799
    disk.write()
800
    return 0
801

    
802

    
803
if __name__ == '__main__':
804
    sys.exit(main())
805

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