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