From 3b2674b32c43497af679f2ab5c1bb946b80309bc Mon Sep 17 00:00:00 2001
From: Nikos Skalkotos <skalkoto@grnet.gr>
Date: Tue, 4 Mar 2014 09:03:07 +0200
Subject: [PATCH] disklabel.py code cleanup

---
 snf-image-helper/disklabel.py |  165 +++++++++++++++++++++--------------------
 1 file changed, 86 insertions(+), 79 deletions(-)

diff --git a/snf-image-helper/disklabel.py b/snf-image-helper/disklabel.py
index 75d4b5c..fda1964 100755
--- a/snf-image-helper/disklabel.py
+++ b/snf-image-helper/disklabel.py
@@ -112,36 +112,28 @@ class MBR(object):
 
             return struct.pack('<BBB', byte0, byte1, byte2)
 
-    format = "<444s2x16s16s16s16s2s"
-    """
-    Offset  Length          Contents
-    0       440(max. 446)   code area
-    440     2(optional)     disk signature
-    444     2               Usually nulls
-    446     16              Partition 0
-    462     16              Partition 1
-    478     16              Partition 2
-    494     16              Partition 3
-    510     2               MBR signature
-    """
     def __init__(self, block):
         """Create an MBR instance"""
-        raw_part = {}
-        (self.code_area,
-         raw_part[0],
-         raw_part[1],
-         raw_part[2],
-         raw_part[3],
-         self.signature) = struct.unpack(self.format, block)
+
+        self.format = "<444s2x16s16s16s16s2s"
+        raw_part = {}     # Offset  Length          Contents
+        (self.code_area,  # 0       440(max. 446)   code area
+                          # 440     2(optional)     disk signature
+                          # 444     2               Usually nulls
+         raw_part[0],     # 446     16              Partition 0
+         raw_part[1],     # 462     16              Partition 1
+         raw_part[2],     # 478     16              Partition 2
+         raw_part[3],     # 494     16              Partition 3
+         self.signature   # 510     2               MBR signature
+         ) = struct.unpack(self.format, block)
 
         self.part = {}
         for i in range(4):
             self.part[i] = self.Partition(raw_part[i])
 
-    @staticmethod
-    def size():
+    def size(self):
         """Return the size of a Master Boot Record."""
-        return struct.calcsize(MBR.format)
+        return struct.calcsize(self.format)
 
     def pack(self):
         """Pack an MBR to a binary string."""
@@ -154,6 +146,7 @@ class MBR(object):
                            self.signature)
 
     def __str__(self):
+        """Print the MBR"""
         ret = ""
         for i in range(4):
             ret += "Partition %d: %s\n" % (i, self.part[i])
@@ -164,7 +157,7 @@ class MBR(object):
 
 
 class Disk(object):
-    """Represents an BSD Disk"""
+    """Represents a BSD Disk"""
 
     def __init__(self, device):
         """Create a Disk instance"""
@@ -200,6 +193,7 @@ class Disk(object):
             self.disklabel.write_to(d)
 
     def __str__(self):
+        """Print the partitioning info of the Disk"""
         return str(self.mbr) + str(self.disklabel)
 
     def enlarge(self, new_size):
@@ -230,7 +224,7 @@ class Disk(object):
         return self.disklabel.get_last_partition_id()
 
     def get_duid(self):
-        """Get the Disk UID (work only on OpenBSD)"""
+        """Get the Disklabel Unique Identifier (works only for OpenBSD)"""
         if 'uid' in self.disklabel.field:
             return self.disklabel.field['uid']
 
@@ -242,13 +236,28 @@ class DisklabelBase(object):
 
     def __init__(self, device):
         """Create a Disklabel instance"""
+        pass
+
+    @property
+    def format(self):
+        """Fields format string"""
+        raise NotImplementedError
+
+    @property
+    def field(self):
+        """Diskalabel Fields data structure"""
+        raise NotImplementedError
+
+    @property
+    def ptable(self):
+        """Partition Table data structure"""
         raise NotImplementedError
 
     def pack(self, checksum=None):
         """Return a binary copy of the Disklabel block"""
 
         out = OrderedDict()
-        for k, v in self.field.items():
+        for k, v in self.field.iteritems():
             out[k] = v
 
         if checksum is not None:
@@ -313,24 +322,24 @@ class PartitionTableBase(object):
 
         raw = cStringIO.StringIO(ptable)
         try:
-            for i in range(pnumber):
-                p = self.Partition(
-                    *struct.unpack(self.format, raw.read(size)))
-                self.part.append(p)
+            for _ in xrange(pnumber):
+                self.part.append(
+                    self.Partition(*struct.unpack(self.format, raw.read(size)))
+                    )
         finally:
             raw.close()
 
     def __str__(self):
         """Print the Partition table"""
         val = ""
-        for i in range(len(self.part)):
+        for i in xrange(len(self.part)):
             val += "%c: %s\n" % (chr(ord('a') + i), str(self.part[i]))
         return val
 
     def pack(self):
         """Packs the partition table into a binary string."""
         ret = ""
-        for i in range(len(self.part)):
+        for i in xrange(len(self.part)):
             ret += struct.pack(self.format, *self.part[i])
         return ret + ((364 - len(self.part) * 16) * '\x00')
 
@@ -340,31 +349,27 @@ class BSD_Disklabel(DisklabelBase):
 
     class PartitionTable(PartitionTableBase):
         """Represents a BSD Partition Table"""
-        format = "<IIIBBH"
-        """
-        Partition Entry:
-        Offset  Length          Contents
-        0       4               Number of sectors in partition
-        4       4               Starting sector
-        8       4               Filesystem basic fragment size
-        12      1               Filesystem type
-        13      1               Filesystem fragments per block
-        14      2               Filesystem cylinders per group
-        """
 
+        format = "<IIIBBH"
         Partition = namedtuple(
-            'Partition', 'size, offset, fsize, fstype, frag, cpg')
-
-    format = "<IHH16s16sIIIIIIHHIHHHHIII20s20sIHHII364s"
+            'Partition',  # Offset  Length Contents
+            ['size',      # 0       4      Number of sectors in partition
+             'offset',    # 4       4      Starting sector
+             'fsize',     # 8       4      Filesystem basic fragment size
+             'fstype',    # 12      1      Filesystem type
+             'frag',      # 13      1      Filesystem fragments per block
+             'cpg'        # 14      2      Filesystem cylinders per group
+             ])
 
     def __init__(self, device):
         """Create a BSD DiskLabel instance"""
+        super(BSD_Disklabel, self).__init__(device)
 
+        # Disklabel starts at offset one
         device.seek(BLOCKSIZE, os.SEEK_CUR)
-        # The offset of the disklabel from the beginning of the partition is
-        # one sector
         sector1 = device.read(BLOCKSIZE)
 
+        self.format = "<IHH16s16sIIIIIIHHIHHHHIII20s20sIHHII364s"
         d_ = OrderedDict()      # Off  Len    Content
         (d_["magic"],           # 0    4      Magic
          d_["dtype"],           # 4    2      Drive Type
@@ -405,6 +410,8 @@ class BSD_Disklabel(DisklabelBase):
         """Print the Disklabel"""
 
         title = "Disklabel"
+
+        # Those fields may contain null bytes
         typename = self.field['typename'].strip('\x00').strip()
         packname = self.field['packname'].strip('\x00').strip()
 
@@ -447,19 +454,16 @@ class OpenBSD_Disklabel(DisklabelBase):
     class PartitionTable(PartitionTableBase):
         """Reprepsents an OpenBSD Partition Table"""
         format = "<IIHHBBH"
-        """
-        Partition Entry:
-        Offset  Length          Contents
-        0       4               Number of sectors in the partition
-        4       4               Starting sector
-        8       2               Starting sector (high part)
-        10      2               Number of sectors (high part)
-        12      1               Filesystem type
-        13      1               Filesystem Fragment per block
-        14      2               FS cylinders per group
-        """
         Partition = namedtuple(
-            'Partition', 'size, offset, offseth, sizeh, fstype, frag, cpg')
+            'Partition',  # Offset  Length Contents
+            ['size',      # 0       4      Number of sectors in the partition
+             'offset',    # 4       4      Starting sector
+             'offseth',   # 8       2      Starting sector (high part)
+             'sizeh',     # 10      2      Number of sectors (high part)
+             'fstype',    # 12      1      Filesystem type
+             'frag',      # 13      1      Filesystem Fragments per block
+             'cpg'        # 14      2      FS cylinders per group
+             ])
 
         def setpsize(self, i, size):
             """Set size for partition i"""
@@ -483,16 +487,15 @@ class OpenBSD_Disklabel(DisklabelBase):
             """Get offset for partition i"""
             return (self.part[i].offseth << 32) + self.part[i].offset
 
-    format = "<IHH16s16sIIIIII8sIHHIII20sHH16sIHHII364s"
-
     def __init__(self, device):
         """Create a DiskLabel instance"""
+        super(OpenBSD_Disklabel, self).__init__(device)
 
+        # Disklabel starts at offset one
         device.seek(BLOCKSIZE, os.SEEK_CUR)
-        # The offset of the disklabel from the beginning of the partition is
-        # one sector
         sector1 = device.read(BLOCKSIZE)
 
+        self.format = "<IHH16s16sIIIIII8sIHHIII20sHH16sIHHII364s"
         d_ = OrderedDict()   # Off  Len    Content
         (d_["magic"],        # 0    4      Magic
          d_["dtype"],        # 4    2      Drive Type
@@ -548,18 +551,18 @@ class OpenBSD_Disklabel(DisklabelBase):
         """Get start of usable region"""
         return (self.field['bstarth'] << 32) + self.field['bstart']
 
-    bstart = property(getbstart, setbstart, None, "usable region start")
+    bstart = property(getbstart, setbstart, None, "start of usable region")
 
     def setbend(self, bend):
-        """Set size of useable region"""
+        """Set end of useable region"""
         self.field['bendh'] = bend >> 32
         self.field['bend'] = bend & 0xffffffff
 
     def getbend(self):
-        """Get size of usable region"""
+        """Get end of usable region"""
         return (self.field['bendh'] << 32) + self.field['bend']
 
-    bend = property(getbend, setbend, None, "usable region size")
+    bend = property(getbend, setbend, None, "end of usable region")
 
     def enlarge(self, new_size):
         """Enlarge the disk and return the last usable sector"""
@@ -580,9 +583,7 @@ class OpenBSD_Disklabel(DisklabelBase):
         # Update the checksum
         self.field['checksum'] = self.compute_checksum()
 
-        # bend is the size and not the end of the usable region. I named it
-        # like this because this is how it is named in OpenBSD. To get the last
-        # usable sector you need to reduce this value by one.
+        # The last usable sector is the end of the usable region minus one
         return self.bend - 1
 
     def get_last_partition_id(self):
@@ -624,10 +625,13 @@ class OpenBSD_Disklabel(DisklabelBase):
     def __str__(self):
         """Print the Disklabel"""
 
-        title = "Disklabel"
+        # Those values may contain null bytes
         typename = self.field['typename'].strip('\x00').strip()
         packname = self.field['packname'].strip('\x00').strip()
+
         duid = "".join(x.encode('hex') for x in self.field['uid'])
+
+        title = "Disklabel"
         return \
             "%s\n%s\n" % (title, len(title) * "=") + \
             "Magic Number: 0x%(magic)x\n" \
@@ -660,8 +664,8 @@ class OpenBSD_Disklabel(DisklabelBase):
             "%s" % self.ptable
 
 
-if __name__ == '__main__':
-
+def main():
+    """Main entry point"""
     usage = "Usage: %prog [options] <input_media>"
     parser = optparse.OptionParser(usage=usage)
 
@@ -671,9 +675,9 @@ if __name__ == '__main__':
     parser.add_option("--get-last-partition", action="store_true",
                       dest="last_part", default=False,
                       help="print the label of the last partition")
-    parser.add_option("--get-duid", action="store_true", dest="duid",
-                      default=False,
-                      help="print the disklabel unique identifier")
+    parser.add_option(
+        "--get-duid", action="store_true", dest="duid", default=False,
+        help="print the Disklabel Unique Identifier (OpenBSD only)")
     parser.add_option("-d", "--enlarge-disk", type="int", dest="disk_size",
                       default=None, metavar="SIZE",
                       help="Enlarge the disk to this SIZE (in sectors)")
@@ -691,11 +695,11 @@ if __name__ == '__main__':
 
     if options.list:
         print disk
-        sys.exit(0)
+        return 0
 
     if options.duid:
         print "%s" % "".join(x.encode('hex') for x in disk.get_duid())
-        sys.exit(0)
+        return 0
 
     if options.last_part:
         print "%c" % chr(ord('a') + disk.get_last_partition_id())
@@ -707,7 +711,10 @@ if __name__ == '__main__':
         disk.enlarge_last_partition()
 
     disk.write()
+    return 0
+
 
-sys.exit(0)
+if __name__ == '__main__':
+    sys.exit(main())
 
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :
-- 
1.7.10.4