Statistics
| Branch: | Tag: | Revision:

root / image_creator / gpt.py @ c71f38be

History | View | Annotate | Download (11.9 kB)

1
#!/usr/bin/env python
2
# -*- coding: utf-8 -*-
3
#
4
# Copyright 2012 GRNET S.A. All rights reserved.
5
#
6
# Redistribution and use in source and binary forms, with or
7
# without modification, are permitted provided that the following
8
# conditions are met:
9
#
10
#   1. Redistributions of source code must retain the above
11
#      copyright notice, this list of conditions and the following
12
#      disclaimer.
13
#
14
#   2. Redistributions in binary form must reproduce the above
15
#      copyright notice, this list of conditions and the following
16
#      disclaimer in the documentation and/or other materials
17
#      provided with the distribution.
18
#
19
# THIS SOFTWARE IS PROVIDED BY GRNET S.A. ``AS IS'' AND ANY EXPRESS
20
# OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
21
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
22
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GRNET S.A OR
23
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
24
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
25
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
26
# USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
27
# AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
28
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
29
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
30
# POSSIBILITY OF SUCH DAMAGE.
31
#
32
# The views and conclusions contained in the software and
33
# documentation are those of the authors and should not be
34
# interpreted as representing official policies, either expressed
35
# or implied, of GRNET S.A.
36

    
37
"""This module provides the code for handling GUID partition tables"""
38

    
39
import struct
40
import sys
41
import uuid
42
import binascii
43

    
44
BLOCKSIZE = 512
45

    
46

    
47
class MBR(object):
48
    """Represents a Master Boot Record."""
49
    class Partition(object):
50
        """Represents a partition entry in MBR"""
51
        format = "<B3sB3sLL"
52

    
53
        def __init__(self, raw_part):
54
            """Create a Partition instance"""
55
            (
56
                self.status,
57
                self.start,
58
                self.type,
59
                self.end,
60
                self.first_sector,
61
                self.sector_count
62
            ) = struct.unpack(self.format, raw_part)
63

    
64
        def pack(self):
65
            """Pack the partition values into a binary string"""
66
            return struct.pack(self.format,
67
                               self.status,
68
                               self.start,
69
                               self.type,
70
                               self.end,
71
                               self.first_sector,
72
                               self.sector_count)
73

    
74
        @staticmethod
75
        def size():
76
            """Returns the size of an MBR partition entry"""
77
            return struct.calcsize(MBR.Partition.format)
78

    
79
        def __str__(self):
80
            start = self.unpack_chs(self.start)
81
            end = self.unpack_chs(self.end)
82
            return "%d %s %d %s %d %d" % (self.status, start, self.type, end,
83
                                          self.first_sector, self.sector_count)
84

    
85
        def unpack_chs(self, chs):
86
            """Unpacks a CHS address string to a tuple."""
87

    
88
            assert len(chs) == 3
89

    
90
            head = struct.unpack('<B', chs[0])[0]
91
            sector = struct.unpack('<B', chs[1])[0] & 0x3f
92
            cylinder = (struct.unpack('<B', chs[1])[0] & 0xC0) << 2 | \
93
                struct.unpack('<B', chs[2])[0]
94

    
95
            return (cylinder, head, sector)
96

    
97
        def pack_chs(self, cylinder, head, sector):
98
            """Packs a CHS tuple to an address string."""
99

    
100
            assert 1 <= sector <= 63
101
            assert 0 <= cylinder <= 1023
102
            assert 0 <= head <= 255
103

    
104
            byte0 = head
105
            byte1 = (cylinder >> 2) & 0xC0 | sector
106
            byte2 = cylinder & 0xff
107

    
108
            return struct.pack('<BBB', byte0, byte1, byte2)
109

    
110
    format = "<444s2x16s16s16s16s2s"
111
    """
112
    Offset  Length          Contents
113
    0       440(max. 446)   code area
114
    440     2(optional)     disk signature
115
    444     2               Usually nulls
116
    446     16              Partition 0
117
    462     16              Partition 1
118
    478     16              Partition 2
119
    494     16              Partition 3
120
    510     2               MBR signature
121
    """
122
    def __init__(self, block):
123
        """Create an MBR instance"""
124
        raw_part = {}
125
        (self.code_area,
126
         raw_part[0],
127
         raw_part[1],
128
         raw_part[2],
129
         raw_part[3],
130
         self.signature) = struct.unpack(self.format, block)
131

    
132
        self.part = {}
133
        for i in range(4):
134
            self.part[i] = self.Partition(raw_part[i])
135

    
136
    @staticmethod
137
    def size():
138
        """Return the size of a Master Boot Record."""
139
        return struct.calcsize(MBR.format)
140

    
141
    def pack(self):
142
        """Pack an MBR to a binary string."""
143
        return struct.pack(self.format,
144
                           self.code_area,
145
                           self.part[0].pack(),
146
                           self.part[1].pack(),
147
                           self.part[2].pack(),
148
                           self.part[3].pack(),
149
                           self.signature)
150

    
151
    def __str__(self):
152
        ret = ""
153
        for i in range(4):
154
            ret += "Partition %d: %s\n" % (i, self.part[i])
155
        ret += "Signature: %s %s\n" % (hex(ord(self.signature[0])),
156
                                       hex(ord(self.signature[1])))
157
        return ret
158

    
159

    
160
class GPTPartitionTable(object):
161
    """Represents a GUID Partition Table."""
162
    class GPTHeader(object):
163
        """Represents a GPT Header of a GUID Partition Table."""
164
        format = "<8s4sII4xQQQQ16sQIII"
165
        """
166
        Offset        Length                 Contents
167
        0       8 bytes         Signature
168
        8       4 bytes         Revision
169
        12      4 bytes         Header size in little endian
170
        16         4 bytes         CRC32 of header
171
        20         4 bytes         Reserved; must be zero
172
        24         8 bytes         Current LBA
173
        32         8 bytes         Backup LBA
174
        40         8 bytes         First usable LBA for partitions
175
        48         8 bytes         Last usable LBA
176
        56         16 bytes         Disk GUID
177
        72         8 bytes         Partition entries starting LBA
178
        80         4 bytes         Number of partition entries
179
        84         4 bytes         Size of a partition entry
180
        88         4 bytes         CRC32 of partition array
181
        92         *                 Reserved; must be zeroes
182
        LBA    size            Total
183
        """
184

    
185
        def __init__(self, block):
186
            """Create a GPTHeader instance"""
187
            (self.signature,
188
             self.revision,
189
             self.hdr_size,
190
             self.header_crc32,
191
             self.current_lba,
192
             self.backup_lba,
193
             self.first_usable_lba,
194
             self.last_usable_lba,
195
             self.uuid,
196
             self.part_entry_start,
197
             self.part_count,
198
             self.part_entry_size,
199
             self.part_crc32) = struct.unpack(self.format, block)
200

    
201
        def pack(self):
202
            """Packs a GPT Header to a binary string."""
203
            return struct.pack(self.format,
204
                               self.signature,
205
                               self.revision,
206
                               self.hdr_size,
207
                               self.header_crc32,
208
                               self.current_lba,
209
                               self.backup_lba,
210
                               self.first_usable_lba,
211
                               self.last_usable_lba,
212
                               self.uuid,
213
                               self.part_entry_start,
214
                               self.part_count,
215
                               self.part_entry_size,
216
                               self.part_crc32)
217

    
218
        @staticmethod
219
        def size():
220
            """Return the size of a GPT Header."""
221
            return struct.calcsize(GPTPartitionTable.GPTHeader.format)
222

    
223
        def __str__(self):
224
            """Print a GPTHeader"""
225
            return "Signature: %s\n" % self.signature + \
226
                   "Revision: %r\n" % self.revision + \
227
                   "Header Size: %d\n" % self.hdr_size + \
228
                   "CRC32: %d\n" % self.header_crc32 + \
229
                   "Current LBA: %d\n" % self.current_lba + \
230
                   "Backup LBA: %d\n" % self.backup_lba + \
231
                   "First Usable LBA: %d\n" % self.first_usable_lba + \
232
                   "Last Usable LBA: %d\n" % self.last_usable_lba + \
233
                   "Disk GUID: %s\n" % uuid.UUID(bytes=self.uuid) + \
234
                   "Partition entries starting LBA: %d\n" % \
235
                   self.part_entry_start + \
236
                   "Number of Partition entries: %d\n" % self.part_count + \
237
                   "Size of a partition entry: %d\n" % self.part_entry_size + \
238
                   "CRC32 of partition array: %s\n" % self.part_crc32
239

    
240
    def __init__(self, disk):
241
        """Create a GPTPartitionTable instance"""
242
        self.disk = disk
243
        with open(disk, "rb") as d:
244
            # MBR (Logical block address 0)
245
            lba0 = d.read(BLOCKSIZE)
246
            self.mbr = MBR(lba0)
247

    
248
            # Primary GPT Header (LBA 1)
249
            raw_header = d.read(self.GPTHeader.size())
250
            self.primary = self.GPTHeader(raw_header)
251

    
252
            # Partition entries (LBA 2...34)
253
            d.seek(self.primary.part_entry_start * BLOCKSIZE)
254
            entries_size = self.primary.part_count * \
255
                self.primary.part_entry_size
256
            self.part_entries = d.read(entries_size)
257

    
258
            # Secondary GPT Header (LBA -1)
259
            d.seek(self.primary.backup_lba * BLOCKSIZE)
260
            raw_header = d.read(self.GPTHeader.size())
261
            self.secondary = self.GPTHeader(raw_header)
262

    
263
    def size(self):
264
        """Return the payload size of GPT partitioned device."""
265
        return (self.primary.backup_lba + 1) * BLOCKSIZE
266

    
267
    def shrink(self, size, old_size):
268
        """Move the secondary GPT Header entries to the address specified by
269
        size parameter.
270
        """
271

    
272
        # Most partition manipulation programs leave 2048 sector after the last
273
        # partition
274
        aligned = size + 2048 * BLOCKSIZE
275

    
276
        # new_size is at least: size + Partition Entries + Secondary GPT Header
277
        new_size = aligned if aligned <= old_size else \
278
            size + len(self.part_entries) + BLOCKSIZE
279

    
280
        assert new_size <= old_size, "The secodary GPT fits in the device"
281

    
282
        if new_size == self.size():
283
            return new_size
284

    
285
        lba_count = new_size // BLOCKSIZE
286

    
287
        # Correct MBR
288
        #TODO: Check if the partition tables is hybrid
289
        self.mbr.part[0].sector_count = (new_size // BLOCKSIZE) - 1
290

    
291
        # Fix Primary header
292
        self.primary.header_crc32 = 0
293
        self.primary.backup_lba = lba_count - 1  # LBA-1
294
        self.primary.last_usable_lba = lba_count - 34  # LBA-34
295
        self.primary.header_crc32 = \
296
            binascii.crc32(self.primary.pack()) & 0xffffffff
297

    
298
        # Fix Secondary header
299
        self.secondary.header_crc32 = 0
300
        self.secondary.current_lba = self.primary.backup_lba
301
        self.secondary.last_usable_lba = lba_count - 34  # LBA-34
302
        self.secondary.part_entry_start = lba_count - 33  # LBA-33
303
        self.secondary.header_crc32 = \
304
            binascii.crc32(self.secondary.pack()) & 0xffffffff
305

    
306
        # Copy the new partition table back to the device
307
        with open(self.disk, "wb") as d:
308
            d.write(self.mbr.pack())
309
            d.write(self.primary.pack())
310
            d.write('\x00' * (BLOCKSIZE - self.primary.size()))
311
            d.write(self.part_entries)
312
            d.seek(self.secondary.part_entry_start * BLOCKSIZE)
313
            d.write(self.part_entries)
314
            d.seek(self.primary.backup_lba * BLOCKSIZE)
315
            d.write(self.secondary.pack())
316
            d.write('\x00' * (BLOCKSIZE - self.secondary.size()))
317

    
318
        return new_size
319

    
320
if __name__ == '__main__':
321
    ptable = GPTPartitionTable(sys.argv[1])
322

    
323
    print "MBR:\n%s" % ptable.mbr
324
    print "Primary partition table:\n%s" % ptable.primary
325
    print "Secondary partition table:\n%s" % ptable.secondary
326

    
327
# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :