code.grnet.gr Git - snf-image/blob - snf-image-helper/disklabel.py

   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             start = self.unpack_chs(self.start)
  78             end = self.unpack_chs(self.end)
  79             return "%d %s %d %s %d %d" % (self.status, start, self.type, end,
  80                                           self.first_sector, self.sector_count)
  81
  82         @staticmethod
  83         def unpack_chs(chs):
  84             """Unpacks a CHS address string to a tuple."""
  85
  86             assert len(chs) == 3
  87
  88             head = struct.unpack('<B', chs[0])[0]
  89             sector = struct.unpack('<B', chs[1])[0] & 0x3f
  90             cylinder = (struct.unpack('<B', chs[1])[0] & 0xC0) << 2 | \
  91                 struct.unpack('<B', chs[2])[0]
  92
  93             return (cylinder, head, sector)
  94
  95         @staticmethod
  96         def pack_chs(cylinder, head, sector):
  97             """Packs a CHS tuple to an address string."""
  98
  99             assert 1 <= sector <= 63
 100             assert 0 <= head <= 255
 101             assert 0 <= cylinder
 102
 103             # If the cylinders overflow then put the value (1023, 254, 63) to
 104             # the tuple. At least this is what OpenBSD does.
 105             if cylinder > 1023:
 106                 cylinder = 1023
 107                 head = 254
 108                 sector = 63
 109
 110             byte0 = head
 111             byte1 = (cylinder >> 2) & 0xC0 | sector
 112             byte2 = cylinder & 0xff
 113
 114             return struct.pack('<BBB', byte0, byte1, byte2)
 115
 116     def __init__(self, block):
 117         """Create an MBR instance"""
 118
 119         self.fmt = "<444s2x16s16s16s16s2s"
 120         raw_part = {}     # Offset  Length          Contents
 121         (self.code_area,  # 0       440(max. 446)   code area
 122                           # 440     2(optional)     disk signature
 123                           # 444     2               Usually nulls
 124          raw_part[0],     # 446     16              Partition 0
 125          raw_part[1],     # 462     16              Partition 1
 126          raw_part[2],     # 478     16              Partition 2
 127          raw_part[3],     # 494     16              Partition 3
 128          self.signature   # 510     2               MBR signature
 129          ) = struct.unpack(self.fmt, block)
 130
 131         self.part = {}
 132         for i in range(4):
 133             self.part[i] = self.Partition(raw_part[i])
 134
 135     def size(self):
 136         """Return the size of a Master Boot Record."""
 137         return struct.calcsize(self.fmt)
 138
 139     def pack(self):
 140         """Pack an MBR to a binary string."""
 141         return struct.pack(self.fmt,
 142                            self.code_area,
 143                            self.part[0].pack(),
 144                            self.part[1].pack(),
 145                            self.part[2].pack(),
 146                            self.part[3].pack(),
 147                            self.signature)
 148
 149     def __str__(self):
 150         """Print the MBR"""
 151         ret = ""
 152         for i in range(4):
 153             ret += "Partition %d: %s\n" % (i, self.part[i])
 154         ret += "Signature: %s %s\n" % (hex(ord(self.signature[0])),
 155                                        hex(ord(self.signature[1])))
 156         title = "Master Boot Record"
 157         return "%s\n%s\n%s\n" % (title, len(title) * "=", ret)
 158
 159
 160 class Disk(object):
 161     """Represents a BSD Disk"""
 162
 163     def __init__(self, device):
 164         """Create a Disk instance"""
 165         self.device = device
 166         self.part_num = None
 167         self.disklabel = None
 168
 169         with open(device, "rb") as d:
 170             sector0 = d.read(BLOCKSIZE)
 171             self.mbr = MBR(sector0)
 172
 173             for i in range(4):
 174                 ptype = self.mbr.part[i].type
 175                 if ptype in (0xa5, 0xa6, 0xa9):
 176                     d.seek(BLOCKSIZE * self.mbr.part[i].first_sector)
 177                     self.part_num = i
 178                     if ptype == 0xa5:  # FreeBSD
 179                         self.disklabel = BSDDisklabel(d)
 180                     elif ptype == 0xa6:  # OpenBSD
 181                         self.disklabel = OpenBSDDisklabel(d)
 182                     else:  # NetBSD
 183                         self.disklabel = BSDDisklabel(d)
 184                     break
 185
 186         assert self.disklabel is not None, "No *BSD partition found"
 187
 188     def write(self):
 189         """Write the changes back to the media"""
 190         with open(self.device, 'rw+b') as d:
 191             d.write(self.mbr.pack())
 192
 193             d.seek(self.mbr.part[self.part_num].first_sector * BLOCKSIZE)
 194             self.disklabel.write_to(d)
 195
 196     def __str__(self):
 197         """Print the partitioning info of the Disk"""
 198         return str(self.mbr) + str(self.disklabel)
 199
 200     def enlarge(self, new_size):
 201         """Enlarge the disk and return the last usable sector"""
 202
 203         # Fix the disklabel
 204         end = self.disklabel.enlarge(new_size)
 205
 206         # Fix the MBR
 207         start = self.mbr.part[self.part_num].first_sector
 208         self.mbr.part[self.part_num].sector_count = end - start + 1
 209
 210         ntracks = self.disklabel.field['ntracks']
 211         nsectors = self.disklabel.field['nsectors']
 212
 213         cylinder = end // (ntracks * nsectors)
 214         header = (end // nsectors) % ntracks
 215         sector = (end % nsectors) + 1
 216         chs = MBR.Partition.pack_chs(cylinder, header, sector)
 217         self.mbr.part[self.part_num].end = chs
 218
 219     def enlarge_last_partition(self):
 220         """Enlarge the last partition to cover up all the free space"""
 221         self.disklabel.enlarge_last_partition()
 222
 223     def get_last_partition_id(self):
 224         """Get the ID of the last partition"""
 225         return self.disklabel.get_last_partition_id()
 226
 227     def get_duid(self):
 228         """Get the Disklabel Unique Identifier (works only for OpenBSD)"""
 229         if 'uid' in self.disklabel.field:
 230             return self.disklabel.field['uid']
 231
 232         return ""
 233
 234
 235 class DisklabelBase(object):
 236     """Disklabel base class"""
 237     __metaclass__ = abc.ABCMeta
 238
 239     def __init__(self, device):
 240         """Create a Disklabel instance"""
 241
 242         # Subclasses need to overwrite this
 243         self.field = None
 244         self.ptable = None
 245
 246     @abc.abstractproperty
 247     def fmt(self):
 248         """Fields format string for the disklabel fields"""
 249         pass
 250
 251     def pack(self, checksum=None):
 252         """Return a binary copy of the Disklabel block"""
 253
 254         out = OrderedDict()
 255         for k, v in self.field.iteritems():
 256             out[k] = v
 257
 258         if checksum is not None:
 259             out['checksum'] = checksum
 260
 261         return struct.pack(self.fmt, * out.values() + [self.ptable.pack()])
 262
 263     def compute_checksum(self):
 264         """Compute the checksum of the disklabel"""
 265
 266         raw = cStringIO.StringIO(self.pack(0))
 267         checksum = 0
 268         try:
 269             uint16 = raw.read(2)
 270             while uint16 != "":
 271                 checksum ^= struct.unpack('<H', uint16)[0]
 272                 uint16 = raw.read(2)
 273         finally:
 274             raw.close()
 275
 276         return checksum
 277
 278     @abc.abstractmethod
 279     def enlarge(self, new_size):
 280         """Enlarge the disk and return the last usable sector"""
 281         pass
 282
 283     def write_to(self, device):
 284         """Write the disklabel to a device"""
 285
 286         # The disklabel starts at sector 1
 287         device.seek(BLOCKSIZE, os.SEEK_CUR)
 288         device.write(self.pack())
 289
 290     @abc.abstractmethod
 291     def enlarge_last_partition(self):
 292         """Enlarge the last partition to consume all the usable space"""
 293         pass
 294
 295     @abc.abstractmethod
 296     def get_last_partition_id(self):
 297         """Get the ID of the last partition"""
 298         pass
 299
 300     @abc.abstractmethod
 301     def __str__(self):
 302         """Print the Disklabel"""
 303         pass
 304
 305
 306 class PartitionTableBase(object):
 307     """Base Class for disklabel partition tables"""
 308     __metaclass__ = abc.ABCMeta
 309
 310     @abc.abstractproperty
 311     def fmt(self):
 312         """Partition fields format string"""
 313         pass
 314
 315     @abc.abstractproperty
 316     def fields(self):
 317         """The partition fields"""
 318         pass
 319
 320     def __init__(self, ptable, pnumber):
 321         """Create a Partition Table instance"""
 322
 323         self.Partition = namedtuple('Partition', self.fields)
 324         self.part = []
 325
 326         size = struct.calcsize(self.fmt)
 327         raw = cStringIO.StringIO(ptable)
 328         try:
 329             for _ in xrange(pnumber):
 330                 self.part.append(
 331                     self.Partition(*struct.unpack(self.fmt, raw.read(size)))
 332                     )
 333         finally:
 334             raw.close()
 335
 336     def __str__(self):
 337         """Print the Partition table"""
 338         val = ""
 339         for i in xrange(len(self.part)):
 340             val += "%c: %s\n" % (chr(ord('a') + i), str(self.part[i]))
 341         return val
 342
 343     def pack(self):
 344         """Packs the partition table into a binary string."""
 345         ret = ""
 346         for i in xrange(len(self.part)):
 347             ret += struct.pack(self.fmt, *self.part[i])
 348         return ret + ((364 - len(self.part) * 16) * '\x00')
 349
 350
 351 class BSDDisklabel(DisklabelBase):
 352     """Represents an BSD Disklabel"""
 353
 354     class PartitionTable(PartitionTableBase):
 355         """Represents a BSD Partition Table"""
 356
 357         @property
 358         def fmt(self):
 359             """Partition fields format string"""
 360             return "<IIIBBH"
 361
 362         @property
 363         def fields(self):
 364             """The partition fields"""
 365             return [    # Offset  Length Contents
 366                 'size',     # 0       4      Number of sectors in partition
 367                 'offset',   # 4       4      Starting sector of the partition
 368                 'fsize',    # 8       4      File system basic fragment size
 369                 'fstype',   # 12      1      File system type
 370                 'frag',     # 13      1      File system fragments per block
 371                 'cpg'       # 14      2      File system cylinders per group
 372                 ]
 373
 374     @property
 375     def fmt(self):
 376         return "<IHH16s16sIIIIIIHHIHHHHIII20s20sIHHII364s"
 377
 378     def __init__(self, device):
 379         """Create a BSD DiskLabel instance"""
 380         super(BSDDisklabel, self).__init__(device)
 381
 382         # Disklabel starts at offset one
 383         device.seek(BLOCKSIZE, os.SEEK_CUR)
 384         sector1 = device.read(BLOCKSIZE)
 385
 386         d_ = OrderedDict()      # Off  Len    Content
 387         (d_["magic"],           # 0    4      Magic
 388          d_["dtype"],           # 4    2      Drive Type
 389          d_["subtype"],         # 6    2      Subtype
 390          d_["typename"],        # 8    16     Type Name
 391          d_["packname"],        # 24   16     Pack Identifier
 392          d_["secsize"],         # 32   4      Bytes per sector
 393          d_["nsectors"],        # 36   4      Data sectors per track
 394          d_["ntracks"],         # 40   4      Tracks per cylinder
 395          d_["ncylinders"],      # 44   4      Data cylinders per unit
 396          d_["secpercyl"],       # 48   4      Data sectors per cylinder
 397          d_["secperunit"],      # 52   4      Data sectors per unit
 398          d_["sparespertrack"],  # 56   2      Spare sectors per track
 399          d_["sparespercyl"],    # 58   2      Spare sectors per cylinder
 400          d_["acylinders"],      # 60   4      Alternative cylinders per unit
 401          d_["rpm"],             # 64   2      Rotation Speed
 402          d_["interleave"],      # 66   2      Hardware sector interleave
 403          d_["trackskew"],       # 68   2      Sector 0 skew, per track
 404          d_["cylskew"],         # 70   2      Sector 0 skew, per cylinder
 405          d_["headswitch"],      # 72   4      Head switch time
 406          d_["trkseek"],         # 76   4      Track-to-track seek
 407          d_["flags"],           # 80   4      Generic Flags
 408          d_["drivedata"],       # 84   5*4    Drive-type specific information
 409          d_["spare"],           # 104  5*4    Reserved for future use
 410          d_["magic2"],          # 124  4      Magic Number
 411          d_["checksum"],        # 128  2      Xor of data including partitions
 412          d_["npartitions"],     # 130  2      Number of partitions following
 413          d_["bbsize"],          # 132  4      size of boot area at sn0, bytes
 414          d_["sbsize"],          # 136  4      Max size of fs superblock, bytes
 415          ptable_raw             # 140  16*16  Partition Table
 416          ) = struct.unpack(self.fmt, sector1)
 417
 418         assert d_['magic'] == d_['magic2'] == DISKMAGIC, "Disklabel not valid"
 419         self.ptable = self.PartitionTable(ptable_raw, d_['npartitions'])
 420         self.field = d_
 421
 422     def enlarge(self, new_size):
 423         raise NotImplementedError
 424
 425     def enlarge_last_partition(self):
 426         raise NotImplementedError
 427
 428     def get_last_partition_id(self):
 429         raise NotImplementedError
 430
 431     def __str__(self):
 432         """Print the Disklabel"""
 433
 434         title = "Disklabel"
 435
 436         # Those fields may contain null bytes
 437         typename = self.field['typename'].strip('\x00').strip()
 438         packname = self.field['packname'].strip('\x00').strip()
 439
 440         return \
 441             "%s\n%s\n" % (title, len(title) * "=") + \
 442             "Magic Number: 0x%(magic)x\n" \
 443             "Drive type: %(dtype)d\n" \
 444             "Subtype: %(subtype)d\n" % self.field + \
 445             "Typename: %s\n" % typename + \
 446             "Pack Identifier: %s\n" % packname + \
 447             "# of bytes per sector: %(secsize)d\n" \
 448             "# of data sectors per track: %(nsectors)d\n" \
 449             "# of tracks per cylinder: %(ntracks)d\n" \
 450             "# of data cylinders per unit: %(ncylinders)d\n" \
 451             "# of data sectors per cylinder: %(secpercyl)d\n" \
 452             "# of data sectors per unit: %(secperunit)d\n" \
 453             "# of spare sectors per track: %(sparespertrack)d\n" \
 454             "# of spare sectors per cylinder: %(sparespercyl)d\n" \
 455             "Alt. cylinders per unit: %(acylinders)d\n" \
 456             "Rotational speed: %(rpm)d\n" \
 457             "Hardware sector interleave: %(interleave)d\n" \
 458             "Sector 0 skew, per track: %(trackskew)d\n" \
 459             "Sector 0 skew, per cylinder: %(cylskew)d\n" \
 460             "Head switch time, usec: %(headswitch)d\n" \
 461             "Track-to-track seek, usec: %(trkseek)d\n" \
 462             "Generic Flags: %(flags)r\n" \
 463             "Drive data: %(drivedata)r\n" \
 464             "Reserved for future use: %(spare)r\n" \
 465             "The magic number again: 0x%(magic2)x\n" \
 466             "Checksum: %(checksum)d\n" \
 467             "Number of partitions: %(npartitions)d\n"  \
 468             "Size of boot aread at sn0: %(bbsize)d\n"  \
 469             "Max size of fs superblock: %(sbsize)d\n" % self.field + \
 470             "%s" % self.ptable
 471
 472
 473 class OpenBSDDisklabel(DisklabelBase):
 474     """Represents an OpenBSD Disklabel"""
 475
 476     class PartitionTable(PartitionTableBase):
 477         """Reprepsents an OpenBSD Partition Table"""
 478
 479         @property
 480         def fmt(self):
 481             return "<IIHHBBH"
 482
 483         @property
 484         def fields(self):
 485             return [    # Offset  Length Contents
 486                 'size',     # 0       4      Number of sectors in the partition
 487                 'offset',   # 4       4      Starting sector of the partition
 488                 'offseth',  # 8       2      Starting sector (high part)
 489                 'sizeh',    # 10      2      Number of sectors (high part)
 490                 'fstype',   # 12      1      Filesystem type
 491                 'frag',     # 13      1      Filesystem Fragments per block
 492                 'cpg'       # 14      2      FS cylinders per group
 493                 ]
 494
 495         def setpsize(self, i, size):
 496             """Set size for partition i"""
 497             tmp = self.part[i]
 498             self.part[i] = self.Partition(
 499                 size & 0xffffffff, tmp.offset, tmp.offseth, size >> 32,
 500                 tmp.fstype, tmp.frag, tmp.cpg)
 501
 502         def getpsize(self, i):
 503             """Get size for partition i"""
 504             return (self.part[i].sizeh << 32) + self.part[i].size
 505
 506         def setpoffset(self, i, offset):
 507             """Set offset for partition i"""
 508             tmp = self.part[i]
 509             self.part[i] = self.Partition(
 510                 tmp.size, offset & 0xffffffff, offset >> 32, tmp.sizeh,
 511                 tmp.frag, tmp.cpg)
 512
 513         def getpoffset(self, i):
 514             """Get offset for partition i"""
 515             return (self.part[i].offseth << 32) + self.part[i].offset
 516
 517     @property
 518     def fmt(self):
 519         return "<IHH16s16sIIIIII8sIHHIII20sHH16sIHHII364s"
 520
 521     def __init__(self, device):
 522         """Create a DiskLabel instance"""
 523
 524         super(OpenBSDDisklabel, self).__init__(device)
 525
 526         # Disklabel starts at offset one
 527         device.seek(BLOCKSIZE, os.SEEK_CUR)
 528         sector1 = device.read(BLOCKSIZE)
 529
 530         d_ = OrderedDict()   # Off  Len    Content
 531         (d_["magic"],        # 0    4      Magic
 532          d_["dtype"],        # 4    2      Drive Type
 533          d_["subtype"],      # 6    2      Subtype
 534          d_["typename"],     # 8    16     Type Name
 535          d_["packname"],     # 24   16     Pack Identifier
 536          d_["secsize"],      # 32   4      Bytes per sector
 537          d_["nsectors"],     # 36   4      Data sectors per track
 538          d_["ntracks"],      # 40   4      Tracks per cylinder
 539          d_["ncylinders"],   # 44   4      Data cylinders per unit
 540          d_["secpercyl"],    # 48   4      Data sectors per cylinder
 541          d_["secperunit"],   # 52   4      Data sectors per unit
 542          d_["uid"],          # 56   8      Unique label identifier
 543          d_["acylinders"],   # 64   4      Alt cylinders per unit
 544          d_["bstarth"],      # 68   2      Start of useable region (high part)
 545          d_["bendh"],        # 70   2      Size of usable region (high part)
 546          d_["bstart"],       # 72   4      Start of useable region
 547          d_["bend"],         # 76   4      End of usable region
 548          d_["flags"],        # 80   4      Generic Flags
 549          d_["drivedata"],    # 84   5*4    Drive-type specific information
 550          d_["secperunith"],  # 104  2      Number of data sectors (high part)
 551          d_["version"],      # 106  2      Version
 552          d_["spare"],        # 108  4*4    Reserved for future use
 553          d_["magic2"],       # 124  4      Magic number
 554          d_["checksum"],     # 128  2      Xor of data including partitions
 555          d_["npartitions"],  # 130  2      Number of partitions in following
 556          d_["bbsize"],       # 132  4      size of boot area at sn0, bytes
 557          d_["sbsize"],       # 136  4      Max size of fs superblock, bytes
 558          ptable_raw          # 140  16*16  Partition Table
 559          ) = struct.unpack(self.fmt, sector1)
 560
 561         assert d_['magic'] == d_['magic2'] == DISKMAGIC, "Disklabel not valid"
 562         self.ptable = self.PartitionTable(ptable_raw, d_['npartitions'])
 563         self.field = d_
 564
 565     def setdsize(self, dsize):
 566         """Set disk size"""
 567         self.field['secperunith'] = dsize >> 32
 568         self.field['secperunit'] = dsize & 0xffffffff
 569
 570     def getdsize(self):
 571         """Get disk size"""
 572         return (self.field['secperunith'] << 32) + self.field['secperunit']
 573
 574     dsize = property(getdsize, setdsize, None, "disk size")
 575
 576     def setbstart(self, bstart):
 577         """Set start of useable region"""
 578         self.field['bstarth'] = bstart >> 32
 579         self.field['bstart'] = bstart & 0xffffffff
 580
 581     def getbstart(self):
 582         """Get start of usable region"""
 583         return (self.field['bstarth'] << 32) + self.field['bstart']
 584
 585     bstart = property(getbstart, setbstart, None, "start of usable region")
 586
 587     def setbend(self, bend):
 588         """Set end of useable region"""
 589         self.field['bendh'] = bend >> 32
 590         self.field['bend'] = bend & 0xffffffff
 591
 592     def getbend(self):
 593         """Get end of usable region"""
 594         return (self.field['bendh'] << 32) + self.field['bend']
 595
 596     bend = property(getbend, setbend, None, "end of usable region")
 597
 598     def enlarge(self, new_size):
 599         """Enlarge the disk and return the last usable sector"""
 600
 601         assert new_size >= self.dsize, \
 602             "New size cannot be smaller that %s" % self.dsize
 603
 604         # Fix the disklabel
 605         self.dsize = new_size
 606         self.field['ncylinders'] = self.dsize // (self.field['nsectors'] *
 607                                                   self.field['ntracks'])
 608         self.bend = (self.field['ncylinders'] * self.field['nsectors'] *
 609                      self.field['ntracks'])
 610
 611         # Partition 'c' descriptes the entire disk
 612         self.ptable.setpsize(2, new_size)
 613
 614         # Update the checksum
 615         self.field['checksum'] = self.compute_checksum()
 616
 617         # The last usable sector is the end of the usable region minus one
 618         return self.bend - 1
 619
 620     def get_last_partition_id(self):
 621         """Returns the id of the last partition"""
 622         part = 0
 623         end = 0
 624         # Don't check partition 'c' which is the whole disk
 625         for i in filter(lambda x: x != 2, range(len(self.ptable.part))):
 626             curr_end = self.ptable.getpsize(i) + self.ptable.getpoffset(i)
 627             if end < curr_end:
 628                 end = curr_end
 629                 part = i
 630
 631         assert end > 0, "No partition found"
 632
 633         return part
 634
 635     def enlarge_last_partition(self):
 636         """Enlarge the last partition to cover up all the free space"""
 637
 638         part_num = self.get_last_partition_id()
 639
 640         end = self.ptable.getpsize(part_num) + self.ptable.getpoffset(part_num)
 641
 642         assert end > 0, "No partition found"
 643
 644         if self.ptable.part[part_num].fstype == 1:  # Swap partition.
 645             #TODO: Maybe create a warning?
 646             return
 647
 648         if end > (self.bend - 1024):
 649             return
 650
 651         self.ptable.setpsize(
 652             part_num, self.bend - self.ptable.getpoffset(part_num) - 1024)
 653
 654         self.field['checksum'] = self.compute_checksum()
 655
 656     def __str__(self):
 657         """Print the Disklabel"""
 658
 659         # Those values may contain null bytes
 660         typename = self.field['typename'].strip('\x00').strip()
 661         packname = self.field['packname'].strip('\x00').strip()
 662
 663         duid = "".join(x.encode('hex') for x in self.field['uid'])
 664
 665         title = "Disklabel"
 666         return \
 667             "%s\n%s\n" % (title, len(title) * "=") + \
 668             "Magic Number: 0x%(magic)x\n" \
 669             "Drive type: %(dtype)d\n" \
 670             "Subtype: %(subtype)d\n" % self.field + \
 671             "Typename: %s\n" % typename + \
 672             "Pack Identifier: %s\n" % packname + \
 673             "# of bytes per sector: %(secsize)d\n" \
 674             "# of data sectors per track: %(nsectors)d\n" \
 675             "# of tracks per cylinder: %(ntracks)d\n" \
 676             "# of data cylinders per unit: %(ncylinders)d\n" \
 677             "# of data sectors per cylinder: %(secpercyl)d\n" \
 678             "# of data sectors per unit: %(secperunit)d\n" % self.field + \
 679             "DUID: %s\n" % duid + \
 680             "Alt. cylinders per unit: %(acylinders)d\n" \
 681             "Start of useable region (high part): %(bstarth)d\n" \
 682             "Size of useable region (high part): %(bendh)d\n" \
 683             "Start of useable region: %(bstart)d\n" \
 684             "End of usable region: %(bend)d\n" \
 685             "Generic Flags: %(flags)r\n" \
 686             "Drive data: %(drivedata)r\n" \
 687             "Number of data sectors (high part): %(secperunith)d\n" \
 688             "Version: %(version)d\n" \
 689             "Reserved for future use: %(spare)r\n" \
 690             "The magic number again: 0x%(magic2)x\n" \
 691             "Checksum: %(checksum)d\n" \
 692             "Number of partitions: %(npartitions)d\n"  \
 693             "Size of boot aread at sn0: %(bbsize)d\n"  \
 694             "Max size of fs superblock: %(sbsize)d\n" % self.field + \
 695             "%s" % self.ptable
 696
 697
 698 def main():
 699     """Main entry point"""
 700     usage = "Usage: %prog [options] <input_media>"
 701     parser = optparse.OptionParser(usage=usage)
 702
 703     parser.add_option("-l", "--list", action="store_true", dest="list",
 704                       default=False,
 705                       help="list the disklabel on the specified media")
 706     parser.add_option("--get-last-partition", action="store_true",
 707                       dest="last_part", default=False,
 708                       help="print the label of the last partition")
 709     parser.add_option(
 710         "--get-duid", action="store_true", dest="duid", default=False,
 711         help="print the Disklabel Unique Identifier (OpenBSD only)")
 712     parser.add_option("-d", "--enlarge-disk", type="int", dest="disk_size",
 713                       default=None, metavar="SIZE",
 714                       help="Enlarge the disk to this SIZE (in sectors)")
 715     parser.add_option(
 716         "-p", "--enlarge-partition", action="store_true",
 717         dest="enlarge_partition", default=False,
 718         help="Enlarge the last partition to cover up the free space")
 719
 720     options, args = parser.parse_args(sys.argv[1:])
 721
 722     if len(args) != 1:
 723         parser.error("Wrong number of arguments")
 724
 725     disk = Disk(args[0])
 726
 727     if options.list:
 728         print disk
 729         return 0
 730
 731     if options.duid:
 732         print "%s" % "".join(x.encode('hex') for x in disk.get_duid())
 733         return 0
 734
 735     if options.last_part:
 736         print "%c" % chr(ord('a') + disk.get_last_partition_id())
 737
 738     if options.disk_size is not None:
 739         disk.enlarge(options.disk_size)
 740
 741     if options.enlarge_partition:
 742         disk.enlarge_last_partition()
 743
 744     disk.write()
 745     return 0
 746
 747
 748 if __name__ == '__main__':
 749     sys.exit(main())
 750
 751 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :