Statistics
| Branch: | Tag: | Revision:

root / image_creator / gpt.py @ 09743d3a

History | View | Annotate | Download (10.9 kB)

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 :