Cleanup code & add comments in gpt module
[snf-image-creator] / image_creator / gpt.py
1 # Copyright 2012 GRNET S.A. All rights reserved.
2 #
3 # Redistribution and use in source and binary forms, with or
4 # without modification, are permitted provided that the following
5 # conditions are met:
6 #
7 #   1. Redistributions of source code must retain the above
8 #      copyright notice, this list of conditions and the following
9 #      disclaimer.
10 #
11 #   2. Redistributions in binary form must reproduce the above
12 #      copyright notice, this list of conditions and the following
13 #      disclaimer in the documentation and/or other materials
14 #      provided with the distribution.
15 #
16 # THIS SOFTWARE IS PROVIDED BY GRNET S.A. ``AS IS'' AND ANY EXPRESS
17 # OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
18 # WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
19 # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GRNET S.A OR
20 # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
21 # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
22 # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
23 # USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
24 # AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
25 # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
26 # ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
27 # POSSIBILITY OF SUCH DAMAGE.
28 #
29 # The views and conclusions contained in the software and
30 # documentation are those of the authors and should not be
31 # interpreted as representing official policies, either expressed
32 # or implied, of GRNET S.A.
33
34 import struct
35 import sys
36 import uuid
37 import binascii
38
39 BLOCKSIZE = 512
40
41
42 class MBR(object):
43     """Represents a Master Boot Record."""
44     class Partition(object):
45         format = "<B3sB3sLL"
46
47         def __init__(self, raw_part):
48             (
49                 self.status,
50                 self.start,
51                 self.type,
52                 self.end,
53                 self.first_sector,
54                 self.sector_count
55             ) = struct.unpack(self.format, raw_part)
56
57         def pack(self):
58             return struct.pack(self.format,
59                 self.status,
60                 self.start,
61                 self.type,
62                 self.end,
63                 self.first_sector,
64                 self.sector_count
65             )
66
67         @staticmethod
68         def size():
69             """Returns the size of an MBR partition entry"""
70             return struct.calcsize(MBR.Partition.format)
71
72         def __str__(self):
73             start = self.unpack_chs(self.start)
74             end = self.unpack_chs(self.end)
75             return "%d %s %d %s %d %d" % (self.status, start, self.type, end,
76                 self.first_sector, self.sector_count)
77
78         def unpack_chs(self, chs):
79             """Unpacks a CHS address string to a tuple."""
80
81             assert len(chs) == 3
82
83             head = struct.unpack('<B', chs[0])[0]
84             sector = struct.unpack('<B', chs[1])[0] & 0x3f
85             cylinder = (struct.unpack('<B', chs[1])[0] & 0xC0) << 2 | \
86                 struct.unpack('<B', chs[2])[0]
87
88             return (cylinder, head, sector)
89
90         def pack_chs(self, cylinder, head, sector):
91             """Packs a CHS tuple to an address string."""
92
93             assert 1 <= sector <= 63
94             assert 0 <= cylinder <= 1023
95             assert 0 <= head <= 255
96
97             byte0 = head
98             byte1 = (cylinder >> 2) & 0xC0 | sector
99             byte2 = cylinder & 0xff
100
101             return struct.pack('<BBB', byte0, byte1, byte2)
102
103     format = "<444s2x16s16s16s16s2s"
104     """
105     Offset  Length          Contents
106     0       440(max. 446)   code area
107     440     2(optional)     disk signature
108     444     2               Usually nulls
109     446     16              Partition 0
110     462     16              Partition 1
111     478     16              Partition 2
112     494     16              Partition 3
113     510     2               MBR signature
114     """
115     def __init__(self, block):
116         raw_part = {}
117         self.code_area, \
118         raw_part[0], \
119         raw_part[1], \
120         raw_part[2], \
121         raw_part[3], \
122         self.signature = struct.unpack(self.format, block)
123
124         self.part = {}
125         for i in range(4):
126             self.part[i] = self.Partition(raw_part[i])
127
128     @staticmethod
129     def size():
130         """Returns the size of a Master Boot Record."""
131         return struct.calcsize(MBR.format)
132
133     def pack(self):
134         """Packs an MBR to a binary string."""
135         return struct.pack(self.format,
136             self.code_area,
137             self.part[0].pack(),
138             self.part[1].pack(),
139             self.part[2].pack(),
140             self.part[3].pack(),
141             self.signature
142         )
143
144     def __str__(self):
145         ret = ""
146         for i in range(4):
147             ret += "Partition %d: %s\n" % (i, self.part[i])
148         ret += "Signature: %s %s\n" % (
149                     hex(ord(self.signature[0])), hex(ord(self.signature[1])))
150         return ret
151
152
153 class GPTPartitionTable(object):
154     """Represents a GUID Partition Table."""
155     class GPTHeader(object):
156         """Represents a GPT Header of a GUID Partition Table."""
157         format = "<8s4sII4xQQQQ16sQIII"
158         """
159         Offset  Length          Contents
160         0       8 bytes         Signature
161         8       4 bytes         Revision
162         12      4 bytes         Header size in little endian
163         16      4 bytes         CRC32 of header
164         20      4 bytes         Reserved; must be zero
165         24      8 bytes         Current LBA
166         32      8 bytes         Backup LBA
167         40      8 bytes         First usable LBA for partitions
168         48      8 bytes         Last usable LBA
169         56      16 bytes        Disk GUID
170         72      8 bytes         Partition entries starting LBA
171         80      4 bytes         Number of partition entries
172         84      4 bytes         Size of a partition entry
173         88      4 bytes         CRC32 of partition array
174         92      *               Reserved; must be zeroes
175         LBA    size            Total
176         """
177
178         def __init__(self, block):
179             self.signature, \
180             self.revision, \
181             self.hdr_size, \
182             self.header_crc32, \
183             self.current_lba, \
184             self.backup_lba, \
185             self.first_usable_lba, \
186             self.last_usable_lba, \
187             self.uuid, \
188             self.part_entry_start, \
189             self.part_count, \
190             self.part_entry_size, \
191             self.part_crc32 = struct.unpack(self.format, block)
192
193         def pack(self):
194             """Packs a GPT Header to a binary string."""
195             return struct.pack(self.format,
196                 self.signature, \
197                 self.revision, \
198                 self.hdr_size, \
199                 self.header_crc32, \
200                 self.current_lba, \
201                 self.backup_lba, \
202                 self.first_usable_lba, \
203                 self.last_usable_lba, \
204                 self.uuid, \
205                 self.part_entry_start, \
206                 self.part_count, \
207                 self.part_entry_size, \
208                 self.part_crc32
209             )
210
211         @staticmethod
212         def size():
213             """Returns the size of a GPT Header."""
214             return struct.calcsize(GPTPartitionTable.GPTHeader.format)
215
216         def __str__(self):
217             return \
218             "Signature: %s\n" % self.signature + \
219             "Revision: %r\n" % self.revision + \
220             "Header Size: %d\n" % self.hdr_size + \
221             "CRC32: %d\n" % self.header_crc32 + \
222             "Current LBA: %d\n" % self.current_lba + \
223             "Backup LBA: %d\n" % self.backup_lba + \
224             "First Usable LBA: %d\n" % self.first_usable_lba + \
225             "Last Usable LBA: %d\n" % self.last_usable_lba + \
226             "Disk GUID: %s\n" % uuid.UUID(bytes=self.uuid) + \
227             "Partition entries starting LBA: %d\n" % self.part_entry_start + \
228             "Number of Partition entries: %d\n" % self.part_count + \
229             "Size of a partition entry: %d\n" % self.part_entry_size + \
230             "CRC32 of partition array: %s\n" % self.part_crc32
231
232     def __init__(self, disk):
233         self.disk = disk
234         with open(disk, "rb") as d:
235             #MBR (Logical block address 0)
236             lba0 = d.read(BLOCKSIZE)
237             self.mbr = MBR(lba0)
238
239             # Primary GPT Header (LBA 1)
240             raw_header = d.read(self.GPTHeader.size())
241             self.primary = self.GPTHeader(raw_header)
242
243             # Partition entries (LBA 2...34)
244             d.seek(self.primary.part_entry_start * BLOCKSIZE)
245             entries_size = self.primary.part_count * \
246                                                 self.primary.part_entry_size
247             self.part_entries = d.read(entries_size)
248
249             # Secondary GPT Header (LBA -1)
250             d.seek(self.primary.backup_lba * BLOCKSIZE)
251             raw_header = d.read(self.GPTHeader.size())
252             self.secondary = self.GPTHeader(raw_header)
253
254     def size(self):
255         """Returns the payload size of GPT partitioned device."""
256         return (self.primary.backup_lba + 1) * BLOCKSIZE
257
258     def shrink(self, size):
259         """Move the secondary GPT Header entries to the address specified by
260         size parameter.
261         """
262         if size == self.size():
263             return size
264
265         assert size < self.size()
266
267         # new_size = size + Partition Entries + Secondary GPT Header
268         new_size = size + len(self.part_entries) + BLOCKSIZE
269         new_size = ((new_size + 4095) // 4096) * 4096  # align to 4K
270         lba_count = new_size // BLOCKSIZE
271
272         # Correct MBR
273         #TODO: Check if the partition tables is hybrid
274         self.mbr.part[0].sector_count = (new_size // BLOCKSIZE) - 1
275
276         # Fix Primary header
277         self.primary.header_crc32 = 0
278         self.primary.backup_lba = lba_count - 1  # LBA-1
279         self.primary.last_usable_lba = lba_count - 34  # LBA-34
280         self.primary.header_crc32 = \
281                             binascii.crc32(self.primary.pack()) & 0xffffffff
282
283         # Fix Secondary header
284         self.secondary.header_crc32 = 0
285         self.secondary.current_lba = self.primary.backup_lba
286         self.secondary.last_usable_lba = lba_count - 34  # LBA-34
287         self.secondary.part_entry_start = lba_count - 33  # LBA-33
288         self.secondary.header_crc32 = \
289                             binascii.crc32(self.secondary.pack()) & 0xffffffff
290
291         # Copy the new partition table back to the device
292         with open(self.disk, "wb") as d:
293             d.write(self.mbr.pack())
294             d.write(self.primary.pack())
295             d.write('\x00' * (BLOCKSIZE - self.primary.size()))
296             d.seek(self.secondary.part_entry_start * BLOCKSIZE)
297             d.write(self.part_entries)
298             d.seek(self.primary.backup_lba * BLOCKSIZE)
299             d.write(self.secondary.pack())
300             d.write('\x00' * (BLOCKSIZE - self.secondary.size()))
301
302         return new_size
303
304 if __name__ == '__main__':
305     ptable = GPTPartitionTable(sys.argv[1])
306
307     print "MBR:\n%s" % ptable.mbr
308     print "Primary partition table:\n%s" % ptable.primary
309     print "Secondary partition table:\n%s" % ptable.secondary
310
311 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :