Statistics
| Branch: | Tag: | Revision:

root / image_creator / gpt.py @ 331aa0ec

History | View | Annotate | Download (9.8 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
    class Partition(object):
44
        format = "<B3sB3sLL"
45

    
46
        def __init__(self, raw_part):
47

    
48
            (   self.status,
49
                self.start,
50
                self.type,
51
                self.end,
52
                self.first_sector,
53
                self.sector_count
54
            ) = struct.unpack(self.format, raw_part)
55

    
56
        def pack(self):
57
            return struct.pack(self.format,
58
                self.status,
59
                self.start,
60
                self.type,
61
                self.end,
62
                self.first_sector,
63
                self.sector_count
64
            )
65

    
66
        def show(self):
67
            start = self.unpack_chs(self.start)
68
            end = self.unpack_chs(self.end)
69
            print "%d %s %d %s %d %d" % (self.status, start, self.type, end,
70
                self.first_sector, self.sector_count)
71

    
72
        def unpack_chs(self, chs):
73

    
74
            assert len(chs) == 3
75

    
76
            head = struct.unpack('<B', chs[0])[0]
77
            sector = struct.unpack('<B', chs[1])[0] & 0x3f
78
            cylinder = (struct.unpack('<B', chs[1])[0] & 0xC0) << 2 | \
79
                struct.unpack('<B', chs[2])[0]
80

    
81
            return (cylinder, head, sector)
82

    
83
        def pack_chs(self, cylinder, head, sector):
84

    
85
            assert 1 <= sector <= 63
86
            assert 0 <= cylinder <= 1023
87
            assert 0 <= head <= 255
88

    
89
            byte0 = head
90
            byte1 = (cylinder >> 2) & 0xC0 | sector
91
            byte2 = cylinder & 0xff
92

    
93
            return struct.pack('<BBB', byte0, byte1, byte2)
94

    
95
    format = "<444s2x16s16s16s16s2s"
96
    """
97
    Offset  Length          Contents
98
    0       440(max. 446)   code area
99
    440     2(optional)     disk signature
100
    444     2               Usually nulls
101
    446     16              Partition 0
102
    462     16              Partition 1
103
    478     16              Partition 2
104
    494     16              Partition 3
105
    510     2               MBR signature
106
    """
107
    def __init__(self, block):
108
        raw_part = {}
109
        self.code_area, \
110
        raw_part[0], \
111
        raw_part[1], \
112
        raw_part[2], \
113
        raw_part[3], \
114
        self.signature = struct.unpack(self.format, block)
115

    
116
        self.part = {}
117
        for i in range(4):
118
            self.part[i] = self.Partition(raw_part[i])
119

    
120
    def pack(self):
121
        return struct.pack(self.format,
122
            self.code_area,
123
            self.part[0].pack(),
124
            self.part[1].pack(),
125
            self.part[2].pack(),
126
            self.part[3].pack(),
127
            self.signature
128
        )
129

    
130
    def show(self):
131
        for i in range(4):
132
            print "Part %d: " % i,
133
            self.part[i].show()
134

    
135

    
136
class GPTPartitionTable(object):
137
    class GPTHeader(object):
138
        format = "<8s4sII4xQQQQ16sQIII"
139
        """
140
        Offset        Length                 Contents
141
        0       8 bytes         Signature
142
        8       4 bytes         Revision
143
        12      4 bytes         Header size in little endian
144
        16         4 bytes         CRC32 of header
145
        20         4 bytes         Reserved; must be zero
146
        24         8 bytes         Current LBA
147
        32         8 bytes         Backup LBA
148
        40         8 bytes         First usable LBA for partitions
149
        48         8 bytes         Last usable LBA
150
        56         16 bytes         Disk GUID
151
        72         8 bytes         Partition entries starting LBA
152
        80         4 bytes         Number of partition entries
153
        84         4 bytes         Size of a partition entry
154
        88         4 bytes         CRC32 of partition array
155
        92         *                 Reserved; must be zeroes
156
        LBA    size            Total
157
        """
158

    
159
        def __init__(self, block):
160
            self.signature, \
161
            self.revision, \
162
            self.size, \
163
            self.header_crc32, \
164
            self.current_lba, \
165
            self.backup_lba, \
166
            self.first_usable_lba, \
167
            self.last_usable_lba, \
168
            self.uuid, \
169
            self.part_entry_start, \
170
            self.part_count, \
171
            self.part_entry_size, \
172
            self.part_crc32 = struct.unpack(self.format, block)
173

    
174
        def pack(self):
175
            return struct.pack(self.format,
176
                self.signature, \
177
                self.revision, \
178
                self.size, \
179
                self.header_crc32, \
180
                self.current_lba, \
181
                self.backup_lba, \
182
                self.first_usable_lba, \
183
                self.last_usable_lba, \
184
                self.uuid, \
185
                self.part_entry_start, \
186
                self.part_count, \
187
                self.part_entry_size, \
188
                self.part_crc32
189
            )
190

    
191
        def show(self):
192
            print "Signature: %s" % self.signature
193
            print "Revision: %r" % self.revision
194
            print "Header Size: %d" % self.size
195
            print "CRC32: %d" % self.header_crc32
196
            print "Current LBA: %d" % self.current_lba
197
            print "Backup LBA: %d" % self.backup_lba
198
            print "First Usable LBA: %d" % self.first_usable_lba
199
            print "Last Usable LBA: %d" % self.last_usable_lba
200
            print "Disk GUID: %s" % uuid.UUID(bytes=self.uuid)
201
            print "Partition entries starting LBA: %d" % self.part_entry_start
202
            print "Number of Partition entries: %d" % self.part_count
203
            print "Size of a partition entry: %d" % self.part_entry_size
204
            print "CRC32 of partition array: %s" % self.part_crc32
205

    
206
    def __init__(self, disk):
207
        self.disk = disk
208
        with open(disk, "rb") as d:
209
            #MBR (Logical block address 0)
210
            lba0 = d.read(BLOCKSIZE)
211
            self.mbr = MBR(lba0)
212
            # Primary GPT Header (LBA 1)
213
            lba1 = d.read(BLOCKSIZE)
214
            self.primary = self.GPTHeader(lba1[:92])
215
            # Partition entries (LBA 2...34)
216
            d.seek(self.primary.part_entry_start * BLOCKSIZE)
217
            entries_size = self.primary.part_count * \
218
                                                self.primary.part_entry_size
219
            self.part_entries = d.read(entries_size)
220
            # Secondary GPT Header (LBA -1)
221
            d.seek(self.primary.backup_lba * BLOCKSIZE)
222
            lba_1 = d.read(BLOCKSIZE)
223
            self.secondary = self.GPTHeader(lba_1[:92])
224

    
225
    def size(self):
226
        return (self.primary.backup_lba + 1) * BLOCKSIZE 
227

    
228
    def shrink(self, size):
229

    
230
        if size == self.size():
231
            return size
232

    
233
        assert size < self.size()
234

    
235
        # new_size = size + Partition Entries + Secondary GPT Header
236
        new_size = size + len(self.part_entries) + BLOCKSIZE
237
        new_size = ((new_size + 4095) // 4096) * 4096  # align to 4K
238
        lba_count = new_size // BLOCKSIZE
239

    
240
        # Correct MBR
241
        #TODO: Check for hybrid partition tables
242
        self.mbr.part[0].sector_count = (new_size // BLOCKSIZE) - 1
243

    
244
        # Correct Primary header
245
        self.primary.header_crc32 = 0
246
        self.primary.backup_lba = lba_count - 1  # LBA-1
247
        self.primary.last_usable_lba = lba_count - 34  # LBA-34
248
        self.primary.header_crc32 = \
249
                            binascii.crc32(self.primary.pack()) & 0xffffffff
250

    
251
        # Correct Secondary header entries
252
        self.secondary.header_crc32 = 0
253
        self.secondary.current_lba = self.primary.backup_lba
254
        self.secondary.last_usable_lba = lba_count - 34  # LBA-34
255
        self.secondary.part_entry_start = lba_count - 33  # LBA-33
256
        self.secondary.header_crc32 = \
257
                            binascii.crc32(self.secondary.pack()) & 0xffffffff
258

    
259
        # Copy the new partition table back to the device
260
        with open(self.disk, "wb") as d:
261
            d.write(self.mbr.pack())
262
            d.write(struct.pack("%ss" % BLOCKSIZE, '\x00' * BLOCKSIZE))
263
            d.seek(BLOCKSIZE)
264
            d.write(self.primary.pack())
265
            d.seek(self.secondary.part_entry_start * BLOCKSIZE)
266
            d.write(self.part_entries)
267
            d.seek(self.primary.backup_lba * BLOCKSIZE)
268
            d.write(struct.pack("%ss" % BLOCKSIZE, '\x00' * BLOCKSIZE))
269
            d.seek(self.primary.backup_lba * BLOCKSIZE)
270
            d.write(self.secondary.pack())
271

    
272
        return new_size
273

    
274
if __name__ == '__main__':
275
    ptable = GPTPartitionTable(sys.argv[1])
276

    
277
    print "MBR:"
278
    ptable.mbr.show()
279
    print
280
    print "Primary partition table:"
281
    ptable.primary.show()
282
    print
283
    print "Secondary partition table:"
284
    ptable.secondary.show()
285

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