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

   1 #!/usr/bin/env python
   2 # -*- coding: utf-8 -*-
   3 #
   4 # Copyright (C) 2013 GRNET S.A.
   5 #
   6 # This program is free software; you can redistribute it and/or modify
   7 # it under the terms of the GNU General Public License as published by
   8 # the Free Software Foundation; either version 2 of the License, or
   9 # (at your option) any later version.
  10 #
  11 # This program is distributed in the hope that it will be useful, but
  12 # WITHOUT ANY WARRANTY; without even the implied warranty of
  13 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14 # General Public License for more details.
  15 #
  16 # You should have received a copy of the GNU General Public License
  17 # along with this program; if not, write to the Free Software
  18 # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
  19 # 02110-1301, USA.
  20
  21 """This module provides the code for handling OpenBSD disklabels"""
  22
  23 import struct
  24 import sys
  25 import cStringIO
  26 import optparse
  27
  28 from collections import namedtuple
  29
  30 BLOCKSIZE = 512
  31
  32 LABELSECTOR = 1
  33 LABELOFFSET = 0
  34
  35 BBSIZE = 8192  # size of boot area with label
  36 SBSIZE = 8192  # max size of fs superblock
  37
  38 DISKMAGIC = 0x82564557
  39
  40
  41 class MBR(object):
  42     """Represents a Master Boot Record."""
  43     class Partition(object):
  44         """Represents a partition entry in MBR"""
  45         format = "<B3sB3sLL"
  46
  47         def __init__(self, raw_part):
  48             """Create a Partition instance"""
  49             (
  50                 self.status,
  51                 self.start,
  52                 self.type,
  53                 self.end,
  54                 self.first_sector,
  55                 self.sector_count
  56             ) = struct.unpack(self.format, raw_part)
  57
  58         def pack(self):
  59             """Pack the partition values into a binary string"""
  60             return struct.pack(self.format,
  61                                self.status,
  62                                self.start,
  63                                self.type,
  64                                self.end,
  65                                self.first_sector,
  66                                self.sector_count)
  67
  68         @staticmethod
  69         def size():
  70             """Returns the size of an MBR partition entry"""
  71             return struct.calcsize(MBR.Partition.format)
  72
  73         def __str__(self):
  74             start = self.unpack_chs(self.start)
  75             end = self.unpack_chs(self.end)
  76             return "%d %s %d %s %d %d" % (self.status, start, self.type, end,
  77                                           self.first_sector, self.sector_count)
  78
  79         @staticmethod
  80         def unpack_chs(chs):
  81             """Unpacks a CHS address string to a tuple."""
  82
  83             assert len(chs) == 3
  84
  85             head = struct.unpack('<B', chs[0])[0]
  86             sector = struct.unpack('<B', chs[1])[0] & 0x3f
  87             cylinder = (struct.unpack('<B', chs[1])[0] & 0xC0) << 2 | \
  88                 struct.unpack('<B', chs[2])[0]
  89
  90             return (cylinder, head, sector)
  91
  92         @staticmethod
  93         def pack_chs(cylinder, head, sector):
  94             """Packs a CHS tuple to an address string."""
  95
  96             assert 1 <= sector <= 63
  97             assert 0 <= cylinder <= 1023
  98             assert 0 <= head <= 255
  99
 100             byte0 = head
 101             byte1 = (cylinder >> 2) & 0xC0 | sector
 102             byte2 = cylinder & 0xff
 103
 104             return struct.pack('<BBB', byte0, byte1, byte2)
 105
 106     format = "<444s2x16s16s16s16s2s"
 107     """
 108     Offset  Length          Contents
 109     0       440(max. 446)   code area
 110     440     2(optional)     disk signature
 111     444     2               Usually nulls
 112     446     16              Partition 0
 113     462     16              Partition 1
 114     478     16              Partition 2
 115     494     16              Partition 3
 116     510     2               MBR signature
 117     """
 118     def __init__(self, block):
 119         """Create an MBR instance"""
 120         raw_part = {}
 121         (self.code_area,
 122          raw_part[0],
 123          raw_part[1],
 124          raw_part[2],
 125          raw_part[3],
 126          self.signature) = struct.unpack(self.format, block)
 127
 128         self.part = {}
 129         for i in range(4):
 130             self.part[i] = self.Partition(raw_part[i])
 131
 132     @staticmethod
 133     def size():
 134         """Return the size of a Master Boot Record."""
 135         return struct.calcsize(MBR.format)
 136
 137     def pack(self):
 138         """Pack an MBR to a binary string."""
 139         return struct.pack(self.format,
 140                            self.code_area,
 141                            self.part[0].pack(),
 142                            self.part[1].pack(),
 143                            self.part[2].pack(),
 144                            self.part[3].pack(),
 145                            self.signature)
 146
 147     def __str__(self):
 148         ret = ""
 149         for i in range(4):
 150             ret += "Partition %d: %s\n" % (i, self.part[i])
 151         ret += "Signature: %s %s\n" % (hex(ord(self.signature[0])),
 152                                        hex(ord(self.signature[1])))
 153         return ret
 154
 155
 156 class Disklabel:
 157     """Represents an OpenBSD Disklabel"""
 158     format = "<IHH16s16sIIIIII8sIHHIII20sHH16sIHHII364s"
 159     """
 160     Offset  Length          Contents
 161     0       4               Magic
 162     4       2               Drive Type
 163     6       2               Subtype
 164     8       16              Type Name
 165     24      16              Pack Identifier
 166     32      4               Bytes per sector
 167     36      4               Data sectors per track
 168     40      4               Tracks per cilinder
 169     44      4               Data cylinders per unit
 170     48      4               Data sectors per cylynder
 171     52      4               Data sectors per unit
 172     56      8               Unique label identifier
 173     64      4               Alt cylinders per unit
 174     68      2               Start of useable region (high part)
 175     70      2               Size of usable region (high part)
 176     72      4               Start of useable region
 177     76      4               End of usable region
 178     80      4               Generic Flags
 179     84      5*4             Drive-type specific information
 180     104     2               Number of data sectors (high part)
 181     106     2               Version
 182     108     4*4             Reserved for future use
 183     124     4               Magic number
 184     128     2               Xor of data Inclu. partitions
 185     130     2               Number of partitions in following
 186     132     4               size of boot area at sn0, bytes
 187     136     4               Max size of fs superblock, bytes
 188     140     16*16           Partition Table
 189     """
 190
 191     class PartitionTable:
 192         """Reprepsents an OpenBSD Partition Table"""
 193         format = "<IIHHBBH"
 194         """
 195         Partition Entry:
 196         Offset  Length          Contents
 197         0       4               Number of sectors in the partition
 198         4       4               Starting sector
 199         8       2               Starting sector (high part)
 200         10      2               Number of sectors (high part)
 201         12      1               Filesystem type
 202         13      1               Filesystem Fragment per block
 203         14      2               FS cylinders per group
 204         """
 205
 206         Partition = namedtuple(
 207             'Partition', 'size, offset, offseth, sizeh, fstype, frag, cpg')
 208
 209         def __init__(self, ptable, pnumber):
 210             """Create a Partition Table instance"""
 211             self.part = []
 212
 213             size = struct.calcsize(self.format)
 214
 215             raw = cStringIO.StringIO(ptable)
 216             try:
 217                 for i in range(pnumber):
 218                     p = self.Partition(
 219                         *struct.unpack(self.format, raw.read(size)))
 220                     self.part.append(p)
 221             finally:
 222                 raw.close()
 223
 224         def __str__(self):
 225             """Print the Partition table"""
 226             val = ""
 227             for i in range(len(self.part)):
 228                 val = "%s%s\n" % (val, str(self.part[i]))
 229             return val
 230
 231         def pack(self):
 232             """Packs the partition table into a binary string."""
 233             ret = ""
 234             for i in range(len(self.part)):
 235                 ret += struct.pack(self.format,
 236                                    self.part[i].size,
 237                                    self.part[i].offset,
 238                                    self.part[i].offseth,
 239                                    self.part[i].sizeh,
 240                                    self.part[i].fstype,
 241                                    self.part[i].frag,
 242                                    self.part[i].cpg)
 243             return ret
 244
 245         def setpsize(self, i, size):
 246             """Set size for partition i"""
 247             tmp = self.part[i]
 248             self.part[i] = self.Partition(size & 0xffffffff, tmp.offset,
 249                                           tmp.offseth, size >> 32, tmp.fstype,
 250                                           tmp.frag, tmp.cpg)
 251
 252         def getpsize(self, i):
 253             return (self.part[i].sizeh << 32) + self.part[i].size
 254
 255         def setpoffset(self, i, offset):
 256             """Set  offset for partition i"""
 257             tmp = self.part[i]
 258             self.part[i] = self.Partition(tmp.size, offset & 0xffffffff,
 259                                           offset >> 32, tmp.sizeh, tmp.frag,
 260                                           tmp.cpg)
 261
 262         def getpoffset(self, i):
 263             return (self.part[i].offseth << 32) + self.part[i].offset
 264
 265     def __init__(self, disk):
 266         """Create a DiskLabel instance"""
 267         self.disk = disk
 268         self.part_num = None
 269
 270         with open(disk, "rb") as d:
 271             sector0 = d.read(BLOCKSIZE)
 272             self.mbr = MBR(sector0)
 273
 274             for i in range(4):
 275                 if self.mbr.part[i].type == 0xa6:  # OpenBSD type
 276                     self.part_num = i
 277                     break
 278
 279             assert self.part_num is not None, "No OpenBSD partition found"
 280
 281             d.seek(BLOCKSIZE * self.mbr.part[self.part_num].first_sector)
 282             part_sector0 = d.read(BLOCKSIZE)
 283             # The offset of the disklabel from the begining of the
 284             # partition is one sector
 285             part_sector1 = d.read(BLOCKSIZE)
 286
 287         (self.magic,
 288          self.dtype,
 289          self.subtype,
 290          self.typename,
 291          self.packname,
 292          self.secsize,
 293          self.nsectors,
 294          self.ntracks,
 295          self.ncylinders,
 296          self.secpercyl,
 297          self.secperunit,
 298          self.uid,
 299          self.acylinders,
 300          self.bstarth,
 301          self.bendh,
 302          self.bstart,
 303          self.bend,
 304          self.flags,
 305          self.drivedata,
 306          self.secperunith,
 307          self.version,
 308          self.spare,
 309          self.magic2,
 310          self.checksum,
 311          self.npartitions,
 312          self.bbsize,
 313          self.sbsize,
 314          ptable_raw) = struct.unpack(self.format, part_sector1)
 315
 316         assert self.magic == DISKMAGIC, "Disklabel is not valid"
 317
 318         self.ptable = self.PartitionTable(ptable_raw, self.npartitions)
 319
 320     def pack(self, checksum=None):
 321         return struct.pack(self.format,
 322                            self.magic,
 323                            self.dtype,
 324                            self.subtype,
 325                            self.typename,
 326                            self.packname,
 327                            self.secsize,
 328                            self.nsectors,
 329                            self.ntracks,
 330                            self.ncylinders,
 331                            self.secpercyl,
 332                            self.secperunit,
 333                            self.uid,
 334                            self.acylinders,
 335                            self.bstarth,
 336                            self.bendh,
 337                            self.bstart,
 338                            self.bend,
 339                            self.flags,
 340                            self.drivedata,
 341                            self.secperunith,
 342                            self.version,
 343                            self.spare,
 344                            self.magic2,
 345                            self.checksum if checksum is None else checksum,
 346                            self.npartitions,
 347                            self.bbsize,
 348                            self.sbsize,
 349                            self.ptable.pack() +
 350                            ((364 - self.npartitions*16) * '\x00'))
 351
 352     def compute_checksum(self):
 353         """Compute the checksum of the disklabel"""
 354
 355         raw = cStringIO.StringIO(self.pack(0))
 356         checksum = 0
 357         try:
 358             uint16 = raw.read(2)
 359             while uint16 != "":
 360                 checksum ^= struct.unpack('<H', uint16)[0]
 361                 uint16 = raw.read(2)
 362         finally:
 363             raw.close()
 364
 365         return checksum
 366
 367     def setdsize(self, dsize):
 368         """Set disk size"""
 369         self.secperunith = dsize >> 32
 370         self.secperunit = dsize & 0xffffffff
 371
 372     def getdsize(self):
 373         """Get disk size"""
 374         return (self.secperunith << 32) + self.secperunit
 375
 376     def setbstart(self, bstart):
 377         """Set start of useable region"""
 378         self.bstarth = bstart >> 32
 379         self.bstart = bstart & 0xffffffff
 380
 381     def getbstart(self):
 382         """Get start of usable region"""
 383         return (self.bstarth << 32) + self.bstart
 384
 385     def setbend(self, bend):
 386         """Set end of useable region"""
 387         self.bendh = bend >> 32
 388         self.bend = bend & 0xffffffff
 389
 390     def getbend(self):
 391         return (self.bendh << 32) + self.bend
 392
 393     def enlarge_disk(self, new_size):
 394         """Enlarge the size of the disk"""
 395
 396         assert new_size >= self.secperunit, \
 397             "New size cannot be smaller that %s" % self.secperunit
 398
 399         # Fix the disklabel
 400         self.setdsize(new_size)
 401         self.ncylinders = self.getdsize() // (self.nsectors * self.ntracks)
 402         self.setbend(self.ncylinders * self.nsectors * self.ntracks)
 403
 404         # Partition 'c' descriptes the entire disk
 405         self.ptable.setpsize(2, new_size)
 406
 407         # Fix the MBR table
 408         start = self.mbr.part[self.part_num].first_sector
 409         self.mbr.part[self.part_num].sector_count = self.getbend() - start
 410
 411         lba = self.getbend() - 1
 412         cylinder = lba // (self.ntracks * self.nsectors)
 413         header = (lba // self.nsectors) % self.ntracks
 414         sector = (lba % self.nsectors) + 1
 415         chs = MBR.Partition.pack_chs(cylinder, header, sector)
 416         self.mbr.part[self.part_num].end = chs
 417
 418         self.checksum = self.compute_checksum()
 419
 420     def write(self):
 421         """Write the disklabel back to the media"""
 422         with open(self.disk, 'rw+b') as d:
 423             d.write(self.mbr.pack())
 424
 425             d.seek((self.mbr.part[self.part_num].first_sector + 1) * BLOCKSIZE)
 426             d.write(self.pack())
 427
 428     def get_last_partition_id(self):
 429         """Returns the id of the last partition"""
 430         part = 0
 431         end = 0
 432         # Don't check partition 'c' which is the whole disk
 433         for i in filter(lambda x: x != 2, range(self.npartitions)):
 434             curr_end = self.ptable.getpsize(i) + self.ptable.getpoffset(i)
 435             if end < curr_end:
 436                 end = curr_end
 437                 part = i
 438
 439         assert end > 0, "No partition found"
 440
 441         return part
 442
 443     def enlarge_last_partition(self):
 444         """Enlarge the last partition to cover up all the free space"""
 445
 446         part_num = self.get_last_partition_id()
 447
 448         end = self.ptable.getpsize(part_num) + self.ptable.getpoffset(part_num)
 449
 450         assert end > 0, "No partition found"
 451
 452         if self.ptable.part[part_num].fstype == 1:  # Swap partition.
 453             #TODO: Maybe create a warning?
 454             return
 455
 456         if end > (self.getbend() - 1024):
 457             return
 458
 459         self.ptable.setpsize(
 460             part_num, self.getbend() - self.ptable.getpoffset(part_num) - 1024)
 461
 462         self.checksum = self.compute_checksum()
 463
 464     def __str__(self):
 465         """Print the Disklabel"""
 466         title1 = "Master Boot Record"
 467         title2 = "Disklabel"
 468
 469         return \
 470             "%s\n%s\n%s\n" % (title1, len(title1) * "=", str(self.mbr)) + \
 471             "%s\n%s\n" % (title2, len(title2) * "=") + \
 472             "Magic Number: 0x%x\n" % self.magic + \
 473             "Drive type: %d\n" % self.dtype + \
 474             "Subtype: %d\n" % self.subtype + \
 475             "Typename: %s\n" % self.typename + \
 476             "Pack Identifier: %s\n" % self.packname + \
 477             "Number of bytes per sector: %d\n" % self.secsize + \
 478             "Number of data sectors per track: %d\n" % self.nsectors + \
 479             "Number of tracks per cylinder: %d\n" % self.ntracks + \
 480             "Number of data cylinders per unit: %d\n" % self.ncylinders + \
 481             "Number of data sectors per cylinder: %d\n" % self.secpercyl + \
 482             "Number of data sectors per unit: %d\n" % self.secperunit + \
 483             "DUID: %s\n" % "".join(x.encode('hex') for x in self.uid) + \
 484             "Alt. cylinders per unit: %d\n" % self.acylinders + \
 485             "Start of useable region (high part): %d\n" % self.bstarth + \
 486             "Size of useable region (high part): %d\n" % self.bendh + \
 487             "Start of useable region: %d\n" % self.bstart + \
 488             "End of usable region: %d\n" % self.bend + \
 489             "Generic Flags: %r\n" % self.flags + \
 490             "Drive data: %r\n" % self.drivedata + \
 491             "Number of data sectors (high part): %d\n" % self.secperunith + \
 492             "Version: %d\n" % self.version + \
 493             "Reserved for future use: %r\n" % self.spare + \
 494             "The magic number again: 0x%x\n" % self.magic2 + \
 495             "Checksum: %d\n" % self.checksum + \
 496             "Number of partitions: %d\n" % self.npartitions + \
 497             "Size of boot aread at sn0: %d\n" % self.bbsize + \
 498             "Max size of fs superblock: %d\n" % self.sbsize + \
 499             "%s" % self.ptable
 500
 501
 502 if __name__ == '__main__':
 503
 504     usage = "Usage: %prog [options] <input_media>"
 505     parser = optparse.OptionParser(usage=usage)
 506
 507     parser.add_option("-l", "--list", action="store_true", dest="list",
 508                       default=False,
 509                       help="list the disklabel on the specified media")
 510     parser.add_option("--print-last", action="store_true", dest="last_part",
 511                       default=False,
 512                       help="print the label of the last partition")
 513     parser.add_option("--print-last-linux", action="store_true",
 514                       dest="last_linux", default=False,
 515                       help="print the linux number for the last partition")
 516     parser.add_option("--print-duid", action="store_true", dest="duid",
 517                       default=False,
 518                       help="print the disklabel unique identifier")
 519     parser.add_option("-d", "--enlarge-disk", type="int", dest="disk_size",
 520                       default=None, metavar="SIZE",
 521                       help="Enlarge the disk to this SIZE (in sectors)")
 522     parser.add_option(
 523         "-p", "--enlarge-partition", action="store_true",
 524         dest="enlarge_partition", default=False,
 525         help="Enlarge the last partition to cover up the free space")
 526
 527     options, args = parser.parse_args(sys.argv[1:])
 528
 529     if len(args) != 1:
 530         parser.error("Wrong number of arguments")
 531
 532     disklabel = Disklabel(args[0])
 533
 534     if options.list:
 535         print disklabel
 536         sys.exit(0)
 537
 538     if options.duid:
 539         print "%s" % "".join(x.encode('hex') for x in disklabel.uid)
 540         sys.exit(0)
 541
 542     if options.last_part:
 543         print "%c" % chr(ord('a') + disklabel.get_last_partition_id())
 544
 545     if options.last_linux:
 546         part_id = disklabel.get_last_partition_id()
 547         # The linux kernel does not assign a partition for label 'c' that
 548         # describes the whole disk
 549         print part_id + (4 if part_id > 2 else 5)
 550
 551     if options.disk_size is not None:
 552         disklabel.enlarge_disk(options.disk_size)
 553
 554     if options.enlarge_partition:
 555         disklabel.enlarge_last_partition()
 556
 557     disklabel.write()
 558
 559 sys.exit(0)
 560
 561 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :