Revision be5f0160

b/snf-image-helper/disklabel.py
37 37
BBSIZE = 8192  # size of boot area with label
38 38
SBSIZE = 8192  # max size of fs superblock
39 39

  
40
DISKMAGIC = 0x82564557
41

  
40 42

  
41 43
class MBR(object):
42 44
    """Represents a Master Boot Record."""
......
161 163
        return "%s\n%s\n%s\n" % (title, len(title) * "=", ret)
162 164

  
163 165

  
164
class Disklabel:
165
    """Represents an BSD Disklabel"""
166
class PartitionTableBase(object):
167
    """Base Class for disklabel partition tables"""
168
    format = ""
166 169

  
167
    def __init__(self, disk):
168
        """Create a DiskLabel instance"""
169
        self.disk = disk
170
    Partition = namedtuple('Partition', '')
171

  
172
    def __init__(self, ptable, pnumber):
173
        """Create a Partition Table instance"""
174
        self.part = []
175

  
176
        size = struct.calcsize(self.format)
177

  
178
        raw = cStringIO.StringIO(ptable)
179
        try:
180
            for i in range(pnumber):
181
                p = self.Partition(
182
                    *struct.unpack(self.format, raw.read(size)))
183
                self.part.append(p)
184
        finally:
185
            raw.close()
186

  
187
    def __str__(self):
188
        """Print the Partition table"""
189
        val = ""
190
        for i in range(len(self.part)):
191
            val += "%c: %s\n" % (chr(ord('a') + i), str(self.part[i]))
192
        return val
193

  
194
    def pack(self):
195
        """Packs the partition table into a binary string."""
196
        ret = ""
197
        for i in range(len(self.part)):
198
            ret += struct.pack(self.format, *self.part[i])
199
        return ret
200

  
201

  
202
class Disk(object):
203
    """Represents an BSD Disk"""
204

  
205
    def __init__(self, device):
206
        """Create a Disk instance"""
207
        self.device = device
170 208
        self.part_num = None
171 209
        self.disklabel = None
172 210

  
173
        with open(disk, "rb") as d:
211
        with open(device, "rb") as d:
174 212
            sector0 = d.read(BLOCKSIZE)
175 213
            self.mbr = MBR(sector0)
176 214

  
......
190 228
        assert self.disklabel is not None, "No *BSD partition found"
191 229

  
192 230
    def write(self):
193
        """Write the disklabel back to the media"""
194
        with open(self.disk, 'rw+b') as d:
231
        """Write the changes back to the media"""
232
        with open(self.device, 'rw+b') as d:
195 233
            d.write(self.mbr.pack())
196 234

  
197 235
            d.seek(self.mbr.part[self.part_num].first_sector * BLOCKSIZE)
......
200 238
    def __str__(self):
201 239
        return str(self.mbr) + str(self.disklabel)
202 240

  
203
    def enlarge_disk(self, new_size):
204
        """Enlarge the size of the disk and return the last usable sector"""
241
    def enlarge(self, new_size):
242
        """Enlarge the disk and return the last usable sector"""
205 243

  
206 244
        # Fix the disklabel
207
        end = self.disklabel.enlarge_disk(new_size)
245
        end = self.disklabel.enlarge(new_size)
208 246

  
209 247
        # Fix the MBR
210 248
        start = self.mbr.part[self.part_num].first_sector
......
220 258
        self.disklabel.enlarge_last_partition()
221 259

  
222 260

  
223
class BSD_Disklabel:
224
    pass
261
class BSD_Disklabel(object):
262
    """Represents an BSD Disklabel"""
225 263

  
264
    class PartitionTable(PartitionTableBase):
265
        """Represents a BSD Partition Table"""
266
        format = "<IIIBBH"
267
        """
268
        Partition Entry:
269
        Offset  Length          Contents
270
        0       4               Number of sectors in partition
271
        4       4               Starting sector
272
        8       4               Filesystem basic fragment size
273
        12      1               Filesystem type
274
        13      1               Filesystem fragments per block
275
        14      2               Filesystem cylinders per group
276
        """
226 277

  
227
class OpenBSD_Disklabel:
228
    """Represents an OpenBSD Disklabel"""
229
    format = "<IHH16s16sIIIIII8sIHHIII20sHH16sIHHII364s"
278
        Partition = namedtuple(
279
            'Partition', 'size, offset, fsize, fstype, frag, cpg')
280

  
281
    format = "<IHH16s16sIIIIIIHHIHHHHIII20s20sIHHII64s"
230 282
    """
231 283
    Offset  Length          Contents
232 284
    0       4               Magic
......
236 288
    24      16              Pack Identifier
237 289
    32      4               Bytes per sector
238 290
    36      4               Data sectors per track
239
    40      4               Tracks per cilinder
291
    40      4               Tracks per cylinder
240 292
    44      4               Data cylinders per unit
241
    48      4               Data sectors per cylynder
293
    48      4               Data sectors per cylinder
242 294
    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
295
    56      2               Spare sectors per track
296
    58      2               Spare sectors per cylinder
297
    60      4               Alternative cylinders per unit
298
    64      2               Rotation Speed
299
    66      2               Hardware sector interleave
300
    68      2               Sector 0 skew, per track
301
    70      2               Sector 0 skew, per cylinder
302
    72      4               Head switch time
303
    76      4               Track-to-track seek
249 304
    80      4               Generic Flags
250 305
    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
306
    104     5*4             Reserved for future use
307
    124     4               Magic Number
308
    128     2               Xor of data including partitions
309
    130     2               Number of partitions following
257 310
    132     4               size of boot area at sn0, bytes
258 311
    136     4               Max size of fs superblock, bytes
259 312
    140     16*16           Partition Table
260 313
    """
261 314

  
262
    class PartitionTable:
315

  
316
class OpenBSD_Disklabel(object):
317
    """Represents an OpenBSD Disklabel"""
318

  
319
    class PartitionTable(PartitionTableBase):
263 320
        """Reprepsents an OpenBSD Partition Table"""
264 321
        format = "<IIHHBBH"
265 322
        """
......
277 334
        Partition = namedtuple(
278 335
            'Partition', 'size, offset, offseth, sizeh, fstype, frag, cpg')
279 336

  
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 337
        def setpsize(self, i, size):
317 338
            """Set size for partition i"""
318 339
            tmp = self.part[i]
......
335 356
            """Get offset for partition i"""
336 357
            return (self.part[i].offseth << 32) + self.part[i].offset
337 358

  
338
    DISKMAGIC = 0x82564557
339

  
359
    format = "<IHH16s16sIIIIII8sIHHIII20sHH16sIHHII364s"
360
    """
361
    Offset  Length          Contents
362
    0       4               Magic
363
    4       2               Drive Type
364
    6       2               Subtype
365
    8       16              Type Name
366
    24      16              Pack Identifier
367
    32      4               Bytes per sector
368
    36      4               Data sectors per track
369
    40      4               Tracks per cylinder
370
    44      4               Data cylinders per unit
371
    48      4               Data sectors per cylinder
372
    52      4               Data sectors per unit
373
    56      8               Unique label identifier
374
    64      4               Alt cylinders per unit
375
    68      2               Start of useable region (high part)
376
    70      2               Size of usable region (high part)
377
    72      4               Start of useable region
378
    76      4               End of usable region
379
    80      4               Generic Flags
380
    84      5*4             Drive-type specific information
381
    104     2               Number of data sectors (high part)
382
    106     2               Version
383
    108     4*4             Reserved for future use
384
    124     4               Magic number
385
    128     2               Xor of data including partitions
386
    130     2               Number of partitions in following
387
    132     4               size of boot area at sn0, bytes
388
    136     4               Max size of fs superblock, bytes
389
    140     16*16           Partition Table
390
    """
340 391
    def __init__(self, device):
341 392
        """Create a DiskLabel instance"""
342 393

  
......
374 425
         self.sbsize,
375 426
         ptable_raw) = struct.unpack(self.format, sector1)
376 427

  
377
        assert self.magic == self.DISKMAGIC, "Disklabel is not valid"
428
        assert self.magic == DISKMAGIC, "Disklabel is not valid"
378 429

  
379 430
        self.ptable = self.PartitionTable(ptable_raw, self.npartitions)
380 431

  
......
452 503
        """Get size of usable region"""
453 504
        return (self.bendh << 32) + self.bend
454 505

  
455
    def enlarge_disk(self, new_size):
506
    def enlarge(self, new_size):
456 507
        """Enlarge the size of the disk and return the last usable sector"""
457 508

  
458 509
        assert new_size >= self.getdsize(), \
......
581 632
    if len(args) != 1:
582 633
        parser.error("Wrong number of arguments")
583 634

  
584
    disklabel = Disklabel(args[0])
635
    disk = Disk(args[0])
585 636

  
586 637
    if options.list:
587
        print disklabel
638
        print disk
588 639
        sys.exit(0)
589 640

  
590 641
    if options.duid:
591
        print "%s" % "".join(x.encode('hex') for x in disklabel.uid)
642
        print "%s" % "".join(x.encode('hex') for x in disk.uid)
592 643
        sys.exit(0)
593 644

  
594 645
    if options.last_part:
595
        print "%c" % chr(ord('a') + disklabel.get_last_partition_id())
646
        print "%c" % chr(ord('a') + disk.get_last_partition_id())
596 647

  
597 648
    if options.disk_size is not None:
598
        disklabel.enlarge_disk(options.disk_size)
649
        disk.enlarge(options.disk_size)
599 650

  
600 651
    if options.enlarge_partition:
601
        disklabel.enlarge_last_partition()
652
        disk.enlarge_last_partition()
602 653

  
603
    disklabel.write()
654
    disk.write()
604 655

  
605 656
sys.exit(0)
606 657

  

Also available in: Unified diff