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