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