Bump version to 0.2.2
[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         @staticmethod
67         def size():
68             """Returns the size of an MBR partition entry"""
69             return struct.calcsize(MBR.Partition.format)
70
71         def __str__(self):
72             start = self.unpack_chs(self.start)
73             end = self.unpack_chs(self.end)
74             return "%d %s %d %s %d %d" % (self.status, start, self.type, end,
75                                           self.first_sector, self.sector_count)
76
77         def unpack_chs(self, chs):
78             """Unpacks a CHS address string to a tuple."""
79
80             assert len(chs) == 3
81
82             head = struct.unpack('<B', chs[0])[0]
83             sector = struct.unpack('<B', chs[1])[0] & 0x3f
84             cylinder = (struct.unpack('<B', chs[1])[0] & 0xC0) << 2 | \
85                 struct.unpack('<B', chs[2])[0]
86
87             return (cylinder, head, sector)
88
89         def pack_chs(self, cylinder, head, sector):
90             """Packs a CHS tuple to an address string."""
91
92             assert 1 <= sector <= 63
93             assert 0 <= cylinder <= 1023
94             assert 0 <= head <= 255
95
96             byte0 = head
97             byte1 = (cylinder >> 2) & 0xC0 | sector
98             byte2 = cylinder & 0xff
99
100             return struct.pack('<BBB', byte0, byte1, byte2)
101
102     format = "<444s2x16s16s16s16s2s"
103     """
104     Offset  Length          Contents
105     0       440(max. 446)   code area
106     440     2(optional)     disk signature
107     444     2               Usually nulls
108     446     16              Partition 0
109     462     16              Partition 1
110     478     16              Partition 2
111     494     16              Partition 3
112     510     2               MBR signature
113     """
114     def __init__(self, block):
115         raw_part = {}
116         (self.code_area,
117          raw_part[0],
118          raw_part[1],
119          raw_part[2],
120          raw_part[3],
121          self.signature) = struct.unpack(self.format, block)
122
123         self.part = {}
124         for i in range(4):
125             self.part[i] = self.Partition(raw_part[i])
126
127     @staticmethod
128     def size():
129         """Returns the size of a Master Boot Record."""
130         return struct.calcsize(MBR.format)
131
132     def pack(self):
133         """Packs an MBR to a binary string."""
134         return struct.pack(self.format,
135                            self.code_area,
136                            self.part[0].pack(),
137                            self.part[1].pack(),
138                            self.part[2].pack(),
139                            self.part[3].pack(),
140                            self.signature)
141
142     def __str__(self):
143         ret = ""
144         for i in range(4):
145             ret += "Partition %d: %s\n" % (i, self.part[i])
146         ret += "Signature: %s %s\n" % (hex(ord(self.signature[0])),
147                                        hex(ord(self.signature[1])))
148         return ret
149
150
151 class GPTPartitionTable(object):
152     """Represents a GUID Partition Table."""
153     class GPTHeader(object):
154         """Represents a GPT Header of a GUID Partition Table."""
155         format = "<8s4sII4xQQQQ16sQIII"
156         """
157         Offset  Length          Contents
158         0       8 bytes         Signature
159         8       4 bytes         Revision
160         12      4 bytes         Header size in little endian
161         16      4 bytes         CRC32 of header
162         20      4 bytes         Reserved; must be zero
163         24      8 bytes         Current LBA
164         32      8 bytes         Backup LBA
165         40      8 bytes         First usable LBA for partitions
166         48      8 bytes         Last usable LBA
167         56      16 bytes        Disk GUID
168         72      8 bytes         Partition entries starting LBA
169         80      4 bytes         Number of partition entries
170         84      4 bytes         Size of a partition entry
171         88      4 bytes         CRC32 of partition array
172         92      *               Reserved; must be zeroes
173         LBA    size            Total
174         """
175
176         def __init__(self, block):
177             (self.signature,
178              self.revision,
179              self.hdr_size,
180              self.header_crc32,
181              self.current_lba,
182              self.backup_lba,
183              self.first_usable_lba,
184              self.last_usable_lba,
185              self.uuid,
186              self.part_entry_start,
187              self.part_count,
188              self.part_entry_size,
189              self.part_crc32) = struct.unpack(self.format, block)
190
191         def pack(self):
192             """Packs a GPT Header to a binary string."""
193             return struct.pack(self.format,
194                                self.signature,
195                                self.revision,
196                                self.hdr_size,
197                                self.header_crc32,
198                                self.current_lba,
199                                self.backup_lba,
200                                self.first_usable_lba,
201                                self.last_usable_lba,
202                                self.uuid,
203                                self.part_entry_start,
204                                self.part_count,
205                                self.part_entry_size,
206                                self.part_crc32)
207
208         @staticmethod
209         def size():
210             """Returns the size of a GPT Header."""
211             return struct.calcsize(GPTPartitionTable.GPTHeader.format)
212
213         def __str__(self):
214             return "Signature: %s\n" % self.signature + \
215                    "Revision: %r\n" % self.revision + \
216                    "Header Size: %d\n" % self.hdr_size + \
217                    "CRC32: %d\n" % self.header_crc32 + \
218                    "Current LBA: %d\n" % self.current_lba + \
219                    "Backup LBA: %d\n" % self.backup_lba + \
220                    "First Usable LBA: %d\n" % self.first_usable_lba + \
221                    "Last Usable LBA: %d\n" % self.last_usable_lba + \
222                    "Disk GUID: %s\n" % uuid.UUID(bytes=self.uuid) + \
223                    "Partition entries starting LBA: %d\n" % \
224                    self.part_entry_start + \
225                    "Number of Partition entries: %d\n" % self.part_count + \
226                    "Size of a partition entry: %d\n" % self.part_entry_size + \
227                    "CRC32 of partition array: %s\n" % self.part_crc32
228
229     def __init__(self, disk):
230         self.disk = disk
231         with open(disk, "rb") as d:
232             # MBR (Logical block address 0)
233             lba0 = d.read(BLOCKSIZE)
234             self.mbr = MBR(lba0)
235
236             # Primary GPT Header (LBA 1)
237             raw_header = d.read(self.GPTHeader.size())
238             self.primary = self.GPTHeader(raw_header)
239
240             # Partition entries (LBA 2...34)
241             d.seek(self.primary.part_entry_start * BLOCKSIZE)
242             entries_size = self.primary.part_count * \
243                 self.primary.part_entry_size
244             self.part_entries = d.read(entries_size)
245
246             # Secondary GPT Header (LBA -1)
247             d.seek(self.primary.backup_lba * BLOCKSIZE)
248             raw_header = d.read(self.GPTHeader.size())
249             self.secondary = self.GPTHeader(raw_header)
250
251     def size(self):
252         """Returns the payload size of GPT partitioned device."""
253         return (self.primary.backup_lba + 1) * BLOCKSIZE
254
255     def shrink(self, size, old_size):
256         """Move the secondary GPT Header entries to the address specified by
257         size parameter.
258         """
259
260         # Most partition manipulation programs leave 2048 sector after the last
261         # partition
262         aligned = size + 2048 * BLOCKSIZE
263
264         # new_size is at least: size + Partition Entries + Secondary GPT Header
265         new_size = aligned if aligned <= old_size else \
266             size + len(self.part_entries) + BLOCKSIZE
267
268         assert new_size <= old_size, "The secodary GPT fits in the device"
269
270         if new_size == self.size():
271             return new_size
272
273         lba_count = new_size // BLOCKSIZE
274
275         # Correct MBR
276         #TODO: Check if the partition tables is hybrid
277         self.mbr.part[0].sector_count = (new_size // BLOCKSIZE) - 1
278
279         # Fix Primary header
280         self.primary.header_crc32 = 0
281         self.primary.backup_lba = lba_count - 1  # LBA-1
282         self.primary.last_usable_lba = lba_count - 34  # LBA-34
283         self.primary.header_crc32 = \
284             binascii.crc32(self.primary.pack()) & 0xffffffff
285
286         # Fix Secondary header
287         self.secondary.header_crc32 = 0
288         self.secondary.current_lba = self.primary.backup_lba
289         self.secondary.last_usable_lba = lba_count - 34  # LBA-34
290         self.secondary.part_entry_start = lba_count - 33  # LBA-33
291         self.secondary.header_crc32 = \
292             binascii.crc32(self.secondary.pack()) & 0xffffffff
293
294         # Copy the new partition table back to the device
295         with open(self.disk, "wb") as d:
296             d.write(self.mbr.pack())
297             d.write(self.primary.pack())
298             d.write('\x00' * (BLOCKSIZE - self.primary.size()))
299             d.seek(self.secondary.part_entry_start * BLOCKSIZE)
300             d.write(self.part_entries)
301             d.seek(self.primary.backup_lba * BLOCKSIZE)
302             d.write(self.secondary.pack())
303             d.write('\x00' * (BLOCKSIZE - self.secondary.size()))
304
305         return new_size
306
307 if __name__ == '__main__':
308     ptable = GPTPartitionTable(sys.argv[1])
309
310     print "MBR:\n%s" % ptable.mbr
311     print "Primary partition table:\n%s" % ptable.primary
312     print "Secondary partition table:\n%s" % ptable.secondary
313
314 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :