disklabel.py 27.8 KB
Newer Older
1
#!/usr/bin/env python
2
#
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# -*- coding: utf-8 -*-
#
# Copyright (C) 2013 GRNET S.A.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
# 02110-1301, USA.

22
"""This module provides the code for handling BSD disklabels"""
23
24
25

import struct
import sys
26
import os
27
28
import cStringIO
import optparse
29
import abc
30
31

from collections import namedtuple
32
from collections import OrderedDict
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

BLOCKSIZE = 512

LABELSECTOR = 1
LABELOFFSET = 0

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

DISKMAGIC = 0x82564557


class MBR(object):
    """Represents a Master Boot Record."""
    class Partition(object):
        """Represents a partition entry in MBR"""
49
        fmt = "<B3sB3sLL"
50
51
52

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

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

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

        def __str__(self):
77
78
79
80
81
82
            return "%02Xh %s %02Xh %s %d %d" % (self.status,
                                                self.unpack_chs(self.start),
                                                self.type,
                                                self.unpack_chs(self.end),
                                                self.first_sector,
                                                self.sector_count)
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

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

            assert len(chs) == 3

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

            return (cylinder, head, sector)

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

            assert 1 <= sector <= 63
            assert 0 <= head <= 255
103
104
105
106
107
108
109
110
            assert 0 <= cylinder

            # If the cylinders overflow then put the value (1023, 254, 63) to
            # the tuple. At least this is what OpenBSD does.
            if cylinder > 1023:
                cylinder = 1023
                head = 254
                sector = 63
111
112
113
114
115
116
117
118
119

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

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

    def __init__(self, block):
        """Create an MBR instance"""
120

121
        self.fmt = "<444s2x16s16s16s16s2s"
122
123
124
125
126
127
128
129
130
        raw_part = {}     # Offset  Length          Contents
        (self.code_area,  # 0       440(max. 446)   code area
                          # 440     2(optional)     disk signature
                          # 444     2               Usually nulls
         raw_part[0],     # 446     16              Partition 0
         raw_part[1],     # 462     16              Partition 1
         raw_part[2],     # 478     16              Partition 2
         raw_part[3],     # 494     16              Partition 3
         self.signature   # 510     2               MBR signature
131
         ) = struct.unpack(self.fmt, block)
132
133
134
135
136

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

137
    def size(self):
138
        """Return the size of a Master Boot Record."""
139
        return struct.calcsize(self.fmt)
140
141
142

    def pack(self):
        """Pack an MBR to a binary string."""
143
        return struct.pack(self.fmt,
144
145
146
147
148
149
150
151
                           self.code_area,
                           self.part[0].pack(),
                           self.part[1].pack(),
                           self.part[2].pack(),
                           self.part[3].pack(),
                           self.signature)

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


162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
class Disk(object):
    """Represents a BSD Disk"""

    def __init__(self, device):
        """Create a Disk instance"""
        self.device = device
        self.part_num = None
        self.disklabel = None

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

            for i in range(4):
                ptype = self.mbr.part[i].type
                if ptype in (0xa5, 0xa6, 0xa9):
                    d.seek(BLOCKSIZE * self.mbr.part[i].first_sector)
                    self.part_num = i
                    if ptype == 0xa5:  # FreeBSD
                        self.disklabel = BSDDisklabel(d)
                    elif ptype == 0xa6:  # OpenBSD
                        self.disklabel = OpenBSDDisklabel(d)
                    else:  # NetBSD
                        self.disklabel = BSDDisklabel(d)
                    break

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

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

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

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

    def enlarge(self, new_size):
        """Enlarge the disk and return the last useable sector"""

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

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

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

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

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

225
226
227
228
229
230
231
232
233
234
235
236
237
238
    def get_last_partition_id(self):
        """Get the ID of the last partition"""
        return self.disklabel.get_last_partition_id()

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

        return ""


class DisklabelBase(object):
    """Disklabel base class"""
239
    __metaclass__ = abc.ABCMeta
240
241
242
243

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

244
245
246
        # Subclasses need to overwrite this
        self.field = None
        self.ptable = None
247

248
249
250
251
    @abc.abstractproperty
    def fmt(self):
        """Fields format string for the disklabel fields"""
        pass
252
253
254
255
256
257
258
259
260
261

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

        if checksum is not None:
            out = self.field.copy()
            out['checksum'] = checksum
        else:
            out = self.field

262
        return struct.pack(self.fmt, * out.values() + [self.ptable.pack()])
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285

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

        raw = cStringIO.StringIO(self.pack(0))
        checksum = 0
        try:
            uint16 = raw.read(2)
            while uint16 != "":
                checksum ^= struct.unpack('<H', uint16)[0]
                uint16 = raw.read(2)
        finally:
            raw.close()

        return checksum

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

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

286
    @abc.abstractmethod
287
288
    def get_last_partition_id(self):
        """Get the ID of the last partition"""
289
        pass
290

291
    @abc.abstractmethod
292
293
    def __str__(self):
        """Print the Disklabel"""
294
        pass
295

296

297
298
class PartitionTableBase(object):
    """Base Class for disklabel partition tables"""
299
    __metaclass__ = abc.ABCMeta
300

301
302
303
304
    @abc.abstractproperty
    def fmt(self):
        """Partition fields format string"""
        pass
305

306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
    @abc.abstractproperty
    def fields(self):
        """The partition fields"""
        pass

    @abc.abstractmethod
    def setpsize(self, i, size):
        """Set size for partition i"""
        pass

    @abc.abstractmethod
    def getpsize(self, i):
        """Get size for partition i"""
        pass

    @abc.abstractmethod
    def setpoffset(self, i, offset):
        """Set offset for partition i"""
        pass

    @abc.abstractmethod
    def getpoffset(self, i):
        """Get offset for partition i"""
        pass
330
331
332
333

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

334
335
        self.Partition = namedtuple('Partition', self.fields)
        self.part = []
336

337
        size = struct.calcsize(self.fmt)
338
339
340
341
        raw = cStringIO.StringIO(ptable)
        try:
            for _ in xrange(pnumber):
                self.part.append(
342
                    self.Partition(*struct.unpack(self.fmt, raw.read(size)))
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
                    )
        finally:
            raw.close()

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

    def pack(self):
        """Packs the partition table into a binary string."""
        ret = ""
        for i in xrange(len(self.part)):
358
            ret += struct.pack(self.fmt, *self.part[i])
359
        return ret + ((364 - len(self.part) * 16) * '\x00')
360
361


362
class BSDDisklabel(DisklabelBase):
363
    """Represents a BSD Disklabel"""
364

365
    class PartitionTable(PartitionTableBase):
366
367
        """Represents a BSD Partition Table"""

368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
        @property
        def fmt(self):
            """Partition fields format string"""
            return "<IIIBBH"

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

        def setpsize(self, i, size):
            """Set size for partition i"""
            tmp = self.part[i]
            self.part[i] = self.Partition(size, tmp.offset, tmp.fsize,
                                          tmp.fstype, tmp.frag, tmp.cpg)

        def getpsize(self, i):
            """Get size for partition i"""
            return self.part[i].size

        def setpoffset(self, i, offset):
            """Set offset for partition i"""
            tmp = self.part[i]
            self.part[i] = self.Partition(tmp.size, offset, tmp.fsize,
                                          tmp.fstype, tmp.frag, tmp.cpg)

        def getpoffset(self, i):
            """Get offset for partition i"""
            return self.part[i].offset

    @property
    def fmt(self):
        return "<IHH16s16sIIIIIIHHIHHHHIII20s20sIHHII364s"
408

409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
    def __init__(self, device):
        """Create a BSD DiskLabel instance"""

        super(BSDDisklabel, self).__init__(device)

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

        d_ = OrderedDict()      # Off  Len    Content
        (d_["magic"],           # 0    4      Magic
         d_["dtype"],           # 4    2      Drive Type
         d_["subtype"],         # 6    2      Subtype
         d_["typename"],        # 8    16     Type Name
         d_["packname"],        # 24   16     Pack Identifier
         d_["secsize"],         # 32   4      Bytes per sector
         d_["nsectors"],        # 36   4      Data sectors per track
         d_["ntracks"],         # 40   4      Tracks per cylinder
         d_["ncylinders"],      # 44   4      Data cylinders per unit
         d_["secpercyl"],       # 48   4      Data sectors per cylinder
         d_["secperunit"],      # 52   4      Data sectors per unit
         d_["sparespertrack"],  # 56   2      Spare sectors per track
         d_["sparespercyl"],    # 58   2      Spare sectors per cylinder
         d_["acylinders"],      # 60   4      Alternative cylinders per unit
         d_["rpm"],             # 64   2      Rotation Speed
         d_["interleave"],      # 66   2      Hardware sector interleave
         d_["trackskew"],       # 68   2      Sector 0 skew, per track
         d_["cylskew"],         # 70   2      Sector 0 skew, per cylinder
         d_["headswitch"],      # 72   4      Head switch time
         d_["trkseek"],         # 76   4      Track-to-track seek
         d_["flags"],           # 80   4      Generic Flags
         d_["drivedata"],       # 84   5*4    Drive-type specific information
         d_["spare"],           # 104  5*4    Reserved for future use
         d_["magic2"],          # 124  4      Magic Number
         d_["checksum"],        # 128  2      Xor of data including partitions
         d_["npartitions"],     # 130  2      Number of partitions following
         d_["bbsize"],          # 132  4      size of boot area at sn0, bytes
         d_["sbsize"],          # 136  4      Max size of fs superblock, bytes
         ptable_raw             # 140  16*16  Partition Table
448
         ) = struct.unpack(self.fmt, sector1)
449
450
451
452
453

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

454
455
456
457
458
459
460
461
462
    def enlarge(self, new_size):
        raise NotImplementedError

    def enlarge_last_partition(self):
        raise NotImplementedError

    def get_last_partition_id(self):
        raise NotImplementedError

463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
    def __str__(self):
        """Print the Disklabel"""

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

        title = "Disklabel"
        return \
            "%s\n%s\n" % (title, len(title) * "=") + \
            "Magic Number: 0x%(magic)x\n" \
            "Drive type: %(dtype)d\n" \
            "Subtype: %(subtype)d\n" % self.field + \
            "Typename: %s\n" % typename + \
            "Pack Identifier: %s\n" % packname + \
            "# of bytes per sector: %(secsize)d\n" \
            "# of data sectors per track: %(nsectors)d\n" \
            "# of tracks per cylinder: %(ntracks)d\n" \
            "# of data cylinders per unit: %(ncylinders)d\n" \
            "# of data sectors per cylinder: %(secpercyl)d\n" \
            "# of data sectors per unit: %(secperunit)d\n" \
            "# of spare sectors per track: %(sparespertrack)d\n" \
            "# of spare sectors per cylinder: %(sparespercyl)d\n" \
            "Alt. cylinders per unit: %(acylinders)d\n" \
            "Rotational speed: %(rpm)d\n" \
            "Hardware sector interleave: %(interleave)d\n" \
            "Sector 0 skew, per track: %(trackskew)d\n" \
            "Sector 0 skew, per cylinder: %(cylskew)d\n" \
            "Head switch time, usec: %(headswitch)d\n" \
            "Track-to-track seek, usec: %(trkseek)d\n" \
            "Generic Flags: %(flags)r\n" \
            "Drive data: %(drivedata)r\n" \
            "Reserved for future use: %(spare)r\n" \
            "The magic number again: 0x%(magic2)x\n" \
            "Checksum: %(checksum)d\n" \
            "Number of partitions: %(npartitions)d\n"  \
            "Size of boot aread at sn0: %(bbsize)d\n"  \
            "Max size of fs superblock: %(sbsize)d\n" % self.field + \
            "%s" % self.ptable
502
503


504
class OpenBSDDisklabel(DisklabelBase):
505
506
    """Represents an OpenBSD Disklabel"""

507
    class PartitionTable(PartitionTableBase):
508
509
        """Reprepsents an OpenBSD Partition Table"""

510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
        @property
        def fmt(self):
            return "<IIHHBBH"

        @property
        def fields(self):
            return [        # Offset  Length Contents
                'size',     # 0       4      Number of sectors in the partition
                'offset',   # 4       4      Starting sector of the partition
                'offseth',  # 8       2      Starting sector (high part)
                'sizeh',    # 10      2      Number of sectors (high part)
                'fstype',   # 12      1      File system type
                'frag',     # 13      1      File system fragments per block
                'cpg'       # 14      2      File system cylinders per group
                ]
525
526
527
528

        def setpsize(self, i, size):
            """Set size for partition i"""
            tmp = self.part[i]
529
530
531
            self.part[i] = self.Partition(
                size & 0xffffffff, tmp.offset, tmp.offseth, size >> 32,
                tmp.fstype, tmp.frag, tmp.cpg)
532
533

        def getpsize(self, i):
534
            """Get size for partition i"""
535
536
537
            return (self.part[i].sizeh << 32) + self.part[i].size

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

        def getpoffset(self, i):
545
            """Get offset for partition i"""
546
547
            return (self.part[i].offseth << 32) + self.part[i].offset

548
549
550
551
    @property
    def fmt(self):
        return "<IHH16s16sIIIIII8sIHHIII20sHH16sIHHII364s"

552
    def __init__(self, device):
553
554
        """Create a DiskLabel instance"""

555
556
557
        super(OpenBSDDisklabel, self).__init__(device)

        # Disklabel starts at offset one
558
559
        device.seek(BLOCKSIZE, os.SEEK_CUR)
        sector1 = device.read(BLOCKSIZE)
560

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

        assert d_['magic'] == d_['magic2'] == DISKMAGIC, "Disklabel not valid"
        self.ptable = self.PartitionTable(ptable_raw, d_['npartitions'])
        self.field = d_
595
596
597

    def setdsize(self, dsize):
        """Set disk size"""
598
599
        self.field['secperunith'] = dsize >> 32
        self.field['secperunit'] = dsize & 0xffffffff
600
601
602

    def getdsize(self):
        """Get disk size"""
603
604
605
        return (self.field['secperunith'] << 32) + self.field['secperunit']

    dsize = property(getdsize, setdsize, None, "disk size")
606
607
608

    def setbstart(self, bstart):
        """Set start of useable region"""
609
610
        self.field['bstarth'] = bstart >> 32
        self.field['bstart'] = bstart & 0xffffffff
611
612

    def getbstart(self):
613
        """Get start of useable region"""
614
615
616
        return (self.field['bstarth'] << 32) + self.field['bstart']

    bstart = property(getbstart, setbstart, None, "start of useable region")
617
618
619

    def setbend(self, bend):
        """Set end of useable region"""
620
621
        self.field['bendh'] = bend >> 32
        self.field['bend'] = bend & 0xffffffff
622
623

    def getbend(self):
624
        """Get end of useable region"""
625
626
627
        return (self.field['bendh'] << 32) + self.field['bend']

    bend = property(getbend, setbend, None, "end of useable region")
628

629
    def enlarge(self, new_size):
630
        """Enlarge the disk and return the last useable sector"""
631

632
633
        assert new_size >= self.dsize, \
            "New size cannot be smaller that %d" % self.dsize
634
635

        # Fix the disklabel
636
637
638
639
640
        self.dsize = new_size
        self.field['ncylinders'] = self.dsize // (self.field['nsectors'] *
                                                  self.field['ntracks'])
        self.bend = (self.field['ncylinders'] * self.field['nsectors'] *
                     self.field['ntracks'])
641

642
        # Partition 'c' describes the entire disk
643
644
        self.ptable.setpsize(2, new_size)

645
        # Update the checksum
646
        self.field['checksum'] = self.compute_checksum()
647

648
        # The last useable sector is the end of the useable region minus one
649
        return self.bend - 1
650
651
652
653
654

    def get_last_partition_id(self):
        """Returns the id of the last partition"""
        end = 0
        # Don't check partition 'c' which is the whole disk
655
        for i in [n for n in range(len(self.ptable.part)) if n != 2]:
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
            curr_end = self.ptable.getpsize(i) + self.ptable.getpoffset(i)
            if end < curr_end:
                end = curr_end
                part = i

        assert end > 0, "No partition found"

        return part

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

        part_num = self.get_last_partition_id()

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

        assert end > 0, "No partition found"

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

678
        if end > (self.bend - 1024):
679
680
681
            return

        self.ptable.setpsize(
682
            part_num, self.bend - self.ptable.getpoffset(part_num) - 1024)
683

684
        self.field['checksum'] = self.compute_checksum()
685
686
687
688

    def __str__(self):
        """Print the Disklabel"""

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

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

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


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

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

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

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

755
    disk = Disk(args[0])
756
757

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

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

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

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

    if options.enlarge_partition:
772
        disk.enlarge_last_partition()
773

774
    disk.write()
775
776
    return 0

777

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

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