Statistics
| Branch: | Tag: | Revision:

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

History | View | Annotate | Download (28 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 __str__(self):
452
        """Print the Disklabel"""
453

    
454
        title = "Disklabel"
455

    
456
        # Those fields may contain null bytes
457
        typename = self.field['typename'].strip('\x00').strip()
458
        packname = self.field['packname'].strip('\x00').strip()
459

    
460
        return \
461
            "%s\n%s\n" % (title, len(title) * "=") + \
462
            "Magic Number: 0x%(magic)x\n" \
463
            "Drive type: %(dtype)d\n" \
464
            "Subtype: %(subtype)d\n" % self.field + \
465
            "Typename: %s\n" % typename + \
466
            "Pack Identifier: %s\n" % packname + \
467
            "# of bytes per sector: %(secsize)d\n" \
468
            "# of data sectors per track: %(nsectors)d\n" \
469
            "# of tracks per cylinder: %(ntracks)d\n" \
470
            "# of data cylinders per unit: %(ncylinders)d\n" \
471
            "# of data sectors per cylinder: %(secpercyl)d\n" \
472
            "# of data sectors per unit: %(secperunit)d\n" \
473
            "# of spare sectors per track: %(sparespertrack)d\n" \
474
            "# of spare sectors per cylinder: %(sparespercyl)d\n" \
475
            "Alt. cylinders per unit: %(acylinders)d\n" \
476
            "Rotational speed: %(rpm)d\n" \
477
            "Hardware sector interleave: %(interleave)d\n" \
478
            "Sector 0 skew, per track: %(trackskew)d\n" \
479
            "Sector 0 skew, per cylinder: %(cylskew)d\n" \
480
            "Head switch time, usec: %(headswitch)d\n" \
481
            "Track-to-track seek, usec: %(trkseek)d\n" \
482
            "Generic Flags: %(flags)r\n" \
483
            "Drive data: %(drivedata)r\n" \
484
            "Reserved for future use: %(spare)r\n" \
485
            "The magic number again: 0x%(magic2)x\n" \
486
            "Checksum: %(checksum)d\n" \
487
            "Number of partitions: %(npartitions)d\n"  \
488
            "Size of boot aread at sn0: %(bbsize)d\n"  \
489
            "Max size of fs superblock: %(sbsize)d\n" % self.field + \
490
            "%s" % self.ptable
491

    
492

    
493
class OpenBSDDisklabel(DisklabelBase):
494
    """Represents an OpenBSD Disklabel"""
495

    
496
    class PartitionTable(PartitionTableBase):
497
        """Reprepsents an OpenBSD Partition Table"""
498

    
499
        @property
500
        def fmt(self):
501
            return "<IIHHBBH"
502

    
503
        @property
504
        def fields(self):
505
            return [        # Offset  Length Contents
506
                'size',     # 0       4      Number of sectors in the partition
507
                'offset',   # 4       4      Starting sector of the partition
508
                'offseth',  # 8       2      Starting sector (high part)
509
                'sizeh',    # 10      2      Number of sectors (high part)
510
                'fstype',   # 12      1      Filesystem type
511
                'frag',     # 13      1      Filesystem Fragments per block
512
                'cpg'       # 14      2      FS cylinders per group
513
                ]
514

    
515
        def setpsize(self, i, size):
516
            """Set size for partition i"""
517
            tmp = self.part[i]
518
            self.part[i] = self.Partition(
519
                size & 0xffffffff, tmp.offset, tmp.offseth, size >> 32,
520
                tmp.fstype, tmp.frag, tmp.cpg)
521

    
522
        def getpsize(self, i):
523
            """Get size for partition i"""
524
            return (self.part[i].sizeh << 32) + self.part[i].size
525

    
526
        def setpoffset(self, i, offset):
527
            """Set offset for partition i"""
528
            tmp = self.part[i]
529
            self.part[i] = self.Partition(
530
                tmp.size, offset & 0xffffffff, offset >> 32, tmp.sizeh,
531
                tmp.frag, tmp.cpg)
532

    
533
        def getpoffset(self, i):
534
            """Get offset for partition i"""
535
            return (self.part[i].offseth << 32) + self.part[i].offset
536

    
537
    @property
538
    def fmt(self):
539
        return "<IHH16s16sIIIIII8sIHHIII20sHH16sIHHII364s"
540

    
541
    def __init__(self, device):
542
        """Create a DiskLabel instance"""
543

    
544
        super(OpenBSDDisklabel, self).__init__(device)
545

    
546
        # Disklabel starts at offset one
547
        device.seek(BLOCKSIZE, os.SEEK_CUR)
548
        sector1 = device.read(BLOCKSIZE)
549

    
550
        d_ = OrderedDict()   # Off  Len    Content
551
        (d_["magic"],        # 0    4      Magic
552
         d_["dtype"],        # 4    2      Drive Type
553
         d_["subtype"],      # 6    2      Subtype
554
         d_["typename"],     # 8    16     Type Name
555
         d_["packname"],     # 24   16     Pack Identifier
556
         d_["secsize"],      # 32   4      Bytes per sector
557
         d_["nsectors"],     # 36   4      Data sectors per track
558
         d_["ntracks"],      # 40   4      Tracks per cylinder
559
         d_["ncylinders"],   # 44   4      Data cylinders per unit
560
         d_["secpercyl"],    # 48   4      Data sectors per cylinder
561
         d_["secperunit"],   # 52   4      Data sectors per unit
562
         d_["uid"],          # 56   8      Unique label identifier
563
         d_["acylinders"],   # 64   4      Alt cylinders per unit
564
         d_["bstarth"],      # 68   2      Start of useable region (high part)
565
         d_["bendh"],        # 70   2      Size of usable region (high part)
566
         d_["bstart"],       # 72   4      Start of useable region
567
         d_["bend"],         # 76   4      End of usable region
568
         d_["flags"],        # 80   4      Generic Flags
569
         d_["drivedata"],    # 84   5*4    Drive-type specific information
570
         d_["secperunith"],  # 104  2      Number of data sectors (high part)
571
         d_["version"],      # 106  2      Version
572
         d_["spare"],        # 108  4*4    Reserved for future use
573
         d_["magic2"],       # 124  4      Magic number
574
         d_["checksum"],     # 128  2      Xor of data including partitions
575
         d_["npartitions"],  # 130  2      Number of partitions in following
576
         d_["bbsize"],       # 132  4      size of boot area at sn0, bytes
577
         d_["sbsize"],       # 136  4      Max size of fs superblock, bytes
578
         ptable_raw          # 140  16*16  Partition Table
579
         ) = struct.unpack(self.fmt, sector1)
580

    
581
        assert d_['magic'] == d_['magic2'] == DISKMAGIC, "Disklabel not valid"
582
        self.ptable = self.PartitionTable(ptable_raw, d_['npartitions'])
583
        self.field = d_
584

    
585
    def setdsize(self, dsize):
586
        """Set disk size"""
587
        self.field['secperunith'] = dsize >> 32
588
        self.field['secperunit'] = dsize & 0xffffffff
589

    
590
    def getdsize(self):
591
        """Get disk size"""
592
        return (self.field['secperunith'] << 32) + self.field['secperunit']
593

    
594
    dsize = property(getdsize, setdsize, None, "disk size")
595

    
596
    def setbstart(self, bstart):
597
        """Set start of useable region"""
598
        self.field['bstarth'] = bstart >> 32
599
        self.field['bstart'] = bstart & 0xffffffff
600

    
601
    def getbstart(self):
602
        """Get start of usable region"""
603
        return (self.field['bstarth'] << 32) + self.field['bstart']
604

    
605
    bstart = property(getbstart, setbstart, None, "start of usable region")
606

    
607
    def setbend(self, bend):
608
        """Set end of useable region"""
609
        self.field['bendh'] = bend >> 32
610
        self.field['bend'] = bend & 0xffffffff
611

    
612
    def getbend(self):
613
        """Get end of usable region"""
614
        return (self.field['bendh'] << 32) + self.field['bend']
615

    
616
    bend = property(getbend, setbend, None, "end of usable region")
617

    
618
    def enlarge(self, new_size):
619
        """Enlarge the disk and return the last usable sector"""
620

    
621
        assert new_size >= self.dsize, \
622
            "New size cannot be smaller that %s" % self.dsize
623

    
624
        # Fix the disklabel
625
        self.dsize = new_size
626
        self.field['ncylinders'] = self.dsize // (self.field['nsectors'] *
627
                                                  self.field['ntracks'])
628
        self.bend = (self.field['ncylinders'] * self.field['nsectors'] *
629
                     self.field['ntracks'])
630

    
631
        # Partition 'c' descriptes the entire disk
632
        self.ptable.setpsize(2, new_size)
633

    
634
        # Update the checksum
635
        self.field['checksum'] = self.compute_checksum()
636

    
637
        # The last usable sector is the end of the usable region minus one
638
        return self.bend - 1
639

    
640
    def get_last_partition_id(self):
641
        """Returns the id of the last partition"""
642
        part = 0
643
        end = 0
644
        # Don't check partition 'c' which is the whole disk
645
        for i in filter(lambda x: x != 2, range(len(self.ptable.part))):
646
            curr_end = self.ptable.getpsize(i) + self.ptable.getpoffset(i)
647
            if end < curr_end:
648
                end = curr_end
649
                part = i
650

    
651
        assert end > 0, "No partition found"
652

    
653
        return part
654

    
655
    def enlarge_last_partition(self):
656
        """Enlarge the last partition to cover up all the free space"""
657

    
658
        part_num = self.get_last_partition_id()
659

    
660
        end = self.ptable.getpsize(part_num) + self.ptable.getpoffset(part_num)
661

    
662
        assert end > 0, "No partition found"
663

    
664
        if self.ptable.part[part_num].fstype == 1:  # Swap partition.
665
            #TODO: Maybe create a warning?
666
            return
667

    
668
        if end > (self.bend - 1024):
669
            return
670

    
671
        self.ptable.setpsize(
672
            part_num, self.bend - self.ptable.getpoffset(part_num) - 1024)
673

    
674
        self.field['checksum'] = self.compute_checksum()
675

    
676
    def lba2chs(self, lba, hpc=None, spt=None):
677
        """Returns the CHS address for a given LBA address"""
678

    
679
        chs = super(OpenBSDDisklabel, self).lba2chs(lba, hpc, spt)
680

    
681
        # If the cylinders overflow, then OpenBSD will put the value
682
        #(1023, 254, 63) to the tuple.
683
        if chs[0] > 1023:
684
            return (1023, 254, 63)
685

    
686
        return chs
687

    
688
    def __str__(self):
689
        """Print the Disklabel"""
690

    
691
        # Those values may contain null bytes
692
        typename = self.field['typename'].strip('\x00').strip()
693
        packname = self.field['packname'].strip('\x00').strip()
694

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

    
697
        title = "Disklabel"
698
        return \
699
            "%s\n%s\n" % (title, len(title) * "=") + \
700
            "Magic Number: 0x%(magic)x\n" \
701
            "Drive type: %(dtype)d\n" \
702
            "Subtype: %(subtype)d\n" % self.field + \
703
            "Typename: %s\n" % typename + \
704
            "Pack Identifier: %s\n" % packname + \
705
            "# of bytes per sector: %(secsize)d\n" \
706
            "# of data sectors per track: %(nsectors)d\n" \
707
            "# of tracks per cylinder: %(ntracks)d\n" \
708
            "# of data cylinders per unit: %(ncylinders)d\n" \
709
            "# of data sectors per cylinder: %(secpercyl)d\n" \
710
            "# of data sectors per unit: %(secperunit)d\n" % self.field + \
711
            "DUID: %s\n" % duid + \
712
            "Alt. cylinders per unit: %(acylinders)d\n" \
713
            "Start of useable region (high part): %(bstarth)d\n" \
714
            "Size of useable region (high part): %(bendh)d\n" \
715
            "Start of useable region: %(bstart)d\n" \
716
            "End of usable region: %(bend)d\n" \
717
            "Generic Flags: %(flags)r\n" \
718
            "Drive data: %(drivedata)r\n" \
719
            "Number of data sectors (high part): %(secperunith)d\n" \
720
            "Version: %(version)d\n" \
721
            "Reserved for future use: %(spare)r\n" \
722
            "The magic number again: 0x%(magic2)x\n" \
723
            "Checksum: %(checksum)d\n" \
724
            "Number of partitions: %(npartitions)d\n"  \
725
            "Size of boot aread at sn0: %(bbsize)d\n"  \
726
            "Max size of fs superblock: %(sbsize)d\n" % self.field + \
727
            "%s" % self.ptable
728

    
729

    
730
def main():
731
    """Main entry point"""
732
    usage = "Usage: %prog [options] <input_media>"
733
    parser = optparse.OptionParser(usage=usage)
734

    
735
    parser.add_option("-l", "--list", action="store_true", dest="list",
736
                      default=False,
737
                      help="list the disklabel on the specified media")
738
    parser.add_option("--get-last-partition", action="store_true",
739
                      dest="last_part", default=False,
740
                      help="print the label of the last partition")
741
    parser.add_option(
742
        "--get-duid", action="store_true", dest="duid", default=False,
743
        help="print the Disklabel Unique Identifier (OpenBSD only)")
744
    parser.add_option("-d", "--enlarge-disk", type="int", dest="disk_size",
745
                      default=None, metavar="SIZE",
746
                      help="Enlarge the disk to this SIZE (in sectors)")
747
    parser.add_option(
748
        "-p", "--enlarge-partition", action="store_true",
749
        dest="enlarge_partition", default=False,
750
        help="Enlarge the last partition to cover up the free space")
751

    
752
    options, args = parser.parse_args(sys.argv[1:])
753

    
754
    if len(args) != 1:
755
        parser.error("Wrong number of arguments")
756

    
757
    disk = Disk(args[0])
758

    
759
    if options.list:
760
        print disk
761
        return 0
762

    
763
    if options.duid:
764
        print "%s" % "".join(x.encode('hex') for x in disk.get_duid())
765
        return 0
766

    
767
    if options.last_part:
768
        print "%c" % chr(ord('a') + disk.get_last_partition_id())
769

    
770
    if options.disk_size is not None:
771
        disk.enlarge(options.disk_size)
772

    
773
    if options.enlarge_partition:
774
        disk.enlarge_last_partition()
775

    
776
    disk.write()
777
    return 0
778

    
779

    
780
if __name__ == '__main__':
781
    sys.exit(main())
782

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